From 2df206199c3a2ca85aaa8ba42ce766866dd0e7eb Mon Sep 17 00:00:00 2001 From: EthanPeng-2048 Date: Sat, 31 Oct 2026 00:45:09 +0800 Subject: [PATCH 1/8] =?UTF-8?q?fix:=20=E5=BF=BD=E7=95=A5=E7=BC=96=E8=BE=91?= =?UTF-8?q?=E5=99=A8=E7=BC=93=E5=AD=98=E6=96=87=E4=BB=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 9f772398..ced58148 100644 --- a/.gitignore +++ b/.gitignore @@ -68,6 +68,7 @@ dist/ # ==================== Editors ==================== .vscode/ +.cache/ .idea/ *.swp *.swo From 8406b39550bbc6f47eb271aa479c3db1d33c7317 Mon Sep 17 00:00:00 2001 From: EthanPeng-2048 Date: Sat, 31 Oct 2026 01:24:00 +0800 Subject: [PATCH 2/8] =?UTF-8?q?fix:=20=E8=A7=84=E8=8C=83=E5=8C=96=E9=94=99?= =?UTF-8?q?=E8=AF=AF=E5=A4=84=E7=90=86?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- include/neuralnet.cpp/text_trainer.hpp | 10 ++++------ src/session_checkpoint_test.cpp | 6 ++---- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/include/neuralnet.cpp/text_trainer.hpp b/include/neuralnet.cpp/text_trainer.hpp index b2c7b85a..72502c99 100644 --- a/include/neuralnet.cpp/text_trainer.hpp +++ b/include/neuralnet.cpp/text_trainer.hpp @@ -480,7 +480,7 @@ inline Result TextTrainer::fit() // 装配回调只产出归一化单位（有效 token）；真正的 batch 装配在 do_update // 里按 micro 子跨度做（global batch 的 micro 切分是文本域语义）。 const std::uint64_t seed = static_cast(std::random_device{}()); - auto source_r = IndexedBatchSource::create( + NN_TRY(source_r, IndexedBatchSource::create( train_.sample_count(), cfg_.batch_size, seed, [this](std::span ix) -> Result> @@ -488,8 +488,7 @@ inline Result TextTrainer::fit() NN_TRY_MSG(u, "统计 global batch 有效 token 数失败", train_.valid_target_count(ix)); return std::make_pair(std::any{}, static_cast(*u)); - }); - if (!source_r) return std::unexpected(source_r.error()); + })); IndexedBatchSource& source = *source_r; // ── 非阻塞 loss 回读状态（设备侧异步槽位；host 不等 GPU）───────────── @@ -557,9 +556,8 @@ inline Result TextTrainer::fit() TextSampleBatch batch; // 可复用的装配缓冲（训练 micro + 评估共用） - auto driver_r = RunDriver::create( - session_, source, std::make_unique(train_state_), std::move(rcfg)); - if (!driver_r) return std::unexpected(driver_r.error()); + NN_TRY(driver_r, RunDriver::create( + session_, source, std::make_unique(train_state_), std::move(rcfg))); RunDriver& driver = *driver_r; RunHooks hooks; diff --git a/src/session_checkpoint_test.cpp b/src/session_checkpoint_test.cpp index 128e6f90..39e1844f 100644 --- a/src/session_checkpoint_test.cpp +++ b/src/session_checkpoint_test.cpp @@ -59,14 +59,12 @@ nn::Result one_step(nn::TrainingSession& s, opts.force_update = true; return s.train_step(0, opts, [&](nn::Model& model, nn::ComputeEngine& eng, int) -> nn::Result { - auto logits = model.forward(x); - if (!logits) return std::unexpected(logits.error()); + NN_TRY(logits, model.forward(x)); // step_fn 契约：返回**对模型输出**的梯度；model.backward 由 session 执行 //（本测试焦点是断点编排，直接给常数输出梯度，不经 loss 层） auto go = eng.create_tensor(logits->rows(), logits->cols(), nn::Precision::F32, nn::InitSpec::constant(0.01f)); - if (!go.valid()) - return std::unexpected(nn::Error{"one_step: go 分配失败"}); + if (!go.valid()) NN_FAIL("one_step: go 分配失败"); return nn::TrainingStepOutput{std::move(go), 1, {}, 0}; }); } From 5854b98c5d6f3204caacb610288bc152643916be Mon Sep 17 00:00:00 2001 From: EthanPeng-2048 Date: Sat, 31 Oct 2026 02:15:49 +0800 Subject: [PATCH 3/8] =?UTF-8?q?style:=20=E5=85=A8=E4=BB=93=E7=BB=9F?= =?UTF-8?q?=E4=B8=80=20clang-format=20=E6=A0=BC=E5=BC=8F=E5=8C=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 对主仓 201 个 C++ 源文件（include/src/tools/examples）统一套用根目录 .clang-format 目标风格（Allman / 4 空格 / 100 列 / include 不排序）， 182 个文件产生 diff。五个 extern/ 子仓同步格式化并各自产生独立提交， 本提交同时更新子仓指针。 验证：cmake --build build/dev 全量重编 159/159 目标通过，零错误。 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- examples/downstream/main.cpp | 13 +- examples/fusion_custom_layer.hpp | 29 +- examples/fusion_custom_layer_main.cpp | 50 +- extern/base | 2 +- extern/dataset | 2 +- extern/kvrecord | 2 +- extern/threadpool | 2 +- extern/tokenizer | 2 +- include/neuralnet.cpp/algebra_matrix.hpp | 1607 +++--- include/neuralnet.cpp/algebra_ops.hpp | 6 +- include/neuralnet.cpp/algebra_span.hpp | 39 +- .../backend/compute_gpu_storage.hpp | 82 +- .../backend/compute_gpu_tensor_impl.hpp | 51 +- .../backend/compute_memory_pool.hpp | 160 +- .../backend/compute_staging_ring.hpp | 105 +- .../backend/compute_vk_backend.hpp | 70 +- .../backend/compute_vk_backend_destroy.inc | 1113 ++--- .../backend/compute_vk_backend_frame.inc | 1333 +++-- .../backend/compute_vk_backend_ops.inc | 4332 ++++++++--------- .../backend/compute_vk_backend_pipelines.inc | 1256 +++-- .../backend/compute_vk_device.hpp | 174 +- .../backend/compute_vk_pipeline_catalog.hpp | 31 +- .../backend/compute_vk_pipeline_specs.hpp | 138 +- include/neuralnet.cpp/batch_source.hpp | 34 +- .../neuralnet.cpp/cli/cli_engine_factory.hpp | 55 +- include/neuralnet.cpp/cli/cli_gpu_option.hpp | 61 +- include/neuralnet.cpp/cli/cli_help.hpp | 426 +- .../neuralnet.cpp/cli/cli_lr_scheduler.hpp | 4 +- include/neuralnet.cpp/cli/cli_mnist_io.hpp | 331 +- .../cli/cli_precision_warning.hpp | 51 +- .../neuralnet.cpp/cli/cli_train_common.hpp | 223 +- .../cli/text_train_checkpoint.hpp | 18 +- .../neuralnet.cpp/cli/text_train_config.hpp | 370 +- include/neuralnet.cpp/cli/text_train_loop.hpp | 124 +- .../neuralnet.cpp/cli/text_train_recovery.hpp | 106 +- include/neuralnet.cpp/compute_cpu_engine.hpp | 1081 ++-- include/neuralnet.cpp/compute_engine.hpp | 777 ++- include/neuralnet.cpp/compute_gpu_engine.hpp | 998 ++-- .../neuralnet.cpp/compute_layer_attention.hpp | 426 +- include/neuralnet.cpp/compute_layer_base.hpp | 249 +- .../neuralnet.cpp/compute_layer_bitlinear.hpp | 242 +- .../neuralnet.cpp/compute_layer_bundle.hpp | 195 +- include/neuralnet.cpp/compute_layer_conv.hpp | 191 +- .../neuralnet.cpp/compute_layer_embedding.hpp | 32 +- .../compute_layer_feedforward.hpp | 64 +- include/neuralnet.cpp/compute_layer_gpt.hpp | 254 +- include/neuralnet.cpp/compute_layer_mlp.hpp | 612 +-- include/neuralnet.cpp/compute_layer_rapt.hpp | 513 +- .../neuralnet.cpp/compute_layer_softmax.hpp | 27 +- .../compute_layer_transformer.hpp | 182 +- include/neuralnet.cpp/compute_loss.hpp | 193 +- include/neuralnet.cpp/compute_optimizer.hpp | 449 +- .../compute_position_encoding.hpp | 233 +- include/neuralnet.cpp/compute_tensor.hpp | 158 +- include/neuralnet.cpp/core_env.hpp | 12 +- include/neuralnet.cpp/core_observer_ptr.hpp | 25 +- include/neuralnet.cpp/core_profile.hpp | 173 +- include/neuralnet.cpp/dataset_samples.hpp | 8 +- include/neuralnet.cpp/domain_cnn.hpp | 84 +- include/neuralnet.cpp/domain_gpt.hpp | 151 +- include/neuralnet.cpp/domain_mnist.hpp | 118 +- include/neuralnet.cpp/domain_rla.hpp | 92 +- include/neuralnet.cpp/expr_dsl.hpp | 677 +-- include/neuralnet.cpp/expr_emitter.hpp | 36 +- include/neuralnet.cpp/expr_fold.hpp | 122 +- include/neuralnet.cpp/expr_glsl_gen.hpp | 1278 +++-- include/neuralnet.cpp/expr_graph.hpp | 300 +- include/neuralnet.cpp/expr_op_table.hpp | 106 +- include/neuralnet.cpp/expr_opt.hpp | 159 +- include/neuralnet.cpp/expr_registry.hpp | 18 +- include/neuralnet.cpp/expr_spec.hpp | 582 ++- include/neuralnet.cpp/lr_scheduler.hpp | 8 +- include/neuralnet.cpp/model_container.hpp | 71 +- include/neuralnet.cpp/model_serialization.hpp | 632 ++- include/neuralnet.cpp/model_spec.hpp | 184 +- include/neuralnet.cpp/named_tensor_set.hpp | 112 +- include/neuralnet.cpp/nn.hpp | 1 - include/neuralnet.cpp/rl_experience.hpp | 383 +- include/neuralnet.cpp/rl_pipeline.hpp | 196 +- include/neuralnet.cpp/rl_policy.hpp | 871 ++-- include/neuralnet.cpp/run_driver.hpp | 57 +- include/neuralnet.cpp/tensor_schema.hpp | 230 +- include/neuralnet.cpp/ternary_pack.hpp | 46 +- include/neuralnet.cpp/text_trainer.hpp | 231 +- include/neuralnet.cpp/train_state.hpp | 138 +- .../neuralnet.cpp/training_group_session.hpp | 78 +- include/neuralnet.cpp/training_session.hpp | 83 +- research/f16_weight_decay/decay_probe.cpp | 86 +- research/ternary_scale/probe_scale.cpp | 117 +- research/ternary_ste/poc_binary_ste.cpp | 42 +- research/ternary_ste/poc_ternary_ste.cpp | 50 +- src/adam_quant_test.cpp | 59 +- src/assert_always_test.cpp | 20 +- src/attn_consistency_test.cpp | 123 +- src/attn_gradcheck.cpp | 68 +- src/attn_test.cpp | 98 +- src/attn_w_batch_test.cpp | 71 +- src/batch_source_test.cpp | 46 +- src/batchnorm_test.cpp | 196 +- src/bitnet_struct_test.cpp | 323 +- src/ce_fusion_test.cpp | 210 +- src/chat_prompt_test.cpp | 174 +- src/cli_precision_warning_test.cpp | 35 +- src/cnn_smoke_test.cpp | 69 +- src/conv2d_gradcheck.cpp | 26 +- src/cross_attention_test.cpp | 51 +- src/dataset_convert.cpp | 38 +- src/dataset_gen.cpp | 73 +- src/dataset_samples_test.cpp | 3 +- src/doc_attn_e2e_test.cpp | 65 +- src/dsl_shape_infer_test.cpp | 117 +- src/embedding_test.cpp | 54 +- src/expr_dsl_test.cpp | 84 +- src/expr_fold_test.cpp | 237 +- src/expr_fuse_shapes.cpp | 13 +- src/expr_fuse_shapes.hpp | 63 +- src/expr_fuse_test.cpp | 124 +- src/expr_fused_key_test.cpp | 104 +- src/expr_graph_test.cpp | 183 +- src/expr_matmul_test.cpp | 208 +- src/expr_opt_test.cpp | 200 +- src/expr_reduce_test.cpp | 119 +- src/f16_compute_test.cpp | 43 +- src/f16_cpu_probe.cpp | 381 +- src/f16_precision_test.cpp | 327 +- src/f16_writeback_probe.cpp | 193 +- src/fused_gpu_test.cpp | 412 +- src/gpt_checkpoint_test.cpp | 28 +- src/gpt_gradcheck.cpp | 111 +- src/gpt_offload_test.cpp | 28 +- src/gpu_batch_scope_test.cpp | 18 +- src/gpu_f16_test.cpp | 196 +- src/gpu_stability_probe.cpp | 640 ++- src/gpu_test.cpp | 370 +- src/host_bridge_test.cpp | 21 +- src/layer_bench.cpp | 319 +- src/layer_broadcast_test.cpp | 68 +- src/layer_lifetime_test.cpp | 40 +- src/lifetime_contract_test.cpp | 17 +- src/lr_scheduler_test.cpp | 2 +- src/maxpool_gradcheck.cpp | 58 +- src/mem_probe.cpp | 178 +- src/mnist_infer.cpp | 94 +- src/mnist_train.cpp | 403 +- src/model_roles_test.cpp | 70 +- src/model_spec_validation_test.cpp | 214 +- src/optimizer_state_test.cpp | 93 +- src/optimizer_test.cpp | 50 +- src/precision_profile_test.cpp | 10 +- src/precision_test.cpp | 26 +- src/rapt_checkpoint_test.cpp | 28 +- src/rapt_gradcheck.cpp | 174 +- src/rapt_offload_test.cpp | 28 +- src/rapt_smoke_test.cpp | 50 +- src/rl_experience_test.cpp | 91 +- src/rl_gpu_smoke_test.cpp | 69 +- src/rl_grpo_test.cpp | 76 +- src/rl_output_test.cpp | 129 +- src/rl_pipeline_test.cpp | 164 +- src/rl_policy_test.cpp | 210 +- src/rl_ppo_test.cpp | 164 +- src/rl_training_integration_test.cpp | 60 +- src/rmsnorm_gradcheck.cpp | 87 +- src/run_driver_test.cpp | 98 +- src/sampling_test.cpp | 37 +- src/scan_bitwise_test.cpp | 55 +- src/session_checkpoint_test.cpp | 59 +- src/softmax_gradcheck.cpp | 32 +- src/swiglu_gradcheck.cpp | 92 +- src/t1_58_test.cpp | 142 +- src/tensor_expr_test.cpp | 100 +- src/tensor_precision_test.cpp | 12 +- src/tensor_schema_test.cpp | 92 +- src/ternary_gemm_test.cpp | 110 +- src/test_common.hpp | 45 +- src/text_infer.cpp | 197 +- src/text_train.cpp | 120 +- src/text_train_modules_test.cpp | 121 +- src/text_trainer_test.cpp | 28 +- src/tokenizer_infer.cpp | 93 +- src/tokenizer_train.cpp | 54 +- src/train_state_test.cpp | 108 +- src/training_group_session_test.cpp | 87 +- src/training_session_test.cpp | 95 +- src/write_epoch_test.cpp | 142 +- tools/fused_generate.hpp | 301 +- tools/scan_exprs.cpp | 1336 +++-- 187 files changed, 19039 insertions(+), 20346 deletions(-) diff --git a/examples/downstream/main.cpp b/examples/downstream/main.cpp index 8f0d7853..1b81d270 100644 --- a/examples/downstream/main.cpp +++ b/examples/downstream/main.cpp @@ -61,8 +61,8 @@ bool run_linear(nn::ComputeEngine& engine, const char* tag, float& first_out) } if (y->rows() != OUT || y->cols() != BATCH) { - std::fprintf(stderr, "[%s] 输出形状错误: 期望 (%zu,%zu)，实际 (%zu,%zu)\n", - tag, OUT, BATCH, y->rows(), y->cols()); + std::fprintf(stderr, "[%s] 输出形状错误: 期望 (%zu,%zu)，实际 (%zu,%zu)\n", tag, OUT, BATCH, + y->rows(), y->cols()); return false; } @@ -95,8 +95,7 @@ int main() nn::CpuEngine& cpu = *cpu_owner; float cpu_out = 0.0f; if (!run_linear(cpu, "cpu", cpu_out)) return 1; - std::printf("[downstream] CPU Linear forward OK（y[0]=%.6f）\n", - static_cast(cpu_out)); + std::printf("[downstream] CPU Linear forward OK（y[0]=%.6f）\n", static_cast(cpu_out)); #ifdef NN_HAS_VULKAN // ── GPU 路径（neuralnet::nn 自动传入 NN_HAS_VULKAN 才编译进来）──────── @@ -105,8 +104,7 @@ int main() nn::GpuBackend& backend = *backend_owner; if (auto ir = backend.initialize(); !ir) { - std::printf("[downstream] GPU 初始化失败，跳过 GPU 校验：%s\n", - ir.error().message.c_str()); + std::printf("[downstream] GPU 初始化失败，跳过 GPU 校验：%s\n", ir.error().message.c_str()); return 77; } auto gpu_owner = nn::GpuEngine::create(backend); @@ -116,8 +114,7 @@ int main() // GPU 与 CPU 用不同随机权重初始化，无法直接对拍数值； // 这里只断言两者都跑通（真正的跨引擎对拍在库内测试覆盖）。 - std::printf("[downstream] GPU Linear forward OK（y[0]=%.6f）\n", - static_cast(gpu_out)); + std::printf("[downstream] GPU Linear forward OK（y[0]=%.6f）\n", static_cast(gpu_out)); // GPU 与 CPU 权重不同 → 仅形状/有限性断言已由 run_linear 完成。 (void)cpu_out; #else diff --git a/examples/fusion_custom_layer.hpp b/examples/fusion_custom_layer.hpp index 3ee65934..1fc52969 100644 --- a/examples/fusion_custom_layer.hpp +++ b/examples/fusion_custom_layer.hpp @@ -22,14 +22,15 @@ namespace nn_example // out = exp(|x| * alpha) / (sqrt(|y|) + alpha) // 两个运行时标量参数（alpha）→ 值不进 key，同结构共享 shader。 // 入参只读：`dsl::leaf` 按值收 Tensor（句柄拷贝），表达式不写输入。 -[[nodiscard]] inline nn::Result fused_custom_op( - nn::ComputeEngine& engine, const nn::Tensor& x, const nn::Tensor& y, float alpha, - nn::Precision P = nn::Precision::F32) +[[nodiscard]] inline nn::Result fused_custom_op(nn::ComputeEngine& engine, + const nn::Tensor& x, + const nn::Tensor& y, float alpha, + nn::Precision P = nn::Precision::F32) { - return nn::dsl::compute(engine, - nn::dsl::exp(nn::dsl::abs(nn::dsl::leaf(x)) * nn::dsl::rparam(alpha)) - / (nn::dsl::sqrt(nn::dsl::abs(nn::dsl::leaf(y))) - + nn::dsl::rparam(alpha)), + return nn::dsl::compute( + engine, + nn::dsl::exp(nn::dsl::abs(nn::dsl::leaf(x)) * nn::dsl::rparam(alpha)) / + (nn::dsl::sqrt(nn::dsl::abs(nn::dsl::leaf(y))) + nn::dsl::rparam(alpha)), x.rows(), x.cols(), P); } @@ -39,13 +40,15 @@ namespace nn_example // 25 条之一——即**锚点的符号实例与运行期实例折叠出不同结构**。 // 库外自定义层没有 dry-run/模型 pass，因此本函数是对 `nn_enable_gpu_fusion` // **可靠性**的直接检验：注册表若漏了运行期那一条，GPU 会闭合世界硬报错。 -[[nodiscard]] inline nn::Result fused_custom_norm( - nn::ComputeEngine& engine, const nn::Tensor& x, float inv_cols, float eps, - nn::Precision P = nn::Precision::F32) +[[nodiscard]] inline nn::Result fused_custom_norm(nn::ComputeEngine& engine, + const nn::Tensor& x, float inv_cols, + float eps, + nn::Precision P = nn::Precision::F32) { - return nn::dsl::compute(engine, - nn::dsl::rsqrt(nn::dsl::col_reduce_sum(nn::dsl::leaf(x)) - * nn::dsl::rparam(inv_cols) + nn::dsl::rparam(eps)), + return nn::dsl::compute( + engine, + nn::dsl::rsqrt(nn::dsl::col_reduce_sum(nn::dsl::leaf(x)) * nn::dsl::rparam(inv_cols) + + nn::dsl::rparam(eps)), x.rows(), x.cols(), P); } diff --git a/examples/fusion_custom_layer_main.cpp b/examples/fusion_custom_layer_main.cpp index c1149986..cd97f88c 100644 --- a/examples/fusion_custom_layer_main.cpp +++ b/examples/fusion_custom_layer_main.cpp @@ -26,7 +26,8 @@ namespace // exp(0.5*0.7) / (sqrt(2.0) + 0.7) constexpr float kExpect = 1.419067549f / (1.414213562f + 0.7f); -float abs_diff(float a, float b) { return a > b ? a - b : b - a; } +float abs_diff(float a, float b) +{ return a > b ? a - b : b - a; } } // namespace @@ -35,15 +36,12 @@ int main() // ── CPU 参考 ───────────────────────────────────────────────────────── auto cpu_owner = nn::CpuEngine::create(); nn::CpuEngine& cpu = *cpu_owner; - nn::Tensor x = cpu.create_tensor(4, 3, nn::Precision::F32, - nn::InitSpec::constant(0.5f)); - nn::Tensor y = cpu.create_tensor(4, 3, nn::Precision::F32, - nn::InitSpec::constant(2.0f)); + nn::Tensor x = cpu.create_tensor(4, 3, nn::Precision::F32, nn::InitSpec::constant(0.5f)); + nn::Tensor y = cpu.create_tensor(4, 3, nn::Precision::F32, nn::InitSpec::constant(2.0f)); auto cout_r = nn_example::fused_custom_op(cpu, x, y, 0.7f); if (!cout_r) { - std::fprintf(stderr, "[example] CPU 求值失败: %s\n", - cout_r.error().message.c_str()); + std::fprintf(stderr, "[example] CPU 求值失败: %s\n", cout_r.error().message.c_str()); return 1; } std::vector cbuf(12, 0.0f); @@ -68,8 +66,7 @@ int main() nn::GpuBackend& backend = *backend_owner; if (auto ir = backend.initialize(); !ir) { - std::fprintf(stderr, "[example] GPU 初始化失败（跳过）：%s\n", - ir.error().message.c_str()); + std::fprintf(stderr, "[example] GPU 初始化失败（跳过）：%s\n", ir.error().message.c_str()); return 77; // ctest 约定：77 = skip } auto gpu_owner = nn::GpuEngine::create(backend); @@ -77,10 +74,8 @@ int main() // ③ f32 路径（注册表命中 V0） { - nn::Tensor gx = gpu.create_tensor(4, 3, nn::Precision::F32, - nn::InitSpec::constant(0.5f)); - nn::Tensor gy = gpu.create_tensor(4, 3, nn::Precision::F32, - nn::InitSpec::constant(2.0f)); + nn::Tensor gx = gpu.create_tensor(4, 3, nn::Precision::F32, nn::InitSpec::constant(0.5f)); + nn::Tensor gy = gpu.create_tensor(4, 3, nn::Precision::F32, nn::InitSpec::constant(2.0f)); auto g_r = nn_example::fused_custom_op(gpu, gx, gy, 0.7f); if (!g_r) { @@ -94,8 +89,7 @@ int main() std::fprintf(stderr, "[example] GPU f32 read 失败\n"); return 1; } - std::printf("[example] GPU f32 (V0) = %.7f（与 CPU 差 %.3e）\n", - static_cast(gb[0]), + std::printf("[example] GPU f32 (V0) = %.7f（与 CPU 差 %.3e）\n", static_cast(gb[0]), static_cast(abs_diff(gb[0], cpu_val))); if (abs_diff(gb[0], cpu_val) > 1e-5f) { @@ -106,16 +100,12 @@ int main() // ② f16 路径（注册表命中 V1 运行期精度分派：双视图 + PC `prec`） { - nn::Tensor gx = gpu.create_tensor(4, 3, nn::Precision::F16, - nn::InitSpec::constant(0.5f)); - nn::Tensor gy = gpu.create_tensor(4, 3, nn::Precision::F16, - nn::InitSpec::constant(2.0f)); - auto g_r = nn_example::fused_custom_op(gpu, gx, gy, 0.7f, - nn::Precision::F16); + nn::Tensor gx = gpu.create_tensor(4, 3, nn::Precision::F16, nn::InitSpec::constant(0.5f)); + nn::Tensor gy = gpu.create_tensor(4, 3, nn::Precision::F16, nn::InitSpec::constant(2.0f)); + auto g_r = nn_example::fused_custom_op(gpu, gx, gy, 0.7f, nn::Precision::F16); if (!g_r) { - std::fprintf(stderr, - "[example] GPU f16 分派变体未命中（闭合世界）: %s\n", + std::fprintf(stderr, "[example] GPU f16 分派变体未命中（闭合世界）: %s\n", g_r.error().message.c_str()); return 1; } @@ -127,8 +117,7 @@ int main() } const float gv = static_cast(hb[0]); std::printf("[example] GPU f16 (V1) = %.7f（与 CPU 差 %.3e，f16 容差内）\n", - static_cast(gv), - static_cast(abs_diff(gv, cpu_val))); + static_cast(gv), static_cast(abs_diff(gv, cpu_val))); if (abs_diff(gv, cpu_val) > 2e-3f) { std::fprintf(stderr, "[example] GPU f16 与 CPU 不符\n"); @@ -140,15 +129,12 @@ int main() // 本层没有 dry-run —— 注册表若只靠锚点而符号实例与运行期不一致， // 这里会闭合世界硬报错。 { - nn::Tensor gx = gpu.create_tensor(4, 3, nn::Precision::F16, - nn::InitSpec::constant(0.5f)); - auto r = nn_example::fused_custom_norm(gpu, gx, 1.0f / 3.0f, 1e-5f, - nn::Precision::F16); + nn::Tensor gx = gpu.create_tensor(4, 3, nn::Precision::F16, nn::InitSpec::constant(0.5f)); + auto r = nn_example::fused_custom_norm(gpu, gx, 1.0f / 3.0f, 1e-5f, nn::Precision::F16); if (!r) { - std::fprintf(stderr, - "[example] 归一化形态 GPU 未命中（闭合世界）：%s\n", - r.error().message.c_str()); + std::fprintf(stderr, "[example] 归一化形态 GPU 未命中（闭合世界）：%s\n", + r.error().message.c_str()); return 1; } std::vector hb(12, nn::f16(0.0f)); diff --git a/extern/base b/extern/base index 47747f81..19264645 160000 --- a/extern/base +++ b/extern/base @@ -1 +1 @@ -Subproject commit 47747f81a40c8813a0de04925a8950f4c8ddeca5 +Subproject commit 19264645222bc5eac4e0daf51bba1736ca5db918 diff --git a/extern/dataset b/extern/dataset index 2966d4b0..e58c91f1 160000 --- a/extern/dataset +++ b/extern/dataset @@ -1 +1 @@ -Subproject commit 2966d4b0bb6b979a6fbbce04ab8db7d35062141f +Subproject commit e58c91f1b648b6667eb765be9b8a808a6364b28f diff --git a/extern/kvrecord b/extern/kvrecord index 084eb489..d043582f 160000 --- a/extern/kvrecord +++ b/extern/kvrecord @@ -1 +1 @@ -Subproject commit 084eb48953b3cac7d80f5a627eb3ce25bd1a5a9a +Subproject commit d043582f0441ec09db3b2618aa3af53ddda45d14 diff --git a/extern/threadpool b/extern/threadpool index 37820e03..657d266a 160000 --- a/extern/threadpool +++ b/extern/threadpool @@ -1 +1 @@ -Subproject commit 37820e033c5482727169fc3aa4298941c4f2f791 +Subproject commit 657d266af0551e0e9504988bfbadfdb7f229707d diff --git a/extern/tokenizer b/extern/tokenizer index 5e0e53f8..eefcc14d 160000 --- a/extern/tokenizer +++ b/extern/tokenizer @@ -1 +1 @@ -Subproject commit 5e0e53f81e71b6db045dd2bb1f8493ec3f68294a +Subproject commit eefcc14d3650e6c0dba8521fec903d75fc58fa46 diff --git a/include/neuralnet.cpp/algebra_matrix.hpp b/include/neuralnet.cpp/algebra_matrix.hpp index 2594b2bb..2dff8347 100644 --- a/include/neuralnet.cpp/algebra_matrix.hpp +++ b/include/neuralnet.cpp/algebra_matrix.hpp @@ -33,933 +33,900 @@ namespace nn { - // ── 宿主直配账本（M5，docs/internal/17 §4.5 D5）────────────────────── - // CPU 侧没有内存池：`MatrixT` 的存储走 `std::make_unique_for_overwrite` - // 直配。为了让 `ComputeEngine::pool_stats()` 在**所有**引擎上语义完整 - // （17 §2 目标 4"内存契约统一"），这里把每一次宿主分配/释放记账： - // 活分配数/字节 + 峰值 + 累计——即"直配模式"的池统计口径。 - // 只在分配/释放路径各做常数次 relaxed 原子操作，元素热循环零参与。 - struct HostAllocLedger - { - std::atomic live_blocks{0}; // 当前活分配数 - std::atomic live_bytes{0}; // 当前活字节 - std::atomic peak_bytes{0}; // 活字节历史峰值 - std::atomic total_blocks{0}; // 累计分配数 - std::atomic total_bytes{0}; // 累计分配字节 - - void on_alloc(std::size_t bytes) noexcept - { - const std::uint64_t b = static_cast(bytes); - live_blocks.fetch_add(1, std::memory_order_relaxed); - const std::uint64_t now = - live_bytes.fetch_add(b, std::memory_order_relaxed) + b; - std::uint64_t pk = peak_bytes.load(std::memory_order_relaxed); - while (now > pk && !peak_bytes.compare_exchange_weak( - pk, now, std::memory_order_relaxed, - std::memory_order_relaxed)) - { - } - total_blocks.fetch_add(1, std::memory_order_relaxed); - total_bytes.fetch_add(b, std::memory_order_relaxed); - } - void on_free(std::size_t bytes) noexcept - { - live_blocks.fetch_sub(1, std::memory_order_relaxed); - live_bytes.fetch_sub(static_cast(bytes), - std::memory_order_relaxed); - } - }; +// ── 宿主直配账本（M5，docs/internal/17 §4.5 D5）────────────────────── +// CPU 侧没有内存池：`MatrixT` 的存储走 `std::make_unique_for_overwrite` +// 直配。为了让 `ComputeEngine::pool_stats()` 在**所有**引擎上语义完整 +// （17 §2 目标 4"内存契约统一"），这里把每一次宿主分配/释放记账： +// 活分配数/字节 + 峰值 + 累计——即"直配模式"的池统计口径。 +// 只在分配/释放路径各做常数次 relaxed 原子操作，元素热循环零参与。 +struct HostAllocLedger +{ + std::atomic live_blocks{0}; // 当前活分配数 + std::atomic live_bytes{0}; // 当前活字节 + std::atomic peak_bytes{0}; // 活字节历史峰值 + std::atomic total_blocks{0}; // 累计分配数 + std::atomic total_bytes{0}; // 累计分配字节 - // 全进程唯一账本（header-only：inline 函数内静态对象跨翻译单元同一实例） - [[nodiscard]] inline HostAllocLedger& host_alloc_ledger() noexcept + void on_alloc(std::size_t bytes) noexcept { - static HostAllocLedger led; - return led; - } - - // ═══════════════════════════════════════════════════════════════════════ - // Matrix 类 — 矩阵存储与运算原语 - // - // 错误处理约定（与上层 L2 计算层不同）： - // - L1 代数层用 **NN_ASSERT_ALWAYS / NN_FATAL** 做形状与索引校验 - // （编程错误检查，Release 也终止——RV-01：断言后仍会按错误形状索引 - // = OOB，NN_ASSERT 在 NDEBUG 下展开为空会静默 UB）。 - // NN_ASSERT（可选诊断）只保留给"失败不影响后续正确性"的检查。 - // - L2+ 计算层使用 Result 进行运行时错误处理（用户输入校验） - // - 原因：Matrix 是底层原语，一次 size_t 比较不影响 Release 性能； - // 上层需要向用户报告错误，故使用 Result - // - 参见 docs/contributors/cpp-standards.md "分层职责单一" 章节 - // ═══════════════════════════════════════════════════════════════════════ - // ── MatrixT

— 类型化矩阵存储（docs/development/05-mixed-precision.md D2：存储精度是 P）── - // element = elem

（f16 / f32）；F32 实例即按 float 连续存储的普通矩阵。 - // F16 实例遵循 §7.2 形式化定义： - // matmul / 归约：f16 载入 / f32 累加 / f16 写出（acc_type = f32） - // 逐元素：f32 参考计算 + 每算子输出舍入到 f16（f16 运算符语义） - // 注意：C++ 不允许模板与同名的 using 别名共存，故模板名为 MatrixT， - // f32 便捷别名保持 Matrix。 - template - class MatrixT - { - public: - static constexpr Precision precision = P; - using element = elem

; - using acc_type = acc

; // F32: float（= element，现状）；F16: float（f32 累加） - - private: - // ── 存储：std::unique_ptr + 显式长度 ────────────────── - // 为何不用 std::vector：vector 的"按尺寸构造"必然值初始化（写满一遍零）， - // 而"输出会被完整覆盖"的路径（如 dsl::compute 的逐元素结果）不需要这一遍。 - // 实测本机单线程写满 1.57MB 要 0.50ms（~3.2 GB/s，内存带宽受限）， - // 是每个全尺寸中间量的纯浪费。std::make_unique_for_overwrite（C++20） - // 正是为"内存马上会被覆盖"设计。 - // 规范 §3.1：不出现 new[]/delete[]/裸指针所有权 —— unique_ptr 是 - // 规范推荐的替代；不穿透接口（storage 仅本类可见）。 - std::unique_ptr data_{}; - std::size_t size_{0}; - std::size_t rows_{0}; - std::size_t cols_{0}; - - // ── 写版本（write epoch）──────────────────────────────────────────── - // 语义：本存储的内容**可能已经变了**的单调计数。唯一的推进入口是 - // ComputeEngine 的写入口护栏（`WriteEpochGuard`），即"经引擎媒介写入 - // 张量存储"这件事本身；宿主侧直接摸 span() 不算（铁律 #11 已禁止 - // 库内那样做）。 - // 为什么放在**存储对象**上而不是 Tensor 上：Tensor 的拷贝是别名 - // （共享同一 storage；GPU 还有 with_shape 的零拷贝重贴标签），版本号 - // 必须随数据走，否则"经别名写入"对原句柄不可见。 - // 为什么 uint64：量化缓存"变了才重算"的判据是**相等**比较，回绕才会 - // 出错；2^64 次写实际不可达（即便回绕，也只是多算一次量化，不产生 - // 错值）。 - std::uint64_t write_epoch_{0}; - - public: - [[nodiscard]] std::uint64_t write_epoch() const noexcept { return write_epoch_; } - - // 写版本只前进：整对象替换（拷贝/移动赋值）也必须抬到两端较大者之后， - // 否则"替换后恰好等于某个更早的缓存值"会让参数派生物静默用旧值。 - void mark_content_changed() noexcept { ++write_epoch_; } - void inherit_write_epoch_(const MatrixT& other) noexcept + const std::uint64_t b = static_cast(bytes); + live_blocks.fetch_add(1, std::memory_order_relaxed); + const std::uint64_t now = live_bytes.fetch_add(b, std::memory_order_relaxed) + b; + std::uint64_t pk = peak_bytes.load(std::memory_order_relaxed); + while (now > pk && !peak_bytes.compare_exchange_weak(pk, now, std::memory_order_relaxed, + std::memory_order_relaxed)) { - const std::uint64_t base = write_epoch_ > other.write_epoch_ ? write_epoch_ - : other.write_epoch_; - write_epoch_ = base + 1; } + total_blocks.fetch_add(1, std::memory_order_relaxed); + total_bytes.fetch_add(b, std::memory_order_relaxed); + } + void on_free(std::size_t bytes) noexcept + { + live_blocks.fetch_sub(1, std::memory_order_relaxed); + live_bytes.fetch_sub(static_cast(bytes), std::memory_order_relaxed); + } +}; - private: +// 全进程唯一账本（header-only：inline 函数内静态对象跨翻译单元同一实例） +[[nodiscard]] inline HostAllocLedger& host_alloc_ledger() noexcept +{ + static HostAllocLedger led; + return led; +} + +// ═══════════════════════════════════════════════════════════════════════ +// Matrix 类 — 矩阵存储与运算原语 +// +// 错误处理约定（与上层 L2 计算层不同）： +// - L1 代数层用 **NN_ASSERT_ALWAYS / NN_FATAL** 做形状与索引校验 +// （编程错误检查，Release 也终止——RV-01：断言后仍会按错误形状索引 +// = OOB，NN_ASSERT 在 NDEBUG 下展开为空会静默 UB）。 +// NN_ASSERT（可选诊断）只保留给"失败不影响后续正确性"的检查。 +// - L2+ 计算层使用 Result 进行运行时错误处理（用户输入校验） +// - 原因：Matrix 是底层原语，一次 size_t 比较不影响 Release 性能； +// 上层需要向用户报告错误，故使用 Result +// - 参见 docs/contributors/cpp-standards.md "分层职责单一" 章节 +// ═══════════════════════════════════════════════════════════════════════ +// ── MatrixT

— 类型化矩阵存储（docs/development/05-mixed-precision.md D2：存储精度是 P）── +// element = elem

（f16 / f32）；F32 实例即按 float 连续存储的普通矩阵。 +// F16 实例遵循 §7.2 形式化定义： +// matmul / 归约：f16 载入 / f32 累加 / f16 写出（acc_type = f32） +// 逐元素：f32 参考计算 + 每算子输出舍入到 f16（f16 运算符语义） +// 注意：C++ 不允许模板与同名的 using 别名共存，故模板名为 MatrixT， +// f32 便捷别名保持 Matrix。 +template class MatrixT +{ +public: + static constexpr Precision precision = P; + using element = elem

; + using acc_type = acc

; // F32: float（= element，现状）；F16: float（f32 累加） + +private: + // ── 存储：std::unique_ptr + 显式长度 ────────────────── + // 为何不用 std::vector：vector 的"按尺寸构造"必然值初始化（写满一遍零）， + // 而"输出会被完整覆盖"的路径（如 dsl::compute 的逐元素结果）不需要这一遍。 + // 实测本机单线程写满 1.57MB 要 0.50ms（~3.2 GB/s，内存带宽受限）， + // 是每个全尺寸中间量的纯浪费。std::make_unique_for_overwrite（C++20） + // 正是为"内存马上会被覆盖"设计。 + // 规范 §3.1：不出现 new[]/delete[]/裸指针所有权 —— unique_ptr 是 + // 规范推荐的替代；不穿透接口（storage 仅本类可见）。 + std::unique_ptr data_{}; + std::size_t size_{0}; + std::size_t rows_{0}; + std::size_t cols_{0}; + + // ── 写版本（write epoch）──────────────────────────────────────────── + // 语义：本存储的内容**可能已经变了**的单调计数。唯一的推进入口是 + // ComputeEngine 的写入口护栏（`WriteEpochGuard`），即"经引擎媒介写入 + // 张量存储"这件事本身；宿主侧直接摸 span() 不算（铁律 #11 已禁止 + // 库内那样做）。 + // 为什么放在**存储对象**上而不是 Tensor 上：Tensor 的拷贝是别名 + // （共享同一 storage；GPU 还有 with_shape 的零拷贝重贴标签），版本号 + // 必须随数据走，否则"经别名写入"对原句柄不可见。 + // 为什么 uint64：量化缓存"变了才重算"的判据是**相等**比较，回绕才会 + // 出错；2^64 次写实际不可达（即便回绕，也只是多算一次量化，不产生 + // 错值）。 + std::uint64_t write_epoch_{0}; + +public: + [[nodiscard]] std::uint64_t write_epoch() const noexcept { return write_epoch_; } + + // 写版本只前进：整对象替换（拷贝/移动赋值）也必须抬到两端较大者之后， + // 否则"替换后恰好等于某个更早的缓存值"会让参数派生物静默用旧值。 + void mark_content_changed() noexcept { ++write_epoch_; } + void inherit_write_epoch_(const MatrixT& other) noexcept + { + const std::uint64_t base = + write_epoch_ > other.write_epoch_ ? write_epoch_ : other.write_epoch_; + write_epoch_ = base + 1; + } - [[nodiscard]] constexpr std::size_t index(std::size_t row, std::size_t col) const noexcept - { - return row * cols_ + col; - } +private: + [[nodiscard]] constexpr std::size_t index(std::size_t row, std::size_t col) const noexcept + { return row * cols_ + col; } - // 只分配、不初始化（元素值不确定；调用方必须保证写满后才读） - [[nodiscard]] static std::unique_ptr allocate_(std::size_t n) - { - if (n == 0) return {}; - auto p = std::make_unique_for_overwrite(n); - host_alloc_ledger().on_alloc(n * sizeof(element)); // 直配账本（M5） - return p; - } - // 释放并记账（直配账本的唯一出口；data_ 与 size_ 同步归零） - void release_storage_() noexcept - { - if (!data_) return; - host_alloc_ledger().on_free(size_ * sizeof(element)); - data_.reset(); - size_ = 0; - } - [[nodiscard]] const element *ptr() const noexcept { return data_.get(); } - [[nodiscard]] element *ptr() noexcept { return data_.get(); } + // 只分配、不初始化（元素值不确定；调用方必须保证写满后才读） + [[nodiscard]] static std::unique_ptr allocate_(std::size_t n) + { + if (n == 0) return {}; + auto p = std::make_unique_for_overwrite(n); + host_alloc_ledger().on_alloc(n * sizeof(element)); // 直配账本（M5） + return p; + } + // 释放并记账（直配账本的唯一出口；data_ 与 size_ 同步归零） + void release_storage_() noexcept + { + if (!data_) return; + host_alloc_ledger().on_free(size_ * sizeof(element)); + data_.reset(); + size_ = 0; + } + [[nodiscard]] const element* ptr() const noexcept { return data_.get(); } + [[nodiscard]] element* ptr() noexcept { return data_.get(); } - static void require_same_shape(const MatrixT &lhs, const MatrixT &rhs, std::string_view message) - { - if (lhs.rows_ != rhs.rows_ || lhs.cols_ != rhs.cols_) - { - // **绝不静默 OOB**（RV-01）：形状不匹配在 L1 是编程错误 - // （用户输入已在引擎入口经 Result 校验，见 ComputeEngine::add_inplace）， - // 此处 fail-fast 中止，代价为两次 size_t 比较，不影响热路径。 - // 统一走 NN_FATAL（= 无条件终止 + file:line），不再手抄三段式； - // 详情经 snprintf 组装（message 是 string_view，**不保证** NUL 结尾， - // 不能直接当 %s 实参）。 - char detail[192]; - std::snprintf(detail, sizeof(detail), - "Matrix 形状不匹配: %.*s (%zux%zu vs %zux%zu)", - static_cast(message.size()), message.data(), - lhs.rows_, lhs.cols_, rhs.rows_, rhs.cols_); - NN_FATAL(detail); - } + static void require_same_shape(const MatrixT& lhs, const MatrixT& rhs, std::string_view message) + { + if (lhs.rows_ != rhs.rows_ || lhs.cols_ != rhs.cols_) + { + // **绝不静默 OOB**（RV-01）：形状不匹配在 L1 是编程错误 + // （用户输入已在引擎入口经 Result 校验，见 ComputeEngine::add_inplace）， + // 此处 fail-fast 中止，代价为两次 size_t 比较，不影响热路径。 + // 统一走 NN_FATAL（= 无条件终止 + file:line），不再手抄三段式； + // 详情经 snprintf 组装（message 是 string_view，**不保证** NUL 结尾， + // 不能直接当 %s 实参）。 + char detail[192]; + std::snprintf(detail, sizeof(detail), "Matrix 形状不匹配: %.*s (%zux%zu vs %zux%zu)", + static_cast(message.size()), message.data(), lhs.rows_, lhs.cols_, + rhs.rows_, rhs.cols_); + NN_FATAL(detail); } + } - public: - MatrixT() = default; +public: + MatrixT() = default; - // 零填充构造（保持既有语义：默认构造出的矩阵全零） - explicit MatrixT(std::size_t rows, std::size_t cols) - : data_(allocate_(rows * cols)), size_(rows * cols), rows_(rows), cols_(cols) - { - if (size_ != 0) - std::fill_n(ptr(), size_, element{}); - } + // 零填充构造（保持既有语义：默认构造出的矩阵全零） + explicit MatrixT(std::size_t rows, std::size_t cols) + : data_(allocate_(rows * cols)), size_(rows * cols), rows_(rows), cols_(cols) + { + if (size_ != 0) std::fill_n(ptr(), size_, element{}); + } - // 从标量值初始化矩阵（host 标量为 f32；F16 实例构造时舍入到 f16） - MatrixT(std::size_t rows, std::size_t cols, element value) - : data_(allocate_(rows * cols)), size_(rows * cols), rows_(rows), cols_(cols) - { - if (size_ != 0) - std::fill_n(ptr(), size_, value); - } + // 从标量值初始化矩阵（host 标量为 f32；F16 实例构造时舍入到 f16） + MatrixT(std::size_t rows, std::size_t cols, element value) + : data_(allocate_(rows * cols)), size_(rows * cols), rows_(rows), cols_(cols) + { + if (size_ != 0) std::fill_n(ptr(), size_, value); + } - // ── 未初始化构造（输出会被完整覆盖的路径专用）────────────────────── - // 语义契约：调用方必须在任何读取之前把**全部** size() 个元素写满。 - // 用于把"分配 + 写满零 + 马上全覆盖"里的那一遍零写掉。 - struct uninitialized_tag {}; - MatrixT(std::size_t rows, std::size_t cols, uninitialized_tag) - : data_(allocate_(rows * cols)), size_(rows * cols), rows_(rows), cols_(cols) {} + // ── 未初始化构造（输出会被完整覆盖的路径专用）────────────────────── + // 语义契约：调用方必须在任何读取之前把**全部** size() 个元素写满。 + // 用于把"分配 + 写满零 + 马上全覆盖"里的那一遍零写掉。 + struct uninitialized_tag + { + }; + MatrixT(std::size_t rows, std::size_t cols, uninitialized_tag) + : data_(allocate_(rows * cols)), size_(rows * cols), rows_(rows), cols_(cols) + { + } - [[nodiscard]] static MatrixT make_uninitialized(std::size_t rows, std::size_t cols) - { - return MatrixT(rows, cols, uninitialized_tag{}); - } + [[nodiscard]] static MatrixT make_uninitialized(std::size_t rows, std::size_t cols) + { return MatrixT(rows, cols, uninitialized_tag{}); } - // ── 拷贝（深拷贝）/ 移动（指针转移）──────────────────────────────── - // 不能用默认实现：unique_ptr 不可拷贝，且默认移动会把 size_ 留在被移对象里 - // （表现为"size 非零但 data 为空"的悬空视图）。 - MatrixT(const MatrixT &other) - : data_(allocate_(other.size_)), size_(other.size_), - rows_(other.rows_), cols_(other.cols_) + // ── 拷贝（深拷贝）/ 移动（指针转移）──────────────────────────────── + // 不能用默认实现：unique_ptr 不可拷贝，且默认移动会把 size_ 留在被移对象里 + // （表现为"size 非零但 data 为空"的悬空视图）。 + MatrixT(const MatrixT& other) + : data_(allocate_(other.size_)), size_(other.size_), rows_(other.rows_), cols_(other.cols_) + { + if (size_ != 0) std::copy_n(other.ptr(), size_, ptr()); + } + MatrixT(MatrixT&& other) noexcept + : data_(std::move(other.data_)), size_(other.size_), rows_(other.rows_), cols_(other.cols_), + write_epoch_(other.write_epoch_) + { + other.size_ = 0; + other.rows_ = 0; + other.cols_ = 0; + } + MatrixT& operator=(const MatrixT& other) + { + if (this != &other) { - if (size_ != 0) - std::copy_n(other.ptr(), size_, ptr()); + MatrixT tmp(other); + *this = std::move(tmp); } - MatrixT(MatrixT &&other) noexcept - : data_(std::move(other.data_)), size_(other.size_), - rows_(other.rows_), cols_(other.cols_), - write_epoch_(other.write_epoch_) - { + return *this; + } + MatrixT& operator=(MatrixT&& other) noexcept + { + if (this != &other) + { + release_storage_(); // 旧缓冲先记账释放（M5 直配账本） + data_ = std::move(other.data_); + size_ = other.size_; + rows_ = other.rows_; + cols_ = other.cols_; + inherit_write_epoch_(other); // 只前进（见成员注释） other.size_ = 0; other.rows_ = 0; other.cols_ = 0; } - MatrixT &operator=(const MatrixT &other) - { - if (this != &other) - { - MatrixT tmp(other); - *this = std::move(tmp); - } - return *this; - } - MatrixT &operator=(MatrixT &&other) noexcept - { - if (this != &other) - { - release_storage_(); // 旧缓冲先记账释放（M5 直配账本） - data_ = std::move(other.data_); - size_ = other.size_; - rows_ = other.rows_; - cols_ = other.cols_; - inherit_write_epoch_(other); // 只前进（见成员注释） - other.size_ = 0; - other.rows_ = 0; - other.cols_ = 0; - } - return *this; - } - ~MatrixT() { release_storage_(); } + return *this; + } + ~MatrixT() { release_storage_(); } - // ── 就地调整大小（复用已有内存；新增部分零填充，与原 vector 语义一致）── - void resize(std::size_t rows, std::size_t cols) - { - if (rows_ == rows && cols_ == cols) return; // 尺寸不变，零开销 - const std::size_t n = rows * cols; - auto nd = allocate_(n); - if (nd && size_ != 0) - std::copy_n(ptr(), std::min(size_, n), nd.get()); - if (nd && n > size_) - std::fill_n(nd.get() + size_, n - size_, element{}); - release_storage_(); // 旧缓冲记账释放（M5 直配账本） - data_ = std::move(nd); - size_ = n; - rows_ = rows; - cols_ = cols; - mark_content_changed(); // 内容已变（见 write_epoch_ 注释） - } + // ── 就地调整大小（复用已有内存；新增部分零填充，与原 vector 语义一致）── + void resize(std::size_t rows, std::size_t cols) + { + if (rows_ == rows && cols_ == cols) return; // 尺寸不变，零开销 + const std::size_t n = rows * cols; + auto nd = allocate_(n); + if (nd && size_ != 0) std::copy_n(ptr(), std::min(size_, n), nd.get()); + if (nd && n > size_) std::fill_n(nd.get() + size_, n - size_, element{}); + release_storage_(); // 旧缓冲记账释放（M5 直配账本） + data_ = std::move(nd); + size_ = n; + rows_ = rows; + cols_ = cols; + mark_content_changed(); // 内容已变（见 write_epoch_ 注释） + } - // ── std::span 访问（C++20 现代接口，推荐使用） ──────────────────── - // 零开销抽象：编译后等价于裸指针 + 大小，可替代所有 data_ptr() 场景 - [[nodiscard]] std::span span() const - { - return {ptr(), size_}; - } - [[nodiscard]] std::span span() noexcept { return {ptr(), size_}; } - - // 访问器 - [[nodiscard]] constexpr std::size_t rows() const noexcept { return rows_; } - [[nodiscard]] constexpr std::size_t cols() const noexcept { return cols_; } - [[nodiscard]] constexpr std::size_t size() const noexcept { return size_; } - [[nodiscard]] constexpr bool empty() const noexcept { return size_ == 0; } - [[nodiscard]] element at(std::size_t row, std::size_t col) const - { - if (row >= rows_ || col >= cols_) - { - NN_FATAL("MatrixT::at index out of range"); - } - return ptr()[index(row, col)]; - } - void set_value(std::size_t row, std::size_t col, element value) - { - if (row >= rows_ || col >= cols_) - { - NN_FATAL("MatrixT::set_value index out of range"); - } - ptr()[index(row, col)] = value; - } - [[nodiscard]] element at_unchecked(std::size_t row, std::size_t col) const noexcept { return ptr()[index(row, col)]; } // 无校验 - void set_value_unchecked(std::size_t row, std::size_t col, element value) noexcept { ptr()[index(row, col)] = value; } // 无校验 + // ── std::span 访问（C++20 现代接口，推荐使用） ──────────────────── + // 零开销抽象：编译后等价于裸指针 + 大小，可替代所有 data_ptr() 场景 + [[nodiscard]] std::span span() const { return {ptr(), size_}; } + [[nodiscard]] std::span span() noexcept { return {ptr(), size_}; } + + // 访问器 + [[nodiscard]] constexpr std::size_t rows() const noexcept { return rows_; } + [[nodiscard]] constexpr std::size_t cols() const noexcept { return cols_; } + [[nodiscard]] constexpr std::size_t size() const noexcept { return size_; } + [[nodiscard]] constexpr bool empty() const noexcept { return size_ == 0; } + [[nodiscard]] element at(std::size_t row, std::size_t col) const + { + if (row >= rows_ || col >= cols_) { NN_FATAL("MatrixT::at index out of range"); } + return ptr()[index(row, col)]; + } + void set_value(std::size_t row, std::size_t col, element value) + { + if (row >= rows_ || col >= cols_) { NN_FATAL("MatrixT::set_value index out of range"); } + ptr()[index(row, col)] = value; + } + [[nodiscard]] element at_unchecked(std::size_t row, std::size_t col) const noexcept + { return ptr()[index(row, col)]; } // 无校验 + void set_value_unchecked(std::size_t row, std::size_t col, element value) noexcept + { ptr()[index(row, col)] = value; } // 无校验 // ── 转置（返回新矩阵） ───────────────────────────────────────────── - [[nodiscard]] MatrixT transpose() const - { - MatrixT result(cols_, rows_); - transpose_to(result); - return result; - } + [[nodiscard]] MatrixT transpose() const + { + MatrixT result(cols_, rows_); + transpose_to(result); + return result; + } - // ── 转置到预分配缓冲区（零分配热路径） ───────────────────────────── - void transpose_to(MatrixT &result) const - { - NN_ASSERT_ALWAYS(&result != this, "transpose_to: self-referencing not supported"); - result.resize(cols_, rows_); - if (rows_ == 0 || cols_ == 0) return; + // ── 转置到预分配缓冲区（零分配热路径） ───────────────────────────── + void transpose_to(MatrixT& result) const + { + NN_ASSERT_ALWAYS(&result != this, "transpose_to: self-referencing not supported"); + result.resize(cols_, rows_); + if (rows_ == 0 || cols_ == 0) return; - const auto src = span(); - auto dst = result.span(); - const std::size_t R = rows_; - const std::size_t C = cols_; + const auto src = span(); + auto dst = result.span(); + const std::size_t R = rows_; + const std::size_t C = cols_; - const std::size_t i_blocks = (R + BLOCK_SIZE - 1) / BLOCK_SIZE; - const std::size_t j_blocks = (C + BLOCK_SIZE - 1) / BLOCK_SIZE; + const std::size_t i_blocks = (R + BLOCK_SIZE - 1) / BLOCK_SIZE; + const std::size_t j_blocks = (C + BLOCK_SIZE - 1) / BLOCK_SIZE; - // 通过 nn::parallel_for_blocks 分发块级并行：分块数很少（1~几十），但其内循环计算量大。 - const auto n_blocks = i_blocks * j_blocks; - if (n_blocks <= 1) - { - for (std::size_t i = 0; i < R; ++i) - for (std::size_t j = 0; j < C; ++j) - dst[j * R + i] = src[i * C + j]; - } - else - { - auto block_indices = std::views::iota(std::size_t{0}, n_blocks); - nn::parallel_for_blocks( - block_indices.begin(), block_indices.end(), - [src, dst, R, C, j_blocks](std::size_t block_idx) noexcept - { - const std::size_t ib = block_idx / j_blocks; - const std::size_t jb = block_idx % j_blocks; - const std::size_t i0 = ib * BLOCK_SIZE; - const std::size_t j0 = jb * BLOCK_SIZE; - const std::size_t i1 = std::min(i0 + BLOCK_SIZE, R); - const std::size_t j1 = std::min(j0 + BLOCK_SIZE, C); - for (std::size_t i = i0; i < i1; ++i) - for (std::size_t j = j0; j < j1; ++j) - dst[j * R + i] = src[i * C + j]; - }); - } + // 通过 nn::parallel_for_blocks 分发块级并行：分块数很少（1~几十），但其内循环计算量大。 + const auto n_blocks = i_blocks * j_blocks; + if (n_blocks <= 1) + { + for (std::size_t i = 0; i < R; ++i) + for (std::size_t j = 0; j < C; ++j) + dst[j * R + i] = src[i * C + j]; } + else + { + auto block_indices = std::views::iota(std::size_t{0}, n_blocks); + nn::parallel_for_blocks(block_indices.begin(), block_indices.end(), + [src, dst, R, C, j_blocks](std::size_t block_idx) noexcept + { + const std::size_t ib = block_idx / j_blocks; + const std::size_t jb = block_idx % j_blocks; + const std::size_t i0 = ib * BLOCK_SIZE; + const std::size_t j0 = jb * BLOCK_SIZE; + const std::size_t i1 = std::min(i0 + BLOCK_SIZE, R); + const std::size_t j1 = std::min(j0 + BLOCK_SIZE, C); + for (std::size_t i = i0; i < i1; ++i) + for (std::size_t j = j0; j < j1; ++j) + dst[j * R + i] = src[i * C + j]; + }); + } + } - // ── GEMM 微内核：NI 行 × RB 列寄存器分块 ─────────────────────────── - // 为何这样写（实测依据）：若每个 (i,j) 只用一个标量累加器，每次 FMA 要 - // 2 次 load（受 load port 限制），且内层循环难以向量化。本微内核一次 - // 处理 4 行 × 8 列（一个 AVX2 向量宽度的列块）： - // · B 以 k-major 打包（b_pack[kk * j_len + j]）→ 固定 kk 时 j 连续， - // j 方向是一次连续向量 load； - // · 每个 k 只从 A 广播 4 个标量、做 4 次向量 FMA → FMA/load 比大幅改善； - // · k 维走**编译期常量** BLOCK_SIZE → 循环可完整向量化/展开。 - // 实测 GFLOPS（3072x768x512 / 768x3072x512 / 1024^3 / 512^3）： - // 401 / 407 / 354 / 281。 - // 累加分组固定（块内 kk 升序、块间 k_start 升序）→ 求和顺序确定， - // 结果可复现。 - static constexpr std::size_t MK_NI = 4; // 微内核行数 - static constexpr std::size_t MK_RB = 8; // 微内核列数 + // ── GEMM 微内核：NI 行 × RB 列寄存器分块 ─────────────────────────── + // 为何这样写（实测依据）：若每个 (i,j) 只用一个标量累加器，每次 FMA 要 + // 2 次 load（受 load port 限制），且内层循环难以向量化。本微内核一次 + // 处理 4 行 × 8 列（一个 AVX2 向量宽度的列块）： + // · B 以 k-major 打包（b_pack[kk * j_len + j]）→ 固定 kk 时 j 连续， + // j 方向是一次连续向量 load； + // · 每个 k 只从 A 广播 4 个标量、做 4 次向量 FMA → FMA/load 比大幅改善； + // · k 维走**编译期常量** BLOCK_SIZE → 循环可完整向量化/展开。 + // 实测 GFLOPS（3072x768x512 / 768x3072x512 / 1024^3 / 512^3）： + // 401 / 407 / 354 / 281。 + // 累加分组固定（块内 kk 升序、块间 k_start 升序）→ 求和顺序确定， + // 结果可复现。 + static constexpr std::size_t MK_NI = 4; // 微内核行数 + static constexpr std::size_t MK_RB = 8; // 微内核列数 // ap：A 行块首地址（第 ii 行为 ap + ii * a_stride）； // bp：k-major 打包块（步长 j_len）；rp：输出行块首地址（行步长 r_stride）。 // NI 用**模板参数**（编译期常量）：运行时传 NI 会让 ii 循环无法展开、 // 累加器可能落内存，实测主路径慢 ~1.7x。 - template - static void gemm_microkernel_(const element *ap, std::size_t a_stride, - const element *bp, std::size_t j_len, - element *rp, std::size_t r_stride) noexcept + template + static void gemm_microkernel_(const element* ap, std::size_t a_stride, const element* bp, + std::size_t j_len, element* rp, std::size_t r_stride) noexcept + { + std::array, NI> acc{}; + for (std::size_t kk = 0; kk < BLOCK_SIZE; ++kk) { - std::array, NI> acc{}; - for (std::size_t kk = 0; kk < BLOCK_SIZE; ++kk) - { - const element *brow = bp + kk * j_len; - for (std::size_t ii = 0; ii < NI; ++ii) - { - const acc_type av = static_cast(ap[ii * a_stride + kk]); - for (std::size_t jj = 0; jj < MK_RB; ++jj) - acc[ii][jj] += av * static_cast(brow[jj]); - } - } + const element* brow = bp + kk * j_len; for (std::size_t ii = 0; ii < NI; ++ii) - for (std::size_t jj = 0; jj < MK_RB; ++jj) - rp[ii * r_stride + jj] += static_cast(acc[ii][jj]); - } - - // ── 定长点积（尾块 / 窄 N 专用）───────────────────────────────── - // 必须独立成函数：NN_VECTORIZE_PRAGMA（loop vectorize(assume_safety)） - // 若与其它循环同处一个函数，会让整函数的向量化一起失败（实测同结构 - // 374 -> 95 GFLOPS）。尾块在窄 N（如 Linear 的 batch=1，N=1）时是唯一 - // 路径，这一段不隔离就会退化。 - static acc_type dot_span_(const element *x, std::size_t x_stride, - const element *y, std::size_t y_stride, - std::size_t len) noexcept - { - acc_type s = acc_type{0}; - NN_VECTORIZE_PRAGMA - for (std::size_t k = 0; k < len; ++k) - s += static_cast(x[k * x_stride]) * static_cast(y[k * y_stride]); - return s; - } - - // ── 基于 span 的矩阵乘法（零拷贝，供 batched_matmul 等场景使用） ── - // 从 a/b 的子区间直接计算，无需构造临时 Matrix 拷贝 - static void multiply_to_span( - std::span r, std::size_t M, std::size_t N, - std::span a, std::size_t /*a_rows*/, std::size_t a_cols, - std::span b, std::size_t b_rows, std::size_t b_cols) - { - // RV-01：内维不匹配后**仍会**按 K = a_cols 索引 B → 必须无条件终止 - NN_ASSERT_ALWAYS(a_cols == b_rows, "multiply_to_span: inner dimension mismatch"); - const std::size_t K = a_cols; - if (M == 0 || N == 0 || K == 0) return; - - // ── 窄 N 专用路径（N < MK_RB）──────────────────────────────── - // 典型是 Linear 的 batch=1（N=1）。主路径此处恒走尾块，且每个 - // (输出块, k 步) 都要把 B 打包进 16KB 栈块（值初始化）——GEMV 的 - // 计算量与之同量级，打包的零填充把内存流量翻倍。窄 N 下 B 很小 - // （N*K），直接按步长读、完全不打包更快。 - if (N < MK_RB) { - const std::size_t nb = (M + BLOCK_SIZE - 1) / BLOCK_SIZE; - nn::parallel_for_samples(nb, [a, b, r, M, N, K, b_cols](std::size_t ib) noexcept - { - const std::size_t i0 = ib * BLOCK_SIZE; - const std::size_t i1 = std::min(i0 + BLOCK_SIZE, M); - for (std::size_t i = i0; i < i1; ++i) - for (std::size_t j = 0; j < N; ++j) - r[i * N + j] += dot_span_(a.data() + i * K, 1, - b.data() + j, b_cols, K); - }); - return; + const acc_type av = static_cast(ap[ii * a_stride + kk]); + for (std::size_t jj = 0; jj < MK_RB; ++jj) + acc[ii][jj] += av * static_cast(brow[jj]); } + } + for (std::size_t ii = 0; ii < NI; ++ii) + for (std::size_t jj = 0; jj < MK_RB; ++jj) + rp[ii * r_stride + jj] += static_cast(acc[ii][jj]); + } - // 直接复用 blocked matmul 内核 - const std::size_t i_blocks = (M + BLOCK_SIZE - 1) / BLOCK_SIZE; - const std::size_t j_blocks = (N + BLOCK_SIZE - 1) / BLOCK_SIZE; - const auto n_blocks = i_blocks * j_blocks; + // ── 定长点积（尾块 / 窄 N 专用）───────────────────────────────── + // 必须独立成函数：NN_VECTORIZE_PRAGMA（loop vectorize(assume_safety)） + // 若与其它循环同处一个函数，会让整函数的向量化一起失败（实测同结构 + // 374 -> 95 GFLOPS）。尾块在窄 N（如 Linear 的 batch=1，N=1）时是唯一 + // 路径，这一段不隔离就会退化。 + static acc_type dot_span_(const element* x, std::size_t x_stride, const element* y, + std::size_t y_stride, std::size_t len) noexcept + { + acc_type s = acc_type{0}; + NN_VECTORIZE_PRAGMA + for (std::size_t k = 0; k < len; ++k) + s += static_cast(x[k * x_stride]) * static_cast(y[k * y_stride]); + return s; + } - auto kernel = [a, b, r, M, N, K, a_cols, b_cols, j_blocks](std::size_t block_idx) noexcept + // ── 基于 span 的矩阵乘法（零拷贝，供 batched_matmul 等场景使用） ── + // 从 a/b 的子区间直接计算，无需构造临时 Matrix 拷贝 + static void multiply_to_span(std::span r, std::size_t M, std::size_t N, + std::span a, std::size_t /*a_rows*/, + std::size_t a_cols, std::span b, std::size_t b_rows, + std::size_t b_cols) + { + // RV-01：内维不匹配后**仍会**按 K = a_cols 索引 B → 必须无条件终止 + NN_ASSERT_ALWAYS(a_cols == b_rows, "multiply_to_span: inner dimension mismatch"); + const std::size_t K = a_cols; + if (M == 0 || N == 0 || K == 0) return; + + // ── 窄 N 专用路径（N < MK_RB）──────────────────────────────── + // 典型是 Linear 的 batch=1（N=1）。主路径此处恒走尾块，且每个 + // (输出块, k 步) 都要把 B 打包进 16KB 栈块（值初始化）——GEMV 的 + // 计算量与之同量级，打包的零填充把内存流量翻倍。窄 N 下 B 很小 + // （N*K），直接按步长读、完全不打包更快。 + if (N < MK_RB) + { + const std::size_t nb = (M + BLOCK_SIZE - 1) / BLOCK_SIZE; + nn::parallel_for_samples(nb, + [a, b, r, M, N, K, b_cols](std::size_t ib) noexcept + { + const std::size_t i0 = ib * BLOCK_SIZE; + const std::size_t i1 = std::min(i0 + BLOCK_SIZE, M); + for (std::size_t i = i0; i < i1; ++i) + for (std::size_t j = 0; j < N; ++j) + r[i * N + j] += dot_span_(a.data() + i * K, 1, + b.data() + j, b_cols, K); + }); + return; + } + + // 直接复用 blocked matmul 内核 + const std::size_t i_blocks = (M + BLOCK_SIZE - 1) / BLOCK_SIZE; + const std::size_t j_blocks = (N + BLOCK_SIZE - 1) / BLOCK_SIZE; + const auto n_blocks = i_blocks * j_blocks; + + auto kernel = [a, b, r, M, N, K, a_cols, b_cols, j_blocks](std::size_t block_idx) noexcept + { + const std::size_t i_block = block_idx / j_blocks; + const std::size_t j_block = block_idx % j_blocks; + const std::size_t i_start = i_block * BLOCK_SIZE; + const std::size_t i_end = std::min(i_start + BLOCK_SIZE, M); + const std::size_t j_start = j_block * BLOCK_SIZE; + const std::size_t j_end = std::min(j_start + BLOCK_SIZE, N); + for (std::size_t k_start = 0; k_start < K; k_start += BLOCK_SIZE) { - const std::size_t i_block = block_idx / j_blocks; - const std::size_t j_block = block_idx % j_blocks; - const std::size_t i_start = i_block * BLOCK_SIZE; - const std::size_t i_end = std::min(i_start + BLOCK_SIZE, M); - const std::size_t j_start = j_block * BLOCK_SIZE; - const std::size_t j_end = std::min(j_start + BLOCK_SIZE, N); - for (std::size_t k_start = 0; k_start < K; k_start += BLOCK_SIZE) + const std::size_t k_end = std::min(k_start + BLOCK_SIZE, K); + const std::size_t k_len = k_end - k_start; + const std::size_t j_len = j_end - j_start; + // B 打包成 k-major：b_pack[kk * j_len + j]（固定 kk 时 j 连续） + std::array b_pack{}; + for (std::size_t kk = 0; kk < k_len; ++kk) + for (std::size_t j = 0; j < j_len; ++j) + b_pack[kk * j_len + j] = b[(k_start + kk) * b_cols + (j_start + j)]; + + for (std::size_t i0 = i_start; i0 < i_end; i0 += MK_NI) { - const std::size_t k_end = std::min(k_start + BLOCK_SIZE, K); - const std::size_t k_len = k_end - k_start; - const std::size_t j_len = j_end - j_start; - // B 打包成 k-major：b_pack[kk * j_len + j]（固定 kk 时 j 连续） - std::array b_pack{}; - for (std::size_t kk = 0; kk < k_len; ++kk) - for (std::size_t j = 0; j < j_len; ++j) - b_pack[kk * j_len + j] = b[(k_start + kk) * b_cols + (j_start + j)]; - - for (std::size_t i0 = i_start; i0 < i_end; i0 += MK_NI) + const std::size_t ni = std::min(MK_NI, i_end - i0); + const element* ap = a.data() + i0 * a_cols + k_start; + for (std::size_t j0 = 0; j0 < j_len; j0 += MK_RB) { - const std::size_t ni = std::min(MK_NI, i_end - i0); - const element *ap = a.data() + i0 * a_cols + k_start; - for (std::size_t j0 = 0; j0 < j_len; j0 += MK_RB) + const std::size_t nj = std::min(MK_RB, j_len - j0); + element* rp = r.data() + i0 * N + j_start + j0; + if (ni == MK_NI && nj == MK_RB && k_len == BLOCK_SIZE) { - const std::size_t nj = std::min(MK_RB, j_len - j0); - element *rp = r.data() + i0 * N + j_start + j0; - if (ni == MK_NI && nj == MK_RB && k_len == BLOCK_SIZE) - { - gemm_microkernel_(ap, a_cols, b_pack.data() + j0, - j_len, rp, N); - } - else - { - // 尾块：kk 放在**最内层** → 归约可向量化。 - // 对 N 很小（典型是 Linear 的 batch=1，N=1）这是关键路径： - // 若把 kk 放外层，jj 只有 1 次迭代，归约无法向量化 → 慢 2x。 - for (std::size_t ii = 0; ii < ni; ++ii) - for (std::size_t jj = 0; jj < nj; ++jj) - { - acc_type sum = acc_type{0}; - NN_VECTORIZE_PRAGMA - for (std::size_t kk = 0; kk < k_len; ++kk) - sum += static_cast(ap[ii * a_cols + kk]) - * static_cast(b_pack[kk * j_len + j0 + jj]); - rp[ii * N + jj] += static_cast(sum); - } - } + gemm_microkernel_(ap, a_cols, b_pack.data() + j0, j_len, rp, N); + } + else + { + // 尾块：kk 放在**最内层** → 归约可向量化。 + // 对 N 很小（典型是 Linear 的 batch=1，N=1）这是关键路径： + // 若把 kk 放外层，jj 只有 1 次迭代，归约无法向量化 → 慢 2x。 + for (std::size_t ii = 0; ii < ni; ++ii) + for (std::size_t jj = 0; jj < nj; ++jj) + { + acc_type sum = acc_type{0}; + NN_VECTORIZE_PRAGMA + for (std::size_t kk = 0; kk < k_len; ++kk) + sum += static_cast(ap[ii * a_cols + kk]) * + static_cast(b_pack[kk * j_len + j0 + jj]); + rp[ii * N + jj] += static_cast(sum); + } } } } - }; - - if (n_blocks <= 1) - kernel(0); - else - { - auto block_indices = std::views::iota(std::size_t{0}, n_blocks); - nn::parallel_for_blocks(block_indices.begin(), block_indices.end(), kernel); } - } + }; - // ── 基于 span 的矩阵乘法（B 转置，零拷贝） ───────────────────────── - static void multiply_transposed_to_span( - std::span r, std::size_t M, std::size_t N, - std::span a, std::size_t /*a_rows*/, std::size_t a_cols, - std::span bt, std::size_t /*bt_rows*/, std::size_t bt_cols) + if (n_blocks <= 1) kernel(0); + else { - // RV-01：同 multiply_to_span——内维不匹配即越界读，必须无条件终止 - NN_ASSERT_ALWAYS(a_cols == bt_cols, - "multiply_transposed_to_span: inner dimension mismatch"); - const std::size_t K = a_cols; - if (M == 0 || N == 0 || K == 0) return; - - // 窄 N 专用路径（见 multiply_to_span 处的说明）。此形态 B^T 存为 - // (N, K)，固定 n 时 k 连续 → 点积两侧都连续。 - if (N < MK_RB) - { - const std::size_t nb = (M + BLOCK_SIZE - 1) / BLOCK_SIZE; - nn::parallel_for_samples(nb, [a, bt, r, M, N, K](std::size_t ib) noexcept - { - const std::size_t i0 = ib * BLOCK_SIZE; - const std::size_t i1 = std::min(i0 + BLOCK_SIZE, M); - for (std::size_t i = i0; i < i1; ++i) - for (std::size_t j = 0; j < N; ++j) - r[i * N + j] += dot_span_(a.data() + i * K, 1, - bt.data() + j * K, 1, K); - }); - return; - } - - const std::size_t i_blocks = (M + BLOCK_SIZE - 1) / BLOCK_SIZE; - const std::size_t j_blocks = (N + BLOCK_SIZE - 1) / BLOCK_SIZE; - const auto n_blocks = i_blocks * j_blocks; + auto block_indices = std::views::iota(std::size_t{0}, n_blocks); + nn::parallel_for_blocks(block_indices.begin(), block_indices.end(), kernel); + } + } - auto kernel = [a, bt, r, M, N, K, j_blocks](std::size_t block_idx) noexcept + // ── 基于 span 的矩阵乘法（B 转置，零拷贝） ───────────────────────── + static void multiply_transposed_to_span(std::span r, std::size_t M, std::size_t N, + std::span a, std::size_t /*a_rows*/, + std::size_t a_cols, std::span bt, + std::size_t /*bt_rows*/, std::size_t bt_cols) + { + // RV-01：同 multiply_to_span——内维不匹配即越界读，必须无条件终止 + NN_ASSERT_ALWAYS(a_cols == bt_cols, + "multiply_transposed_to_span: inner dimension mismatch"); + const std::size_t K = a_cols; + if (M == 0 || N == 0 || K == 0) return; + + // 窄 N 专用路径（见 multiply_to_span 处的说明）。此形态 B^T 存为 + // (N, K)，固定 n 时 k 连续 → 点积两侧都连续。 + if (N < MK_RB) + { + const std::size_t nb = (M + BLOCK_SIZE - 1) / BLOCK_SIZE; + nn::parallel_for_samples(nb, + [a, bt, r, M, N, K](std::size_t ib) noexcept + { + const std::size_t i0 = ib * BLOCK_SIZE; + const std::size_t i1 = std::min(i0 + BLOCK_SIZE, M); + for (std::size_t i = i0; i < i1; ++i) + for (std::size_t j = 0; j < N; ++j) + r[i * N + j] += dot_span_(a.data() + i * K, 1, + bt.data() + j * K, 1, K); + }); + return; + } + + const std::size_t i_blocks = (M + BLOCK_SIZE - 1) / BLOCK_SIZE; + const std::size_t j_blocks = (N + BLOCK_SIZE - 1) / BLOCK_SIZE; + const auto n_blocks = i_blocks * j_blocks; + + auto kernel = [a, bt, r, M, N, K, j_blocks](std::size_t block_idx) noexcept + { + const std::size_t i_block = block_idx / j_blocks; + const std::size_t j_block = block_idx % j_blocks; + const std::size_t i_start = i_block * BLOCK_SIZE; + const std::size_t i_end = std::min(i_start + BLOCK_SIZE, M); + const std::size_t j_start = j_block * BLOCK_SIZE; + const std::size_t j_end = std::min(j_start + BLOCK_SIZE, N); + for (std::size_t k_start = 0; k_start < K; k_start += BLOCK_SIZE) { - const std::size_t i_block = block_idx / j_blocks; - const std::size_t j_block = block_idx % j_blocks; - const std::size_t i_start = i_block * BLOCK_SIZE; - const std::size_t i_end = std::min(i_start + BLOCK_SIZE, M); - const std::size_t j_start = j_block * BLOCK_SIZE; - const std::size_t j_end = std::min(j_start + BLOCK_SIZE, N); - for (std::size_t k_start = 0; k_start < K; k_start += BLOCK_SIZE) + const std::size_t k_len = std::min(BLOCK_SIZE, K - k_start); + const std::size_t j_len = j_end - j_start; + // B^T 打包成 k-major（固定 kk 时 j 连续）： + // b_pack[kk * j_len + j] = bt[(j_start + j) * K + (k_start + kk)] + std::array b_pack{}; + for (std::size_t kk = 0; kk < k_len; ++kk) + for (std::size_t j = 0; j < j_len; ++j) + b_pack[kk * j_len + j] = bt[(j_start + j) * K + (k_start + kk)]; + + for (std::size_t i0 = i_start; i0 < i_end; i0 += MK_NI) { - const std::size_t k_len = std::min(BLOCK_SIZE, K - k_start); - const std::size_t j_len = j_end - j_start; - // B^T 打包成 k-major（固定 kk 时 j 连续）： - // b_pack[kk * j_len + j] = bt[(j_start + j) * K + (k_start + kk)] - std::array b_pack{}; - for (std::size_t kk = 0; kk < k_len; ++kk) - for (std::size_t j = 0; j < j_len; ++j) - b_pack[kk * j_len + j] = bt[(j_start + j) * K + (k_start + kk)]; - - for (std::size_t i0 = i_start; i0 < i_end; i0 += MK_NI) + const std::size_t ni = std::min(MK_NI, i_end - i0); + const element* ap = a.data() + i0 * K + k_start; + for (std::size_t j0 = 0; j0 < j_len; j0 += MK_RB) { - const std::size_t ni = std::min(MK_NI, i_end - i0); - const element *ap = a.data() + i0 * K + k_start; - for (std::size_t j0 = 0; j0 < j_len; j0 += MK_RB) + const std::size_t nj = std::min(MK_RB, j_len - j0); + element* rp = r.data() + i0 * N + j_start + j0; + if (ni == MK_NI && nj == MK_RB && k_len == BLOCK_SIZE) { - const std::size_t nj = std::min(MK_RB, j_len - j0); - element *rp = r.data() + i0 * N + j_start + j0; - if (ni == MK_NI && nj == MK_RB && k_len == BLOCK_SIZE) - { - gemm_microkernel_(ap, K, b_pack.data() + j0, j_len, - rp, N); - } - else - { - // 尾块：kk 在最内层 → 归约可向量化 - for (std::size_t ii = 0; ii < ni; ++ii) - for (std::size_t jj = 0; jj < nj; ++jj) - { - acc_type sum = acc_type{0}; - NN_VECTORIZE_PRAGMA - for (std::size_t kk = 0; kk < k_len; ++kk) - sum += static_cast(ap[ii * K + kk]) - * static_cast(b_pack[kk * j_len + j0 + jj]); - rp[ii * N + jj] += static_cast(sum); - } - } + gemm_microkernel_(ap, K, b_pack.data() + j0, j_len, rp, N); + } + else + { + // 尾块：kk 在最内层 → 归约可向量化 + for (std::size_t ii = 0; ii < ni; ++ii) + for (std::size_t jj = 0; jj < nj; ++jj) + { + acc_type sum = acc_type{0}; + NN_VECTORIZE_PRAGMA + for (std::size_t kk = 0; kk < k_len; ++kk) + sum += static_cast(ap[ii * K + kk]) * + static_cast(b_pack[kk * j_len + j0 + jj]); + rp[ii * N + jj] += static_cast(sum); + } } } } - }; - - if (n_blocks <= 1) - kernel(0); - else - { - auto block_indices = std::views::iota(std::size_t{0}, n_blocks); - nn::parallel_for_blocks(block_indices.begin(), block_indices.end(), kernel); } - } + }; - // ── 基于 span 的矩阵乘法（A 转置，零拷贝） ───────────────────────── - static void transpose_multiply_to_span( - std::span r, std::size_t M, std::size_t N, - std::span a, std::size_t a_rows, std::size_t a_cols, - std::span b, std::size_t b_rows, std::size_t b_cols) + if (n_blocks <= 1) kernel(0); + else { - // RV-01：同 multiply_to_span——内维不匹配即越界读，必须无条件终止 - NN_ASSERT_ALWAYS(a_rows == b_rows, - "transpose_multiply_to_span: inner dimension mismatch"); - const std::size_t K = a_rows; // a is (K, M) stored - if (M == 0 || N == 0 || K == 0) return; - - // 窄 N 专用路径（见 multiply_to_span 处的说明）。此形态 A 存为 (K, M) - // → 固定 m 时以 a_cols 为步长读列（这是本形态的固有代价，无法避免）。 - if (N < MK_RB) - { - const std::size_t nb = (M + BLOCK_SIZE - 1) / BLOCK_SIZE; - nn::parallel_for_samples(nb, [a, b, r, M, N, K, a_cols, b_cols](std::size_t ib) noexcept - { - const std::size_t i0 = ib * BLOCK_SIZE; - const std::size_t i1 = std::min(i0 + BLOCK_SIZE, M); - for (std::size_t i = i0; i < i1; ++i) - for (std::size_t j = 0; j < N; ++j) - r[i * N + j] += dot_span_(a.data() + i, a_cols, - b.data() + j, b_cols, K); - }); - return; - } - - const std::size_t i_blocks = (M + BLOCK_SIZE - 1) / BLOCK_SIZE; - const std::size_t j_blocks = (N + BLOCK_SIZE - 1) / BLOCK_SIZE; - const auto n_blocks = i_blocks * j_blocks; + auto block_indices = std::views::iota(std::size_t{0}, n_blocks); + nn::parallel_for_blocks(block_indices.begin(), block_indices.end(), kernel); + } + } - auto kernel = [a, b, r, M, N, K, a_cols, b_cols, j_blocks](std::size_t block_idx) noexcept + // ── 基于 span 的矩阵乘法（A 转置，零拷贝） ───────────────────────── + static void transpose_multiply_to_span(std::span r, std::size_t M, std::size_t N, + std::span a, std::size_t a_rows, + std::size_t a_cols, std::span b, + std::size_t b_rows, std::size_t b_cols) + { + // RV-01：同 multiply_to_span——内维不匹配即越界读，必须无条件终止 + NN_ASSERT_ALWAYS(a_rows == b_rows, "transpose_multiply_to_span: inner dimension mismatch"); + const std::size_t K = a_rows; // a is (K, M) stored + if (M == 0 || N == 0 || K == 0) return; + + // 窄 N 专用路径（见 multiply_to_span 处的说明）。此形态 A 存为 (K, M) + // → 固定 m 时以 a_cols 为步长读列（这是本形态的固有代价，无法避免）。 + if (N < MK_RB) + { + const std::size_t nb = (M + BLOCK_SIZE - 1) / BLOCK_SIZE; + nn::parallel_for_samples(nb, + [a, b, r, M, N, K, a_cols, b_cols](std::size_t ib) noexcept + { + const std::size_t i0 = ib * BLOCK_SIZE; + const std::size_t i1 = std::min(i0 + BLOCK_SIZE, M); + for (std::size_t i = i0; i < i1; ++i) + for (std::size_t j = 0; j < N; ++j) + r[i * N + j] += dot_span_(a.data() + i, a_cols, + b.data() + j, b_cols, K); + }); + return; + } + + const std::size_t i_blocks = (M + BLOCK_SIZE - 1) / BLOCK_SIZE; + const std::size_t j_blocks = (N + BLOCK_SIZE - 1) / BLOCK_SIZE; + const auto n_blocks = i_blocks * j_blocks; + + auto kernel = [a, b, r, M, N, K, a_cols, b_cols, j_blocks](std::size_t block_idx) noexcept + { + const std::size_t i_block = block_idx / j_blocks; + const std::size_t j_block = block_idx % j_blocks; + const std::size_t i_start = i_block * BLOCK_SIZE; + const std::size_t i_end = std::min(i_start + BLOCK_SIZE, M); + const std::size_t j_start = j_block * BLOCK_SIZE; + const std::size_t j_end = std::min(j_start + BLOCK_SIZE, N); + for (std::size_t k_start = 0; k_start < K; k_start += BLOCK_SIZE) { - const std::size_t i_block = block_idx / j_blocks; - const std::size_t j_block = block_idx % j_blocks; - const std::size_t i_start = i_block * BLOCK_SIZE; - const std::size_t i_end = std::min(i_start + BLOCK_SIZE, M); - const std::size_t j_start = j_block * BLOCK_SIZE; - const std::size_t j_end = std::min(j_start + BLOCK_SIZE, N); - for (std::size_t k_start = 0; k_start < K; k_start += BLOCK_SIZE) - { - const std::size_t k_len = std::min(BLOCK_SIZE, K - k_start); - const std::size_t j_len = j_end - j_start; - const std::size_t i_len = i_end - i_start; - // A 存储为 (K, M)：按 (i, k) 打包成行内 k 连续的块， - // 微内核以 k_len 为行步长读它（避免 stride=a_cols 的跨行访问） - std::array a_pack{}; - for (std::size_t ii = 0; ii < i_len; ++ii) - for (std::size_t kk = 0; kk < k_len; ++kk) - a_pack[ii * k_len + kk] = a[(k_start + kk) * a_cols + (i_start + ii)]; - // B 打包成 k-major（固定 kk 时 j 连续） - std::array b_pack{}; + const std::size_t k_len = std::min(BLOCK_SIZE, K - k_start); + const std::size_t j_len = j_end - j_start; + const std::size_t i_len = i_end - i_start; + // A 存储为 (K, M)：按 (i, k) 打包成行内 k 连续的块， + // 微内核以 k_len 为行步长读它（避免 stride=a_cols 的跨行访问） + std::array a_pack{}; + for (std::size_t ii = 0; ii < i_len; ++ii) for (std::size_t kk = 0; kk < k_len; ++kk) - for (std::size_t j = 0; j < j_len; ++j) - b_pack[kk * j_len + j] = b[(k_start + kk) * b_cols + (j_start + j)]; - - for (std::size_t ii0 = 0; ii0 < i_len; ii0 += MK_NI) + a_pack[ii * k_len + kk] = a[(k_start + kk) * a_cols + (i_start + ii)]; + // B 打包成 k-major（固定 kk 时 j 连续） + std::array b_pack{}; + for (std::size_t kk = 0; kk < k_len; ++kk) + for (std::size_t j = 0; j < j_len; ++j) + b_pack[kk * j_len + j] = b[(k_start + kk) * b_cols + (j_start + j)]; + + for (std::size_t ii0 = 0; ii0 < i_len; ii0 += MK_NI) + { + const std::size_t ni = std::min(MK_NI, i_len - ii0); + const element* ap = a_pack.data() + ii0 * k_len; + element* rbase = r.data() + (i_start + ii0) * N + j_start; + for (std::size_t j0 = 0; j0 < j_len; j0 += MK_RB) { - const std::size_t ni = std::min(MK_NI, i_len - ii0); - const element *ap = a_pack.data() + ii0 * k_len; - element *rbase = r.data() + (i_start + ii0) * N + j_start; - for (std::size_t j0 = 0; j0 < j_len; j0 += MK_RB) + const std::size_t nj = std::min(MK_RB, j_len - j0); + element* rp = rbase + j0; + if (ni == MK_NI && nj == MK_RB && k_len == BLOCK_SIZE) { - const std::size_t nj = std::min(MK_RB, j_len - j0); - element *rp = rbase + j0; - if (ni == MK_NI && nj == MK_RB && k_len == BLOCK_SIZE) - { - gemm_microkernel_(ap, k_len, b_pack.data() + j0, j_len, - rp, N); - } - else - { - // 尾块：kk 在最内层 → 归约可向量化 - for (std::size_t ii = 0; ii < ni; ++ii) - for (std::size_t jj = 0; jj < nj; ++jj) - { - acc_type sum = acc_type{0}; - NN_VECTORIZE_PRAGMA - for (std::size_t kk = 0; kk < k_len; ++kk) - sum += static_cast(ap[ii * k_len + kk]) - * static_cast(b_pack[kk * j_len + j0 + jj]); - rp[ii * N + jj] += static_cast(sum); - } - } + gemm_microkernel_(ap, k_len, b_pack.data() + j0, j_len, rp, N); + } + else + { + // 尾块：kk 在最内层 → 归约可向量化 + for (std::size_t ii = 0; ii < ni; ++ii) + for (std::size_t jj = 0; jj < nj; ++jj) + { + acc_type sum = acc_type{0}; + NN_VECTORIZE_PRAGMA + for (std::size_t kk = 0; kk < k_len; ++kk) + sum += static_cast(ap[ii * k_len + kk]) * + static_cast(b_pack[kk * j_len + j0 + jj]); + rp[ii * N + jj] += static_cast(sum); + } } } } - }; - - if (n_blocks <= 1) - kernel(0); - else - { - auto block_indices = std::views::iota(std::size_t{0}, n_blocks); - nn::parallel_for_blocks(block_indices.begin(), block_indices.end(), kernel); } - } + }; - // ── 矩阵乘法到预分配缓冲区（零分配热路径） ───────────────────────── - // 使用 std::span（C++20）提供类型安全的非拥有视图 - void multiply_to(MatrixT &result, const MatrixT &other) const + if (n_blocks <= 1) kernel(0); + else { - NN_ASSERT_ALWAYS(&result != this, "multiply_to: self-referencing not supported"); - const std::size_t M = rows_; - const std::size_t N = other.cols_; - const std::size_t K = cols_; - result.resize(M, N); - if (M == 0 || N == 0 || K == 0) return; - - // 清零结果矩阵（使用 RAII 封装的方法） - result.zero(); - - // 委托给 span 版内核：GEMM 内核只保留一份，避免两处实现漂移 - multiply_to_span(result.span(), M, N, span(), M, K, other.span(), K, N); + auto block_indices = std::views::iota(std::size_t{0}, n_blocks); + nn::parallel_for_blocks(block_indices.begin(), block_indices.end(), kernel); } + } - // ── 矩阵乘法（B 转置）到预分配缓冲区 ───────────────────────── - // 计算 result = this * B^T，其中 B 存储为 (N, K) 行主序 - // 无需实际转置 B，直接从 B 的行读取列（节省 O(N*K) 的拷贝） - // - // 维度要求：this=(M,K), b_trans=(N,K) → result=(M,N) - // 即 C[m][n] = Σ_k A[m][k] * B[n][k] - // ── 矩阵乘法（B 转置）到预分配缓冲区 ───────────────────────── - // result = this * B^T，其中 B 存储为 (N, K) 行主序；C[m][n] = Σ_k A[m][k] * B[n][k]。 - // 委托给 span 版内核：GEMM 内核只保留一份，避免两处实现漂移。 - void multiply_transposed_to(MatrixT &result, const MatrixT &b_trans) const - { - NN_ASSERT_ALWAYS(&result != this && &result != &b_trans, - "multiply_transposed_to: self-referencing not supported"); - NN_ASSERT_ALWAYS(cols_ == b_trans.cols_, - "multiply_transposed_to: inner dimensions mismatch"); - const std::size_t M = rows_; - const std::size_t K = cols_; - const std::size_t N = b_trans.rows_; - result.resize(M, N); - if (M == 0 || N == 0 || K == 0) return; - result.zero(); - multiply_transposed_to_span(result.span(), M, N, span(), M, K, - b_trans.span(), N, K); - } + // ── 矩阵乘法到预分配缓冲区（零分配热路径） ───────────────────────── + // 使用 std::span（C++20）提供类型安全的非拥有视图 + void multiply_to(MatrixT& result, const MatrixT& other) const + { + NN_ASSERT_ALWAYS(&result != this, "multiply_to: self-referencing not supported"); + const std::size_t M = rows_; + const std::size_t N = other.cols_; + const std::size_t K = cols_; + result.resize(M, N); + if (M == 0 || N == 0 || K == 0) return; + + // 清零结果矩阵（使用 RAII 封装的方法） + result.zero(); + + // 委托给 span 版内核：GEMM 内核只保留一份，避免两处实现漂移 + multiply_to_span(result.span(), M, N, span(), M, K, other.span(), K, N); + } - // ── 矩阵乘法（A 转置）到预分配缓冲区 ───────────────────────── - // 计算 result = this^T * B，其中 this 存储为 (K, M) 行主序 - // 无需实际转置 this，直接从 this 的列读取行（节省 O(K*M) 的拷贝） - // - // 维度要求：this=(K,M), b=(K,N) → result=(M,N) - // 即 C[m][n] = Σ_k A[k][m] * B[k][n] - // 委托给 span 版内核：GEMM 内核只保留一份，避免两处实现漂移。 - void transpose_multiply_to(MatrixT &result, const MatrixT &b) const - { - NN_ASSERT_ALWAYS(&result != this && &result != &b, - "transpose_multiply_to: self-referencing not supported"); - NN_ASSERT_ALWAYS(rows_ == b.rows_, - "transpose_multiply_to: inner dimensions mismatch"); - const std::size_t K = rows_; // this is (K, M) - const std::size_t M = cols_; - const std::size_t N = b.cols_; - result.resize(M, N); - if (M == 0 || N == 0 || K == 0) return; - result.zero(); - transpose_multiply_to_span(result.span(), M, N, span(), K, cols_, - b.span(), K, b.cols()); - } + // ── 矩阵乘法（B 转置）到预分配缓冲区 ───────────────────────── + // 计算 result = this * B^T，其中 B 存储为 (N, K) 行主序 + // 无需实际转置 B，直接从 B 的行读取列（节省 O(N*K) 的拷贝） + // + // 维度要求：this=(M,K), b_trans=(N,K) → result=(M,N) + // 即 C[m][n] = Σ_k A[m][k] * B[n][k] + // ── 矩阵乘法（B 转置）到预分配缓冲区 ───────────────────────── + // result = this * B^T，其中 B 存储为 (N, K) 行主序；C[m][n] = Σ_k A[m][k] * B[n][k]。 + // 委托给 span 版内核：GEMM 内核只保留一份，避免两处实现漂移。 + void multiply_transposed_to(MatrixT& result, const MatrixT& b_trans) const + { + NN_ASSERT_ALWAYS(&result != this && &result != &b_trans, + "multiply_transposed_to: self-referencing not supported"); + NN_ASSERT_ALWAYS(cols_ == b_trans.cols_, + "multiply_transposed_to: inner dimensions mismatch"); + const std::size_t M = rows_; + const std::size_t K = cols_; + const std::size_t N = b_trans.rows_; + result.resize(M, N); + if (M == 0 || N == 0 || K == 0) return; + result.zero(); + multiply_transposed_to_span(result.span(), M, N, span(), M, K, b_trans.span(), N, K); + } - void scale_inplace(Scalar scalar) noexcept - { - auto s = span(); - nn::for_each(s.begin(), s.end(), - [scalar](element &value) noexcept { value *= scalar; }); - } + // ── 矩阵乘法（A 转置）到预分配缓冲区 ───────────────────────── + // 计算 result = this^T * B，其中 this 存储为 (K, M) 行主序 + // 无需实际转置 this，直接从 this 的列读取行（节省 O(K*M) 的拷贝） + // + // 维度要求：this=(K,M), b=(K,N) → result=(M,N) + // 即 C[m][n] = Σ_k A[k][m] * B[k][n] + // 委托给 span 版内核：GEMM 内核只保留一份，避免两处实现漂移。 + void transpose_multiply_to(MatrixT& result, const MatrixT& b) const + { + NN_ASSERT_ALWAYS(&result != this && &result != &b, + "transpose_multiply_to: self-referencing not supported"); + NN_ASSERT_ALWAYS(rows_ == b.rows_, "transpose_multiply_to: inner dimensions mismatch"); + const std::size_t K = rows_; // this is (K, M) + const std::size_t M = cols_; + const std::size_t N = b.cols_; + result.resize(M, N); + if (M == 0 || N == 0 || K == 0) return; + result.zero(); + transpose_multiply_to_span(result.span(), M, N, span(), K, cols_, b.span(), K, b.cols()); + } - // 逐元素加法 inplace - void add_inplace(const MatrixT &other) - { - require_same_shape(*this, other, "add_inplace dimension mismatch"); - auto s = span(); - auto o = other.span(); - nn::transform(s.begin(), s.end(), o.begin(), - s.begin(), std::plus<>{}); - } + void scale_inplace(Scalar scalar) noexcept + { + auto s = span(); + nn::for_each(s.begin(), s.end(), [scalar](element& value) noexcept { value *= scalar; }); + } - // 填充零 - void zero() noexcept - { - if (size_ != 0) - std::fill_n(ptr(), size_, element{}); - } + // 逐元素加法 inplace + void add_inplace(const MatrixT& other) + { + require_same_shape(*this, other, "add_inplace dimension mismatch"); + auto s = span(); + auto o = other.span(); + nn::transform(s.begin(), s.end(), o.begin(), s.begin(), std::plus<>{}); + } - // ── 归约操作 ──────────────────────────────────────────────────── - // result = reduce_op(init, transform_op(data[0]), transform_op(data[1]), ...) - template - [[nodiscard]] T reduce(T init, ReduceOp&& reduce_op, TransformOp&& transform_op) const - { - auto s = span(); - return nn::transform_reduce(s.begin(), s.end(), init, - std::forward(reduce_op), std::forward(transform_op)); - } + // 填充零 + void zero() noexcept + { + if (size_ != 0) std::fill_n(ptr(), size_, element{}); + } + + // ── 归约操作 ──────────────────────────────────────────────────── + // result = reduce_op(init, transform_op(data[0]), transform_op(data[1]), ...) + template + [[nodiscard]] T reduce(T init, ReduceOp&& reduce_op, TransformOp&& transform_op) const + { + auto s = span(); + return nn::transform_reduce(s.begin(), s.end(), init, std::forward(reduce_op), + std::forward(transform_op)); + } + + // ── 按行归约（通用数学原语，不是算法） ────────────────────────── + // 对每一行独立归约，返回 (rows, 1) 矩阵。 + // result[r][0] = reduce_op(init, transform_op(this[r][0]), ..., transform_op(this[r][cols-1])) + // 上层可基于此表达 softmax 行最大值/行求和、按行范数等算法。 + template + [[nodiscard]] MatrixT row_reduce(T init, ReduceOp&& reduce_op, TransformOp&& transform_op) const + { + MatrixT result(rows_, 1); + if (rows_ == 0) return result; - // ── 按行归约（通用数学原语，不是算法） ────────────────────────── - // 对每一行独立归约，返回 (rows, 1) 矩阵。 - // result[r][0] = reduce_op(init, transform_op(this[r][0]), ..., transform_op(this[r][cols-1])) - // 上层可基于此表达 softmax 行最大值/行求和、按行范数等算法。 - template - [[nodiscard]] MatrixT row_reduce(T init, ReduceOp&& reduce_op, TransformOp&& transform_op) const + const auto self = span(); + auto out = result.span(); + const std::size_t C = cols_; + + auto process_row = + [self, out, C, init, reduce_op = std::forward(reduce_op), + transform_op = std::forward(transform_op)](std::size_t r) noexcept { - MatrixT result(rows_, 1); - if (rows_ == 0) return result; + const auto row = self.subspan(r * C, C); + T acc = init; + for (std::size_t c = 0; c < C; ++c) + acc = reduce_op(acc, transform_op(row[c])); + out[r] = static_cast(acc); + }; + + // 并行门控按**元素数**（R*C），与 broadcast_* 一致：行数本身远低于 + // PARALLEL_THRESHOLD（512K），按行数比较会让本原语**恒定串行**， + // 故必须按 R*C 分片触发。 + // 每行独立累加、行内顺序不变 → 并行与串行逐字节一致（铁律 8）。 + if (rows_ * cols_ >= PARALLEL_THRESHOLD && rows_ > 1) + nn::parallel_for_samples(rows_, process_row); + else + for (std::size_t r = 0; r < rows_; ++r) + process_row(r); + return result; + } - const auto self = span(); - auto out = result.span(); - const std::size_t C = cols_; + // ── 按列归约（通用数学原语，不是算法） ────────────────────────── + // 对每一列独立归约，返回 (1, cols) 矩阵。 + // result[0][c] = reduce_op(init, transform_op(this[0][c]), ..., transform_op(this[rows-1][c])) + // 上层可基于此表达 LayerNorm 列均值/列方差等算法。 + // + // 实现策略：cache-friendly blocked + 行块并行。 + // 实测：blocked 全面优于 naive（按列跨行扫描）， + // 行块并行仅在 R >= 256 且 R*C >= PARALLEL_THRESHOLD 时启用。 + template + [[nodiscard]] MatrixT col_reduce(T init, ReduceOp&& reduce_op, TransformOp&& transform_op) const + { + MatrixT result(1, cols_); + if (cols_ == 0) return result; - auto process_row = [self, out, C, init, - reduce_op = std::forward(reduce_op), - transform_op = std::forward(transform_op)](std::size_t r) noexcept { - const auto row = self.subspan(r * C, C); + const auto self = span(); + auto out = result.span(); + const std::size_t R = rows_; + const std::size_t C = cols_; + + // 极小矩阵直接 naive（避免清零开销） + if (R * C < 64) + { + for (std::size_t c = 0; c < C; ++c) + { T acc = init; - for (std::size_t c = 0; c < C; ++c) - acc = reduce_op(acc, transform_op(row[c])); - out[r] = static_cast(acc); - }; - - // 并行门控按**元素数**（R*C），与 broadcast_* 一致：行数本身远低于 - // PARALLEL_THRESHOLD（512K），按行数比较会让本原语**恒定串行**， - // 故必须按 R*C 分片触发。 - // 每行独立累加、行内顺序不变 → 并行与串行逐字节一致（铁律 8）。 - if (rows_ * cols_ >= PARALLEL_THRESHOLD && rows_ > 1) - nn::parallel_for_samples(rows_, process_row); - else - for (std::size_t r = 0; r < rows_; ++r) - process_row(r); + for (std::size_t r = 0; r < R; ++r) + acc = reduce_op(acc, transform_op(self[r * C + c])); + out[c] = static_cast(acc); + } return result; } - // ── 按列归约（通用数学原语，不是算法） ────────────────────────── - // 对每一列独立归约，返回 (1, cols) 矩阵。 - // result[0][c] = reduce_op(init, transform_op(this[0][c]), ..., transform_op(this[rows-1][c])) - // 上层可基于此表达 LayerNorm 列均值/列方差等算法。 - // - // 实现策略：cache-friendly blocked + 行块并行。 - // 实测：blocked 全面优于 naive（按列跨行扫描）， - // 行块并行仅在 R >= 256 且 R*C >= PARALLEL_THRESHOLD 时启用。 - template - [[nodiscard]] MatrixT col_reduce(T init, ReduceOp&& reduce_op, TransformOp&& transform_op) const - { - MatrixT result(1, cols_); - if (cols_ == 0) return result; + // 行块并行启用条件：R >= COL_REDUCE_PARALLEL_ROWS 且 R*C >= PARALLEL_THRESHOLD。 + // 行数门槛取 256：R*C >= 512K 时即使 R=256 每线程也有 >=16K 元素 + // 的工作量（32 线程假设），同步开销不占主导；256 覆盖常见 d_model=768 场景。 + constexpr std::size_t COL_REDUCE_PARALLEL_ROWS = 256; // 行数门槛 + const std::size_t hw_threads = std::thread::hardware_concurrency(); + const std::size_t n_threads = (hw_threads == 0) ? 1 : hw_threads; + const bool use_parallel = + R >= COL_REDUCE_PARALLEL_ROWS && R * C >= PARALLEL_THRESHOLD && n_threads > 1; - const auto self = span(); - auto out = result.span(); - const std::size_t R = rows_; - const std::size_t C = cols_; - - // 极小矩阵直接 naive（避免清零开销） - if (R * C < 64) + if (!use_parallel) + { + // ── 单线程行主序扫描 ── + for (std::size_t c = 0; c < C; ++c) + out[c] = static_cast(init); + for (std::size_t r = 0; r < R; ++r) { + // 行视图：std::span::subspan 零成本（ptr+len），替代裸指针行起点 + const auto row = self.subspan(r * C, C); for (std::size_t c = 0; c < C; ++c) { - T acc = init; - for (std::size_t r = 0; r < R; ++r) - acc = reduce_op(acc, transform_op(self[r * C + c])); - out[c] = static_cast(acc); + element v = static_cast(transform_op(row[c])); + out[c] = static_cast(reduce_op(static_cast(out[c]), v)); } - return result; } + return result; + } - // 行块并行启用条件：R >= COL_REDUCE_PARALLEL_ROWS 且 R*C >= PARALLEL_THRESHOLD。 - // 行数门槛取 256：R*C >= 512K 时即使 R=256 每线程也有 >=16K 元素 - // 的工作量（32 线程假设），同步开销不占主导；256 覆盖常见 d_model=768 场景。 - constexpr std::size_t COL_REDUCE_PARALLEL_ROWS = 256; // 行数门槛 - const std::size_t hw_threads = std::thread::hardware_concurrency(); - const std::size_t n_threads = (hw_threads == 0) ? 1 : hw_threads; - const bool use_parallel = - R >= COL_REDUCE_PARALLEL_ROWS && - R * C >= PARALLEL_THRESHOLD && - n_threads > 1; - - if (!use_parallel) + // ── 行块并行路径 ── + // 分配 n_threads 组本地累加器（连续存储，cache 友好） + std::vector local_acc(n_threads * C); + for (std::size_t t = 0; t < n_threads; ++t) + for (std::size_t c = 0; c < C; ++c) + local_acc[t * C + c] = init; + + // 按行分块并行扫描 + auto& pool = global_thread_pool(); + const std::size_t base = R / n_threads; + const std::size_t rem = R % n_threads; + auto row_blocks = std::views::iota(std::size_t{0}, n_threads); + pool.parallel_for_blocks( + row_blocks.begin(), row_blocks.end(), + [self, &local_acc, &reduce_op, &transform_op, C, base, rem](std::size_t t) noexcept { - // ── 单线程行主序扫描 ── - for (std::size_t c = 0; c < C; ++c) - out[c] = static_cast(init); - for (std::size_t r = 0; r < R; ++r) + const std::size_t r0 = t * base + std::min(t, rem); + const std::size_t r_end = (t + 1) * base + std::min(t + 1, rem); + // 本线程累加器行视图（lambda 体内局部，捕获语义不受影响） + auto acc = std::span(local_acc).subspan(t * C, C); + for (std::size_t r = r0; r < r_end; ++r) { - // 行视图：std::span::subspan 零成本（ptr+len），替代裸指针行起点 const auto row = self.subspan(r * C, C); for (std::size_t c = 0; c < C; ++c) { element v = static_cast(transform_op(row[c])); - out[c] = static_cast(reduce_op(static_cast(out[c]), v)); + acc[c] = reduce_op(acc[c], v); } } - return result; - } + }); - // ── 行块并行路径 ── - // 分配 n_threads 组本地累加器（连续存储，cache 友好） - std::vector local_acc(n_threads * C); - for (std::size_t t = 0; t < n_threads; ++t) - for (std::size_t c = 0; c < C; ++c) - local_acc[t * C + c] = init; - - // 按行分块并行扫描 - auto& pool = global_thread_pool(); - const std::size_t base = R / n_threads; - const std::size_t rem = R % n_threads; - auto row_blocks = std::views::iota(std::size_t{0}, n_threads); - pool.parallel_for_blocks(row_blocks.begin(), row_blocks.end(), - [self, &local_acc, &reduce_op, &transform_op, C, base, rem](std::size_t t) noexcept { - const std::size_t r0 = t * base + std::min(t, rem); - const std::size_t r_end = (t + 1) * base + std::min(t + 1, rem); - // 本线程累加器行视图（lambda 体内局部，捕获语义不受影响） - auto acc = std::span(local_acc).subspan(t * C, C); - for (std::size_t r = r0; r < r_end; ++r) - { - const auto row = self.subspan(r * C, C); - for (std::size_t c = 0; c < C; ++c) - { - element v = static_cast(transform_op(row[c])); - acc[c] = reduce_op(acc[c], v); - } - } - }); - - // 归并阶段：串行合并 n_threads 组累加器到 out[c] - // 第 0 组直接写入，其余组归并进来（reduce_op 满足结合律，结果与单线程一致） + // 归并阶段：串行合并 n_threads 组累加器到 out[c] + // 第 0 组直接写入，其余组归并进来（reduce_op 满足结合律，结果与单线程一致） + for (std::size_t c = 0; c < C; ++c) + out[c] = static_cast(local_acc[c]); // 组 0 + for (std::size_t t = 1; t < n_threads; ++t) + { + const auto acc = std::span(local_acc).subspan(t * C, C); for (std::size_t c = 0; c < C; ++c) - out[c] = static_cast(local_acc[c]); // 组 0 - for (std::size_t t = 1; t < n_threads; ++t) - { - const auto acc = std::span(local_acc).subspan(t * C, C); - for (std::size_t c = 0; c < C; ++c) - out[c] = static_cast(reduce_op(static_cast(out[c]), acc[c])); - } - return result; + out[c] = static_cast(reduce_op(static_cast(out[c]), acc[c])); } + return result; + } +}; - }; - - // ── 便捷类型别名 ──────────────────────────────────────────────────── - using Matrix = MatrixT; - using MatrixF32 = MatrixT; - using MatrixF16 = MatrixT; +// ── 便捷类型别名 ──────────────────────────────────────────────────── +using Matrix = MatrixT; +using MatrixF32 = MatrixT; +using MatrixF16 = MatrixT; } // namespace nn diff --git a/include/neuralnet.cpp/algebra_ops.hpp b/include/neuralnet.cpp/algebra_ops.hpp index d58db100..4ce9ff90 100644 --- a/include/neuralnet.cpp/algebra_ops.hpp +++ b/include/neuralnet.cpp/algebra_ops.hpp @@ -135,13 +135,15 @@ struct Ne struct Max { - [[nodiscard]] static constexpr Scalar apply(Scalar a, Scalar b) noexcept { return std::max(a, b); } + [[nodiscard]] static constexpr Scalar apply(Scalar a, Scalar b) noexcept + { return std::max(a, b); } static constexpr ExprOp op_id() noexcept { return ExprOp::Max; } }; struct Min { - [[nodiscard]] static constexpr Scalar apply(Scalar a, Scalar b) noexcept { return std::min(a, b); } + [[nodiscard]] static constexpr Scalar apply(Scalar a, Scalar b) noexcept + { return std::min(a, b); } static constexpr ExprOp op_id() noexcept { return ExprOp::Min; } }; diff --git a/include/neuralnet.cpp/algebra_span.hpp b/include/neuralnet.cpp/algebra_span.hpp index ec04e4a6..e5d00cd6 100644 --- a/include/neuralnet.cpp/algebra_span.hpp +++ b/include/neuralnet.cpp/algebra_span.hpp @@ -21,13 +21,13 @@ namespace nn // 原语的契约，DSL 只是其唯一使用者）──────────────────────────────────── /// 表达式：可按索引求值为 Scalar 的类型（Span/ConstSpan 与所有 DSL 节点满足） template -concept Expression = requires(const T &t, std::size_t i) { +concept Expression = requires (const T& t, std::size_t i) { { t.eval(i) } -> std::convertible_to; }; /// 布尔表达式：可按索引求值为 bool 的类型（用于条件选择 select） template -concept BoolExpression = requires(const T &t, std::size_t i) { +concept BoolExpression = requires (const T& t, std::size_t i) { { t.eval(i) } -> std::convertible_to; }; @@ -38,25 +38,25 @@ concept BoolExpression = requires(const T &t, std::size_t i) { class Span { private: - Scalar *data_; + Scalar* data_; std::size_t size_; public: // ── 构造函数 ────────────────────────────────────────────────────────── constexpr Span() noexcept : data_(nullptr), size_(0) {} - constexpr Span(Scalar *data, std::size_t size) noexcept : data_(data), size_(size) {} + constexpr Span(Scalar* data, std::size_t size) noexcept : data_(data), size_(size) {} /// 从 std::span 隐式构造（兼容 Matrix::span()） constexpr Span(std::span sp) noexcept : data_(sp.data()), size_(sp.size()) {} // ── 访问器 ─────────────────────────────────────────────────────────── - [[nodiscard]] constexpr Scalar *data() noexcept { return data_; } - [[nodiscard]] constexpr const Scalar *data() const noexcept { return data_; } + [[nodiscard]] constexpr Scalar* data() noexcept { return data_; } + [[nodiscard]] constexpr const Scalar* data() const noexcept { return data_; } [[nodiscard]] constexpr std::size_t size() const noexcept { return size_; } [[nodiscard]] constexpr bool empty() const noexcept { return size_ == 0; } // ── 索引访问 ───────────────────────────────────────────────────────── - [[nodiscard]] constexpr Scalar &operator[](std::size_t i) noexcept { return data_[i]; } + [[nodiscard]] constexpr Scalar& operator[](std::size_t i) noexcept { return data_[i]; } [[nodiscard]] constexpr Scalar operator[](std::size_t i) const noexcept { return data_[i]; } // ── 表达式树叶子节点的 eval ────────────────────────────────────────── @@ -64,14 +64,10 @@ class Span // ── 子视图 ─────────────────────────────────────────────────────────── [[nodiscard]] constexpr Span subspan(std::size_t offset, std::size_t count) const noexcept - { - return Span{data_ + offset, count}; - } + { return Span{data_ + offset, count}; } [[nodiscard]] constexpr Span subspan(std::size_t offset) const noexcept - { - return Span{data_ + offset, size_ - offset}; - } + { return Span{data_ + offset, size_ - offset}; } // 注：Span 不定义成员运算符；逐元素算法一律通过表达式 DSL // （nn::dsl，见 expr_dsl.hpp）表达。 @@ -95,13 +91,13 @@ class Span class ConstSpan { private: - const Scalar *data_; + const Scalar* data_; std::size_t size_; public: // ── 构造函数 ────────────────────────────────────────────────────────── constexpr ConstSpan() noexcept : data_(nullptr), size_(0) {} - constexpr ConstSpan(const Scalar *data, std::size_t size) noexcept : data_(data), size_(size) {} + constexpr ConstSpan(const Scalar* data, std::size_t size) noexcept : data_(data), size_(size) {} /// 从 std::span 隐式构造（兼容 const Matrix::span()） constexpr ConstSpan(std::span sp) noexcept : data_(sp.data()), size_(sp.size()) {} @@ -112,10 +108,10 @@ class ConstSpan constexpr ConstSpan(std::span sp) noexcept : data_(sp.data()), size_(sp.size()) {} /// 从 Span 隐式构造（Span 可作为只读输入） - constexpr ConstSpan(const Span &s) noexcept : data_(s.data()), size_(s.size()) {} + constexpr ConstSpan(const Span& s) noexcept : data_(s.data()), size_(s.size()) {} // ── 访问器 ─────────────────────────────────────────────────────────── - [[nodiscard]] constexpr const Scalar *data() const noexcept { return data_; } + [[nodiscard]] constexpr const Scalar* data() const noexcept { return data_; } [[nodiscard]] constexpr std::size_t size() const noexcept { return size_; } [[nodiscard]] constexpr bool empty() const noexcept { return size_ == 0; } @@ -127,14 +123,10 @@ class ConstSpan // ── 子视图 ─────────────────────────────────────────────────────────── [[nodiscard]] constexpr ConstSpan subspan(std::size_t offset, std::size_t count) const noexcept - { - return ConstSpan{data_ + offset, count}; - } + { return ConstSpan{data_ + offset, count}; } [[nodiscard]] constexpr ConstSpan subspan(std::size_t offset) const noexcept - { - return ConstSpan{data_ + offset, size_ - offset}; - } + { return ConstSpan{data_ + offset, size_ - offset}; } // 注：ConstSpan 不定义成员运算符；逐元素算法一律通过表达式 DSL // （nn::dsl，见 expr_dsl.hpp）表达，DSL 运算符只接受可折叠节点 @@ -142,4 +134,3 @@ class ConstSpan }; } // namespace nn - diff --git a/include/neuralnet.cpp/backend/compute_gpu_storage.hpp b/include/neuralnet.cpp/backend/compute_gpu_storage.hpp index d79aa65d..bdcd7b8a 100644 --- a/include/neuralnet.cpp/backend/compute_gpu_storage.hpp +++ b/include/neuralnet.cpp/backend/compute_gpu_storage.hpp @@ -66,9 +66,7 @@ class GpuBuffer // 设备域标识（生命周期 v3）：backend 对象地址。buffer 持 backend keepalive // ⇒ 张量存活期内该 backend 不会被析构、其地址不会被新 backend 复用。 [[nodiscard]] DomainId backend_identity() const noexcept - { - return reinterpret_cast(backend_keepalive_.get()); - } + { return reinterpret_cast(backend_keepalive_.get()); } [[nodiscard]] std::uint64_t write_epoch() const noexcept { return write_epoch_; } void mark_content_changed() noexcept { ++write_epoch_; } @@ -78,7 +76,9 @@ class GpuBuffer GpuBuffer(VkDevice device, VkBuffer buffer, const MemoryPool::Allocation& alloc, MemoryPool& pool, std::shared_ptr backend_keepalive) : device_(device), buffer_(buffer), alloc_(alloc), pool_(pool), - backend_keepalive_(std::move(backend_keepalive)) {} + backend_keepalive_(std::move(backend_keepalive)) + { + } // 析构定义在 GpuBackend 之后：batch 录制期间已录制的 descriptor 仍引用 // 本 buffer，需延迟到 end_batch 提交完成后再销毁（见 defer_buffer_destroy）。 @@ -86,8 +86,7 @@ class GpuBuffer GpuBuffer(GpuBuffer&& o) noexcept : device_(o.device_), buffer_(o.buffer_), alloc_(o.alloc_), pool_(o.pool_), - backend_keepalive_(std::move(o.backend_keepalive_)), - write_epoch_(o.write_epoch_) + backend_keepalive_(std::move(o.backend_keepalive_)), write_epoch_(o.write_epoch_) { o.device_ = VK_NULL_HANDLE; o.buffer_ = VK_NULL_HANDLE; @@ -116,10 +115,9 @@ class GpuBuffer // 创建 Device Local 缓冲区（尺寸语义 = 字节数，§6.3） // usage 恒补 SHADER_DEVICE_ADDRESS（P3 实例表按设备地址引用张量缓冲）： // 与 MemoryPool 分配侧的 DEVICE_ADDRESS 标志成对（都随 bda_available 门控）。 - [[nodiscard]] static Result create_device_local( - VkDevice device, MemoryPool& pool, - std::size_t byte_count, VkBufferUsageFlags usage, - std::shared_ptr backend_keepalive) + [[nodiscard]] static Result + create_device_local(VkDevice device, MemoryPool& pool, std::size_t byte_count, + VkBufferUsageFlags usage, std::shared_ptr backend_keepalive) { VkBufferCreateInfo buf_info{}; buf_info.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO; @@ -131,13 +129,11 @@ class GpuBuffer VkBuffer buffer = VK_NULL_HANDLE; VkResult res = vkCreateBuffer(device, &buf_info, nullptr, &buffer); - if (res != VK_SUCCESS) - NN_FAIL("vkCreateBuffer failed: " + std::to_string(res)); + if (res != VK_SUCCESS) NN_FAIL("vkCreateBuffer failed: " + std::to_string(res)); VkMemoryRequirements mem_reqs; vkGetBufferMemoryRequirements(device, buffer, &mem_reqs); - auto alloc_r = pool.allocate( - mem_reqs, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT); + auto alloc_r = pool.allocate(mem_reqs, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT); if (!alloc_r) { vkDestroyBuffer(device, buffer, nullptr); @@ -152,20 +148,19 @@ class GpuBuffer NN_FAIL("vkBindBufferMemory failed: " + std::to_string(res)); } - return GpuBuffer(device, buffer, *alloc_r, pool, - std::move(backend_keepalive)); + return GpuBuffer(device, buffer, *alloc_r, pool, std::move(backend_keepalive)); } // 创建 Host Visible 缓冲区（尺寸语义 = 字节数，§6.3） - [[nodiscard]] static Result create_host_visible( - VkDevice device, MemoryPool& pool, - std::size_t byte_count, VkBufferUsageFlags usage, - std::shared_ptr backend_keepalive) + [[nodiscard]] static Result + create_host_visible(VkDevice device, MemoryPool& pool, std::size_t byte_count, + VkBufferUsageFlags usage, std::shared_ptr backend_keepalive) { VkBufferCreateInfo buf_info{}; buf_info.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO; buf_info.size = byte_count; - buf_info.usage = usage | VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT; + buf_info.usage = + usage | VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT; // P3：与 create_device_local 同款补 SHADER_DEVICE_ADDRESS（池分配侧 // 已配对 DEVICE_ADDRESS 标志，见 MemoryPool::create_block） if (pool.device_address_alloc()) @@ -174,8 +169,7 @@ class GpuBuffer VkBuffer buffer = VK_NULL_HANDLE; VkResult res = vkCreateBuffer(device, &buf_info, nullptr, &buffer); - if (res != VK_SUCCESS) - NN_FAIL("vkCreateBuffer failed: " + std::to_string(res)); + if (res != VK_SUCCESS) NN_FAIL("vkCreateBuffer failed: " + std::to_string(res)); VkMemoryRequirements mem_reqs; vkGetBufferMemoryRequirements(device, buffer, &mem_reqs); @@ -183,18 +177,15 @@ class GpuBuffer if (pool.is_uma()) { // 真 UMA/共享显存：slab 落在统一内存池（device-local 即 host-visible） - preferred = VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT - | VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT - | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT; - fallback = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT - | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT; + preferred = VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT | VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | + VK_MEMORY_PROPERTY_HOST_COHERENT_BIT; + fallback = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT; } else { // 独显：纯 HOST_VISIBLE（系统 RAM），避免 slab 落在 GPU 可驻留内存 - preferred = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT - | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT; - fallback = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT; + preferred = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT; + fallback = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT; } auto alloc_r = pool.allocate(mem_reqs, preferred, fallback); @@ -212,8 +203,7 @@ class GpuBuffer NN_FAIL("vkBindBufferMemory failed: " + std::to_string(res)); } - return GpuBuffer(device, buffer, *alloc_r, pool, - std::move(backend_keepalive)); + return GpuBuffer(device, buffer, *alloc_r, pool, std::move(backend_keepalive)); } [[nodiscard]] VkBuffer impl() const noexcept { return buffer_; } @@ -225,8 +215,7 @@ class GpuBuffer // buffer 字节数 = rows * cols * sizeof(elem

)（f32=4B / f16=2B） // ══════════════════════════════════════════════════════════════════════════ -template -class GpuTensorT +template class GpuTensorT { private: std::shared_ptr buffer_; @@ -237,19 +226,21 @@ class GpuTensorT GpuTensorT() = default; GpuTensorT(std::shared_ptr buffer, std::size_t rows, std::size_t cols) - : buffer_(std::move(buffer)), rows_(rows), cols_(cols) {} + : buffer_(std::move(buffer)), rows_(rows), cols_(cols) + { + } // 从 CPU Matrix 创建（上传数据；同精度跨设备 = 原始字节拷贝，§6.5） - [[nodiscard]] static Result> from_matrix( - const MatrixT

& cpu_mat, GpuBackend& backend); + [[nodiscard]] static Result> from_matrix(const MatrixT

& cpu_mat, + GpuBackend& backend); // 创建空的 GPU Tensor（用于输出） - [[nodiscard]] static Result> create_empty( - std::size_t rows, std::size_t cols, GpuBackend& backend); + [[nodiscard]] static Result> create_empty(std::size_t rows, std::size_t cols, + GpuBackend& backend); // 创建空的 Host Visible Tensor（activation offload 存储，§6.6） - [[nodiscard]] static Result> create_host_visible_empty( - std::size_t rows, std::size_t cols, GpuBackend& backend); + [[nodiscard]] static Result> + create_host_visible_empty(std::size_t rows, std::size_t cols, GpuBackend& backend); // 转换为 CPU Matrix（下载数据） [[nodiscard]] Result> to_matrix(GpuBackend& backend) const; @@ -260,14 +251,11 @@ class GpuTensorT [[nodiscard]] const GpuBuffer& buffer() const noexcept { return *buffer_; } // 底层 buffer 的共享句柄（把同一 buffer 以另一种元素类型重贴标签用， // 见 run_fused_gpu 的 out_f16：in-kernel f16 的输出是 f16 字节布局） - [[nodiscard]] std::shared_ptr shared_buffer() const noexcept - { return buffer_; } + [[nodiscard]] std::shared_ptr shared_buffer() const noexcept { return buffer_; } // 零拷贝 reshape：共享底层 buffer，只改变形状元数据 [[nodiscard]] GpuTensorT

with_shape(std::size_t new_rows, std::size_t new_cols) const - { - return GpuTensorT

(buffer_, new_rows, new_cols); - } + { return GpuTensorT

(buffer_, new_rows, new_cols); } }; } // namespace nn diff --git a/include/neuralnet.cpp/backend/compute_gpu_tensor_impl.hpp b/include/neuralnet.cpp/backend/compute_gpu_tensor_impl.hpp index caba369e..c30d13d2 100644 --- a/include/neuralnet.cpp/backend/compute_gpu_tensor_impl.hpp +++ b/include/neuralnet.cpp/backend/compute_gpu_tensor_impl.hpp @@ -14,24 +14,23 @@ namespace nn // STORAGE_BUFFER: 绑定到 shader 进行 GPU 计算 // TRANSFER_DST: 允许 vkCmdFillBuffer(zero) / vkCmdCopyBuffer(insert_rows) 写入 // TRANSFER_SRC: 允许 vkCmdCopyBuffer(clone / slice_rows) 读取 -static constexpr VkBufferUsageFlags TENSOR_BUFFER_USAGE = - VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT; +static constexpr VkBufferUsageFlags TENSOR_BUFFER_USAGE = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | + VK_BUFFER_USAGE_TRANSFER_DST_BIT | + VK_BUFFER_USAGE_TRANSFER_SRC_BIT; template -inline Result> GpuTensorT

::from_matrix(const MatrixT

& cpu_mat, GpuBackend& backend) +inline Result> GpuTensorT

::from_matrix(const MatrixT

& cpu_mat, + GpuBackend& backend) { - if (cpu_mat.empty()) - NN_FAIL("Empty matrix"); + if (cpu_mat.empty()) NN_FAIL("Empty matrix"); // buffer 字节数 = 元素数 × sizeof(elem

)（§6.3） const std::size_t byte_count = cpu_mat.size() * sizeof(elem

); - auto buf_res = GpuBuffer::create_device_local( - backend.device().device(), backend.alloc_pool(), - byte_count, - TENSOR_BUFFER_USAGE, backend.keepalive()); + auto buf_res = + GpuBuffer::create_device_local(backend.device().device(), backend.alloc_pool(), byte_count, + TENSOR_BUFFER_USAGE, backend.keepalive()); NN_TRY_CHECK(buf_res); - auto tensor = GpuTensorT

( - std::make_shared(std::move(*buf_res)), - cpu_mat.rows(), cpu_mat.cols()); + auto tensor = GpuTensorT

(std::make_shared(std::move(*buf_res)), cpu_mat.rows(), + cpu_mat.cols()); // 同精度跨设备 = 原始字节拷贝（§6.5） auto upload_res = backend.upload_blocking(tensor, cpu_mat.span()); @@ -40,41 +39,37 @@ inline Result> GpuTensorT

::from_matrix(const MatrixT

& cpu_ma } template -inline Result> GpuTensorT

::create_empty( - std::size_t rows, std::size_t cols, GpuBackend& backend) +inline Result> GpuTensorT

::create_empty(std::size_t rows, std::size_t cols, + GpuBackend& backend) { // buffer 字节数 = 元素数 × sizeof(elem

)（§6.3） const std::size_t byte_count = rows * cols * sizeof(elem

); // TRANSFER_DST_BIT: 允许 vkCmdFillBuffer (zero) 和 vkCmdCopyBuffer (insert_rows) 写入 // TRANSFER_SRC_BIT: 允许 vkCmdCopyBuffer (clone/slice_rows) 读取 - auto buf_res = GpuBuffer::create_device_local( - backend.device().device(), backend.alloc_pool(), - byte_count, - TENSOR_BUFFER_USAGE, backend.keepalive()); + auto buf_res = + GpuBuffer::create_device_local(backend.device().device(), backend.alloc_pool(), byte_count, + TENSOR_BUFFER_USAGE, backend.keepalive()); NN_TRY_CHECK(buf_res); return GpuTensorT

(std::make_shared(std::move(*buf_res)), rows, cols); } template -inline Result> GpuTensorT

::create_host_visible_empty( - std::size_t rows, std::size_t cols, GpuBackend& backend) +inline Result> +GpuTensorT

::create_host_visible_empty(std::size_t rows, std::size_t cols, GpuBackend& backend) { // Host Visible 存储（activation offload）：数据驻留主机可见内存（§6.6）。 // 仅用于 vkCmdCopyBuffer 中转（TRANSFER_DST/SRC）。 const std::size_t byte_count = rows * cols * sizeof(elem

); - auto buf_res = GpuBuffer::create_host_visible( - backend.device().device(), backend.alloc_pool(), - byte_count, - TENSOR_BUFFER_USAGE, backend.keepalive()); + auto buf_res = + GpuBuffer::create_host_visible(backend.device().device(), backend.alloc_pool(), byte_count, + TENSOR_BUFFER_USAGE, backend.keepalive()); NN_TRY_CHECK(buf_res); return GpuTensorT

(std::make_shared(std::move(*buf_res)), rows, cols); } -template -inline Result> GpuTensorT

::to_matrix(GpuBackend& backend) const +template inline Result> GpuTensorT

::to_matrix(GpuBackend& backend) const { - if (!valid()) - NN_FAIL("Invalid GpuTensor"); + if (!valid()) NN_FAIL("Invalid GpuTensor"); MatrixT

cpu_mat(rows_, cols_); auto dl_res = backend.download_blocking(*this, cpu_mat.span()); NN_TRY_CHECK(dl_res); diff --git a/include/neuralnet.cpp/backend/compute_memory_pool.hpp b/include/neuralnet.cpp/backend/compute_memory_pool.hpp index c4dd8cae..b23b40a5 100644 --- a/include/neuralnet.cpp/backend/compute_memory_pool.hpp +++ b/include/neuralnet.cpp/backend/compute_memory_pool.hpp @@ -119,9 +119,7 @@ class MemoryPool VkDeviceSize offset; bool operator==(const SuballocKey& o) const noexcept - { - return memory == o.memory && offset == o.offset; - } + { return memory == o.memory && offset == o.offset; } }; struct SuballocKeyHash @@ -176,8 +174,8 @@ class MemoryPool // 相比 first-fit，best-fit 降低把一个大小合适的洞切成两个小洞的概率， // 且跳过对齐 padding 过大的候选（避免产生不可复用的残片）。 // 返回 (offset, padding_bytes)；无合适候选返回 nullopt。 - [[nodiscard]] static std::optional> find_best( - Block& block, VkDeviceSize alignment, VkDeviceSize size) + [[nodiscard]] static std::optional> + find_best(Block& block, VkDeviceSize alignment, VkDeviceSize size) { const VkDeviceSize align_mask = alignment - 1; // 对齐 padding 预算：超过该值即视为"残片"，跳过该候选。 @@ -193,10 +191,8 @@ class MemoryPool { const VkDeviceSize aligned = (r.offset + align_mask) & ~align_mask; const VkDeviceSize padding = aligned - r.offset; - if (padding > max_padding) - continue; // 对齐残片过大，跳过 - if (r.size < padding + size) - continue; // 放不下 + if (padding > max_padding) continue; // 对齐残片过大，跳过 + if (r.size < padding + size) continue; // 放不下 // 剩余碎片 = 该区分配后剩下的字节数（越小越贴合） const VkDeviceSize waste = r.size - (padding + size); if (waste < best_waste) @@ -207,8 +203,7 @@ class MemoryPool best_region_start = r.offset; } } - if (!best_offset) - return std::nullopt; + if (!best_offset) return std::nullopt; // 定位被选中的区域并从集合中剔除，按 padding/remaining 重插 auto it = block.free_regions.lower_bound(FreeRegion{best_region_start, 1}); @@ -216,18 +211,16 @@ class MemoryPool return std::nullopt; // 防御：理论不可达 const VkDeviceSize remaining = it->size - best_padding - size; block.free_regions.erase(it); - if (best_padding > 0) - block.free_regions.insert({best_region_start, best_padding}); - if (remaining > 0) - block.free_regions.insert({*best_offset + size, remaining}); + if (best_padding > 0) block.free_regions.insert({best_region_start, best_padding}); + if (remaining > 0) block.free_regions.insert({*best_offset + size, remaining}); return std::make_pair(*best_offset, best_padding); } // 注：Block 的移动赋值被 delete（防止 vector 操作引发意外释放）， // 但移动构造可用。此处 `return block;` 经由移动构造构造 std::expected， // 不触发移动赋值，故可正常编译。 - [[nodiscard]] Result create_block( - uint32_t memory_type_index, VkMemoryPropertyFlags flags, VkDeviceSize size = 0) + [[nodiscard]] Result create_block(uint32_t memory_type_index, + VkMemoryPropertyFlags flags, VkDeviceSize size = 0) { VkDeviceSize alloc_size = (size > 0) ? size : block_size_; VkMemoryAllocateInfo alloc_info{}; @@ -240,14 +233,14 @@ class MemoryPool VkMemoryAllocateFlagsInfo addr_flags{}; addr_flags.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_FLAGS_INFO; addr_flags.flags = VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT; - if (device_address_alloc_) - alloc_info.pNext = &addr_flags; + if (device_address_alloc_) alloc_info.pNext = &addr_flags; VkDeviceMemory memory = VK_NULL_HANDLE; const auto t0 = std::chrono::steady_clock::now(); VkResult res = vkAllocateMemory(device_, &alloc_info, nullptr, &memory); - c_vkalloc_ms_ += std::chrono::duration( - std::chrono::steady_clock::now() - t0).count(); + c_vkalloc_ms_ += + std::chrono::duration(std::chrono::steady_clock::now() - t0) + .count(); if (res != VK_SUCCESS) { // issue #28：VK_ERROR_OUT_OF_DEVICE_MEMORY(-2) 的常见根因不是显存 @@ -272,20 +265,17 @@ class MemoryPool return block; } - [[nodiscard]] std::optional find_memory_type( - uint32_t type_bits, VkMemoryPropertyFlags preferred, - VkMemoryPropertyFlags fallback = 0) const + [[nodiscard]] std::optional find_memory_type(uint32_t type_bits, + VkMemoryPropertyFlags preferred, + VkMemoryPropertyFlags fallback = 0) const { std::optional best, fb; for (uint32_t i = 0; i < mem_props_.memoryTypeCount; ++i) { - if (!(type_bits & (1u << i))) - continue; + if (!(type_bits & (1u << i))) continue; auto flags = mem_props_.memoryTypes[i].propertyFlags; - if ((flags & preferred) == preferred && !best) - best = i; - if (fallback && (flags & fallback) == fallback && !fb) - fb = i; + if ((flags & preferred) == preferred && !best) best = i; + if (fallback && (flags & fallback) == fallback && !fb) fb = i; } return best ? best : fb; } @@ -296,11 +286,8 @@ class MemoryPool VkDeviceSize small_block_size = DEFAULT_SMALL_BLOCK_SIZE, VkDeviceSize ladder_max = 0) : device_(device), physical_device_(physical_device), block_size_(block_size), - small_block_size_(small_block_size), retain_free_bytes_(0), - ladder_max_(ladder_max) - { - vkGetPhysicalDeviceMemoryProperties(physical_device_, &mem_props_); - } + small_block_size_(small_block_size), retain_free_bytes_(0), ladder_max_(ladder_max) + { vkGetPhysicalDeviceMemoryProperties(physical_device_, &mem_props_); } ~MemoryPool() { blocks_.clear(); } @@ -310,24 +297,22 @@ class MemoryPool MemoryPool(MemoryPool&&) = delete; MemoryPool& operator=(MemoryPool&&) = delete; - [[nodiscard]] Result allocate( - VkMemoryRequirements requirements, - VkMemoryPropertyFlags preferred_flags, - VkMemoryPropertyFlags fallback_flags = 0) + [[nodiscard]] Result allocate(VkMemoryRequirements requirements, + VkMemoryPropertyFlags preferred_flags, + VkMemoryPropertyFlags fallback_flags = 0) { std::lock_guard lock(mutex_); c_alloc_calls_++; - auto mem_type = find_memory_type(requirements.memoryTypeBits, preferred_flags, fallback_flags); - if (!mem_type) - NN_FAIL("No suitable GPU memory type found"); + auto mem_type = + find_memory_type(requirements.memoryTypeBits, preferred_flags, fallback_flags); + if (!mem_type) NN_FAIL("No suitable GPU memory type found"); VkDeviceSize alignment = requirements.alignment > 0 ? requirements.alignment : 1; // matmul_tiled 等 vec4 SSBO 快路径要求基址 16B 对齐（std430 vec4 读的 // 对齐门槛）。buffer requirements 可能小于 16，这里统一抬到 ≥16； // 仍保持 2 的幂，find_best 的位掩码取整不受影响。代价：每次分配 // 最多 12B padding。 - if (alignment < 16u) - alignment = 16u; + if (alignment < 16u) alignment = 16u; const VkDeviceSize alloc_size = requirements.size; // 尺寸分类分池：超大分配独占整块；小分配走小块池；其余走大块池。 // 顺序：alloc > 大块阈值 → 独占超大块；alloc < 小块阈值 → 小块池。 @@ -340,25 +325,22 @@ class MemoryPool // 最小 2 的幂"，夹在 [small_block_size_, ladder_max_]；超出上限 // 则独占精确尺寸块。等价于把原来的单一中等类别拆成一串类别。 VkDeviceSize c = small_block_size_; - while (c < alloc_size && c < ladder_max_) c <<= 1; + while (c < alloc_size && c < ladder_max_) + c <<= 1; pool_size = (alloc_size > ladder_max_) ? alloc_size : c; } - else if (alloc_size > block_size_) - pool_size = alloc_size; // 超出大块底材：独占整块 + else if (alloc_size > block_size_) pool_size = alloc_size; // 超出大块底材：独占整块 else if (alloc_size < SMALL_ALLOC_THRESHOLD) - pool_size = small_block_size_; // 小分配：小块底材池（抗碎片） - else - pool_size = block_size_; + pool_size = small_block_size_; // 小分配：小块底材池（抗碎片） + else pool_size = block_size_; // 尝试在现有同类块中分配（best-fit） for (auto& block_ptr : blocks_) { auto& block = *block_ptr; c_block_scans_++; - if (block.memory_type_index != *mem_type) - continue; - if (block.size != pool_size) - continue; // 只服务同尺寸分类的块，保持大块连续性 + if (block.memory_type_index != *mem_type) continue; + if (block.size != pool_size) continue; // 只服务同尺寸分类的块，保持大块连续性 c_region_scans_ += block.free_regions.size(); auto best = find_best(block, alignment, alloc_size); if (best) @@ -374,8 +356,7 @@ class MemoryPool NN_TRY_CHECK(block_result); auto new_block = std::make_unique(std::move(*block_result)); auto best = find_best(*new_block, alignment, alloc_size); - if (!best) - NN_FAIL("Suballocation failed in new block"); + if (!best) NN_FAIL("Suballocation failed in new block"); VkDeviceMemory mem = new_block->memory; VkMemoryPropertyFlags flags = new_block->property_flags; new_block->allocation_count++; @@ -387,21 +368,18 @@ class MemoryPool void free(const Allocation& alloc) { - if (!alloc.valid()) - return; + if (!alloc.valid()) return; std::lock_guard lock(mutex_); c_free_calls_++; SuballocKey key{alloc.memory, alloc.offset}; auto active_it = active_allocs_.find(key); - if (active_it == active_allocs_.end()) - return; + if (active_it == active_allocs_.end()) return; for (auto& block_ptr : blocks_) { auto& block = *block_ptr; - if (block.memory != alloc.memory) - continue; + if (block.memory != alloc.memory) continue; FreeRegion new_reg{alloc.offset, alloc.size}; auto ins_it = block.free_regions.insert(new_reg).first; @@ -455,10 +433,10 @@ class MemoryPool { has_device_local = true; if (!(flags & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT)) - return false; // 存在纯 device-local（独显专用 VRAM） + return false; // 存在纯 device-local（独显专用 VRAM） } } - return has_device_local; // 全部 device-local 均可 host 映射 → UMA + return has_device_local; // 全部 device-local 均可 host 映射 → UMA } // ── P3：块分配的 buffer device address 标志 ───────────────────────── @@ -466,21 +444,20 @@ class MemoryPool // 之前）；读取端用于 buffer 创建侧配对 SHADER_DEVICE_ADDRESS usage 位 // （usage 与 alloc 标志必须成对，见 create_block）。 void set_device_address_alloc(bool on) noexcept { device_address_alloc_ = on; } - [[nodiscard]] bool device_address_alloc() const noexcept - { return device_address_alloc_; } + [[nodiscard]] bool device_address_alloc() const noexcept { return device_address_alloc_; } // ── 池统计（L2 仪器化，供显存采样/逐项归因） ────────────────────── struct PoolStats { - std::size_t block_count = 0; // 活跃 block 数 - std::size_t allocation_count = 0; // 活跃子分配数 - VkDeviceSize total_bytes = 0; // 所有 block 底材总大小 - VkDeviceSize free_bytes = 0; // 空闲区域总大小 - VkDeviceSize allocated_bytes = 0; // total - free - VkDeviceSize max_contiguous_free = 0; // 最大连续空闲块 - double fragmentation = 0.0; // 1 - 最大连续空闲/总空闲 - VkDeviceSize device_bytes = 0; // DEVICE_LOCAL 块（真实显存） - VkDeviceSize host_bytes = 0; // HOST_VISIBLE 块（host RAM，offload/staging 用） + std::size_t block_count = 0; // 活跃 block 数 + std::size_t allocation_count = 0; // 活跃子分配数 + VkDeviceSize total_bytes = 0; // 所有 block 底材总大小 + VkDeviceSize free_bytes = 0; // 空闲区域总大小 + VkDeviceSize allocated_bytes = 0; // total - free + VkDeviceSize max_contiguous_free = 0; // 最大连续空闲块 + double fragmentation = 0.0; // 1 - 最大连续空闲/总空闲 + VkDeviceSize device_bytes = 0; // DEVICE_LOCAL 块（真实显存） + VkDeviceSize host_bytes = 0; // HOST_VISIBLE 块（host RAM，offload/staging 用） // 池账本计数器（诊断，见 MemoryPool::c_* 注释） std::size_t c_alloc_calls = 0; @@ -505,8 +482,8 @@ class MemoryPool // 池账本计数器（诊断；把"粒度变细 → 耗时翻倍"归因到线性扫描/底材申请） s += " | calls a/f=" + std::to_string(c_alloc_calls) + "/" + std::to_string(c_free_calls); - s += " scans blk=" + std::to_string(c_block_scans) + " reg=" + - std::to_string(c_region_scans); + s += " scans blk=" + std::to_string(c_block_scans) + + " reg=" + std::to_string(c_region_scans); s += " blk_new/free=" + std::to_string(c_blocks_created) + "/" + std::to_string(c_blocks_released); s += " vkalloc=" + std::to_string(vkalloc_ms) + "ms"; @@ -531,21 +508,18 @@ class MemoryPool const auto& b = *bp; s.allocation_count += b.allocation_count; s.total_bytes += b.size; - if (b.property_flags & VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT) - s.device_bytes += b.size; - if (b.property_flags & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT) - s.host_bytes += b.size; + if (b.property_flags & VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT) s.device_bytes += b.size; + if (b.property_flags & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT) s.host_bytes += b.size; for (const auto& r : b.free_regions) { s.free_bytes += r.size; - if (r.size > s.max_contiguous_free) - s.max_contiguous_free = r.size; + if (r.size > s.max_contiguous_free) s.max_contiguous_free = r.size; } } s.allocated_bytes = s.total_bytes - s.free_bytes; if (s.free_bytes > 0) - s.fragmentation = 1.0 - static_cast(s.max_contiguous_free) - / static_cast(s.free_bytes); + s.fragmentation = 1.0 - static_cast(s.max_contiguous_free) / + static_cast(s.free_bytes); return s; } @@ -582,13 +556,11 @@ class MemoryPool // 单次释放扫描保持完全相同的安全条件，并避免额外 set 遍历。 if (retain_free_bytes_ == 0) { - for (auto it = blocks_.begin(); it != blocks_.end(); ) + for (auto it = blocks_.begin(); it != blocks_.end();) { const auto& b = **it; - if (b.allocation_count == 0 && - b.free_regions.size() == 1 && - b.free_regions.begin()->offset == 0 && - b.free_regions.begin()->size == b.size) + if (b.allocation_count == 0 && b.free_regions.size() == 1 && + b.free_regions.begin()->offset == 0 && b.free_regions.begin()->size == b.size) { it = blocks_.erase(it); c_blocks_released_++; @@ -605,14 +577,12 @@ class MemoryPool for (const auto& r : bp->free_regions) total_free += r.size; - for (auto it = blocks_.begin(); it != blocks_.end(); ) + for (auto it = blocks_.begin(); it != blocks_.end();) { const auto& b = **it; // 完全空闲：无活跃子分配，且整个 block 被单个空闲区覆盖 - if (b.allocation_count == 0 && - b.free_regions.size() == 1 && - b.free_regions.begin()->offset == 0 && - b.free_regions.begin()->size == b.size) + if (b.allocation_count == 0 && b.free_regions.size() == 1 && + b.free_regions.begin()->offset == 0 && b.free_regions.begin()->size == b.size) { // 仅当释放后总空闲仍 ≥ 保留目标时才释放（避免抖动） if (total_free - b.size >= retain_free_bytes_) diff --git a/include/neuralnet.cpp/backend/compute_staging_ring.hpp b/include/neuralnet.cpp/backend/compute_staging_ring.hpp index f91af35b..b404e612 100644 --- a/include/neuralnet.cpp/backend/compute_staging_ring.hpp +++ b/include/neuralnet.cpp/backend/compute_staging_ring.hpp @@ -38,7 +38,7 @@ class StagingRing // 默认 64MB/region × DEFAULT_NUM_REGIONS=4（实际分配大小由 initialize() // 动态计算，会根据 host-visible 显存大小自动扩容，此值仅作为下限参考） static constexpr std::size_t DEFAULT_REGION_SIZE = 64ull * 1024 * 1024; // 64MB - static constexpr std::size_t DEFAULT_NUM_REGIONS = 4; // region 轮转深度，降低 acquire 阻塞概率 + static constexpr std::size_t DEFAULT_NUM_REGIONS = 4; // region 轮转深度，降低 acquire 阻塞概率 // host-visible 显存预算占比的倒数（实际取 1/HOST_VISIBLE_FRACTION） static constexpr VkDeviceSize HOST_VISIBLE_FRACTION = 16; // 动态计算 staging 大小的下限，避免小显存机器分配过小 @@ -94,17 +94,16 @@ class StagingRing std::vector regions_; std::size_t region_size_; std::atomic current_{0}; - bool timeline_ = false; // 用时间线信号量（否则 host 等 fence 回退） + bool timeline_ = false; // 用时间线信号量（否则 host 等 fence 回退） public: StagingRing() = default; - [[nodiscard]] Result initialize( - VkDevice device, VkPhysicalDevice physical_device, - VkCommandPool cmd_pool, MemoryPool& pool, - bool use_timeline = false, - std::size_t region_size = DEFAULT_REGION_SIZE, - std::size_t num_regions = DEFAULT_NUM_REGIONS) + [[nodiscard]] Result initialize(VkDevice device, VkPhysicalDevice physical_device, + VkCommandPool cmd_pool, MemoryPool& pool, + bool use_timeline = false, + std::size_t region_size = DEFAULT_REGION_SIZE, + std::size_t num_regions = DEFAULT_NUM_REGIONS) { device_ = device; physical_device_ = physical_device; @@ -134,8 +133,7 @@ class StagingRing // 最终取 用户请求值 与 动态计算值 的较大者 // （上限已封顶，避免小显存机器被 staging 占满预算） - region_size_ = std::max( - static_cast(region_size), calculated_size); + region_size_ = std::max(static_cast(region_size), calculated_size); regions_.resize(num_regions); for (std::size_t i = 0; i < num_regions; ++i) @@ -154,33 +152,29 @@ class StagingRing buf_info.sharingMode = VK_SHARING_MODE_EXCLUSIVE; VkResult res = vkCreateBuffer(device_, &buf_info, nullptr, &r.buffer); - if (res != VK_SUCCESS) - NN_FAIL("vkCreateBuffer failed: " + std::to_string(res)); + if (res != VK_SUCCESS) NN_FAIL("vkCreateBuffer failed: " + std::to_string(res)); // 分配内存 VkMemoryRequirements mem_reqs; vkGetBufferMemoryRequirements(device_, r.buffer, &mem_reqs); - auto alloc_r = pool_->allocate( - mem_reqs, - VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT); + auto alloc_r = pool_->allocate(mem_reqs, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | + VK_MEMORY_PROPERTY_HOST_COHERENT_BIT); NN_TRY_CHECK(alloc_r); r.alloc = *alloc_r; // 绑定内存 res = vkBindBufferMemory(device_, r.buffer, r.alloc.memory, r.alloc.offset); - if (res != VK_SUCCESS) - NN_FAIL("vkBindBufferMemory failed: " + std::to_string(res)); + if (res != VK_SUCCESS) NN_FAIL("vkBindBufferMemory failed: " + std::to_string(res)); // 映射内存 - res = vkMapMemory(device_, r.alloc.memory, r.alloc.offset, r.alloc.size, 0, &r.mapped_ptr); - if (res != VK_SUCCESS) - NN_FAIL("vkMapMemory failed: " + std::to_string(res)); + res = vkMapMemory(device_, r.alloc.memory, r.alloc.offset, r.alloc.size, 0, + &r.mapped_ptr); + if (res != VK_SUCCESS) NN_FAIL("vkMapMemory failed: " + std::to_string(res)); // 创建 fence VkFenceCreateInfo fence_info{}; fence_info.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO; fence_info.flags = VK_FENCE_CREATE_SIGNALED_BIT; // 初始状态为 signaled res = vkCreateFence(device_, &fence_info, nullptr, &r.fence); - if (res != VK_SUCCESS) - NN_FAIL("vkCreateFence failed: " + std::to_string(res)); + if (res != VK_SUCCESS) NN_FAIL("vkCreateFence failed: " + std::to_string(res)); // 创建跨 submit 信号量（仅时间线模式；初始 value 0 = 未 signal， // 见 Region::semaphore 注释） if (timeline_) @@ -214,19 +208,14 @@ class StagingRing ~StagingRing() { - if (device_ == VK_NULL_HANDLE) - return; + if (device_ == VK_NULL_HANDLE) return; for (auto& r : regions_) { - if (r.fence != VK_NULL_HANDLE) - vkDestroyFence(device_, r.fence, nullptr); - if (r.semaphore != VK_NULL_HANDLE) - vkDestroySemaphore(device_, r.semaphore, nullptr); - if (r.cmd != VK_NULL_HANDLE) - vkFreeCommandBuffers(device_, cmd_pool_, 1, &r.cmd); - if (r.buffer != VK_NULL_HANDLE) - vkDestroyBuffer(device_, r.buffer, nullptr); + if (r.fence != VK_NULL_HANDLE) vkDestroyFence(device_, r.fence, nullptr); + if (r.semaphore != VK_NULL_HANDLE) vkDestroySemaphore(device_, r.semaphore, nullptr); + if (r.cmd != VK_NULL_HANDLE) vkFreeCommandBuffers(device_, cmd_pool_, 1, &r.cmd); + if (r.buffer != VK_NULL_HANDLE) vkDestroyBuffer(device_, r.buffer, nullptr); // MemoryPool 会自动释放内存 } } @@ -264,8 +253,7 @@ class StagingRing { for (auto& r : regions_) { - if (!r.in_flight) - continue; + if (!r.in_flight) continue; vkWaitForFences(device_, 1, &r.fence, VK_TRUE, UINT64_MAX); vkResetFences(device_, 1, &r.fence); r.in_flight = false; @@ -274,62 +262,50 @@ class StagingRing // 上传数据到 staging region（元素类型无关，memcpy 字节级操作，§6.3） template - [[nodiscard]] Result upload( - std::size_t region_idx, std::span data, VkDeviceSize offset = 0) + [[nodiscard]] Result upload(std::size_t region_idx, std::span data, + VkDeviceSize offset = 0) { - if (region_idx >= regions_.size()) - NN_FAIL("Invalid region index"); + if (region_idx >= regions_.size()) NN_FAIL("Invalid region index"); auto& r = regions_[region_idx]; const std::size_t byte_size = data.size() * sizeof(T); - if (offset + byte_size > region_size_) - NN_FAIL("Upload exceeds staging region size"); + if (offset + byte_size > region_size_) NN_FAIL("Upload exceeds staging region size"); std::memcpy(static_cast(r.mapped_ptr) + offset, data.data(), byte_size); return {}; } // 从 staging region 下载数据（元素类型无关，memcpy 字节级操作，§6.3） template - [[nodiscard]] Result download( - std::size_t region_idx, std::span data, VkDeviceSize offset = 0) + [[nodiscard]] Result download(std::size_t region_idx, std::span data, + VkDeviceSize offset = 0) { - if (region_idx >= regions_.size()) - NN_FAIL("Invalid region index"); + if (region_idx >= regions_.size()) NN_FAIL("Invalid region index"); auto& r = regions_[region_idx]; const std::size_t byte_size = data.size() * sizeof(T); - if (offset + byte_size > region_size_) - NN_FAIL("Download exceeds staging region size"); + if (offset + byte_size > region_size_) NN_FAIL("Download exceeds staging region size"); std::memcpy(data.data(), static_cast(r.mapped_ptr) + offset, byte_size); return {}; } // 获取缓冲区句柄 [[nodiscard]] VkBuffer buffer(std::size_t region_idx) const noexcept - { - return regions_[region_idx].buffer; - } + { return regions_[region_idx].buffer; } // 获取 fence [[nodiscard]] VkFence fence(std::size_t region_idx) const noexcept - { - return regions_[region_idx].fence; - } + { return regions_[region_idx].fence; } // 获取 region 的跨 submit 信号量（消费方 submit 的 pWaitSemaphores 用） [[nodiscard]] VkSemaphore semaphore(std::size_t region_idx) const noexcept - { - return regions_[region_idx].semaphore; - } + { return regions_[region_idx].semaphore; } // 是否使用时间线信号量（false = 走 drain_in_flight 阻塞回退） [[nodiscard]] bool timeline() const noexcept { return timeline_; } // 本 region 最近一次上传的信号 value（消费方 pWaitSemaphoreValues 用） [[nodiscard]] std::uint64_t signal_value(std::size_t region_idx) const noexcept - { - return regions_[region_idx].signal_value; - } + { return regions_[region_idx].signal_value; } // 为本次上传分配下一个信号 value（严格递增；必须在 mark_in_flight 前调用） [[nodiscard]] std::uint64_t next_signal_value(std::size_t region_idx) noexcept @@ -342,21 +318,14 @@ class StagingRing // 获取 region 的专属 command buffer（上传 copy 录制用；acquire 等完 // fence 后 vkResetCommandBuffer 复用，禁止在 pending 状态释放/重建） [[nodiscard]] VkCommandBuffer command_buffer(std::size_t region_idx) const noexcept - { - return regions_[region_idx].cmd; - } + { return regions_[region_idx].cmd; } // region 是否 in flight（已提交、fence 尚未被等待） [[nodiscard]] bool in_flight(std::size_t region_idx) const noexcept - { - return regions_[region_idx].in_flight; - } + { return regions_[region_idx].in_flight; } // 标记 region 为正在使用 - void mark_in_flight(std::size_t region_idx) noexcept - { - regions_[region_idx].in_flight = true; - } + void mark_in_flight(std::size_t region_idx) noexcept { regions_[region_idx].in_flight = true; } // 获取 region 大小 [[nodiscard]] std::size_t region_size() const noexcept { return region_size_; } diff --git a/include/neuralnet.cpp/backend/compute_vk_backend.hpp b/include/neuralnet.cpp/backend/compute_vk_backend.hpp index 842a38cc..31e3882f 100644 --- a/include/neuralnet.cpp/backend/compute_vk_backend.hpp +++ b/include/neuralnet.cpp/backend/compute_vk_backend.hpp @@ -68,20 +68,20 @@ // 所有 get__spirv() 由列表统一展开（见 GpuBackend 内）。 #define NN_SPV_FN_MISSING nn::detail::spirv_missing -#define NN_SPV_GETTER_LIST(X) \ -X(matmul) X(matmul_tiled) X(matmul_gemv) X(matmul_gemv_f16)\ -X(batched_matmul) X(batched_matmul_f16) X(matmul_tiled_f16)\ -X(matmul_tiled_fine) X(rearrange_3d) X(elementwise_v2) X(reduce)\ -X(reduce_f16) X(transpose) X(im2col) X(im2col_f16) X(col2im)\ -X(col2im_f16) X(group_reduce) X(group_reduce_f16) X(gather)\ -X(transpose_f16) X(rearrange_3d_f16) X(gather_f16) X(elementwise_v2_f16)\ -X(scatter_add) X(scatter_add_f16) X(scan_prefix_outer)\ -X(scan_prefix_outer_f16) X(scan_suffix_outer) X(scan_suffix_outer_f16)\ -X(scan_prefix_outer_d32) X(scan_prefix_outer_d32_f16)\ -X(scan_prefix_outer_gen) X(scan_prefix_outer_gen_f16)\ -X(scan_suffix_outer_gen) X(scan_suffix_outer_gen_f16) X(outer_col)\ -X(outer_col_f16) X(cast) X(trit_pack) X(matmul_q) X(matmul_q_f16)\ -X(matmul_q_gemv) X(matmul_q_gemv_f16) +#define NN_SPV_GETTER_LIST(X) \ + X(matmul) \ + X(matmul_tiled) X(matmul_gemv) X(matmul_gemv_f16) X(batched_matmul) X(batched_matmul_f16) \ + X(matmul_tiled_f16) X(matmul_tiled_fine) X(rearrange_3d) X(elementwise_v2) X(reduce) \ + X(reduce_f16) X(transpose) X(im2col) X(im2col_f16) X(col2im) X(col2im_f16) \ + X(group_reduce) X(group_reduce_f16) X(gather) X(transpose_f16) X(rearrange_3d_f16) \ + X(gather_f16) X(elementwise_v2_f16) X(scatter_add) X(scatter_add_f16) \ + X(scan_prefix_outer) X(scan_prefix_outer_f16) X(scan_suffix_outer) \ + X(scan_suffix_outer_f16) X(scan_prefix_outer_d32) \ + X(scan_prefix_outer_d32_f16) X(scan_prefix_outer_gen) \ + X(scan_prefix_outer_gen_f16) X(scan_suffix_outer_gen) \ + X(scan_suffix_outer_gen_f16) X(outer_col) X(outer_col_f16) \ + X(cast) X(trit_pack) X(matmul_q) X(matmul_q_f16) \ + X(matmul_q_gemv) X(matmul_q_gemv_f16) #if __has_include("matmul_spv.hpp") #include "matmul_spv.hpp" #define NN_SPV_FN_matmul nn_matmul_spirv_bytecode @@ -456,10 +456,10 @@ namespace detail struct FusedBatchInst { std::vector inputs; // 与 spec.views 同序 - const GpuBuffer* out = nullptr; // 主输出（槽 n_in） + const GpuBuffer* out = nullptr; // 主输出（槽 n_in） std::vector extra_outs; // extras 输出（槽 n_in+1+k） - std::uint32_t count = 0; // rows*cols（元素数） - std::uint32_t cols = 0; // 每实例列数（row/col 推导） + std::uint32_t count = 0; // rows*cols（元素数） + std::uint32_t cols = 0; // 每实例列数（row/col 推导） }; // ══════════════════════════════════════════════════════════════════════════ @@ -685,11 +685,9 @@ class GpuBackend : public std::enable_shared_from_this // ── SPIR-V 字节码 getter（由文件头的单一注册表展开；RV-08）──────────── // 每个 shader 的字节码访问函数在文件头绑定为 NN_SPV_FN_（有嵌入头 = 真实 // 函数；纯 CPU 构建 = 占位空向量）；此处由列表统一展开为 get__spirv()。 -#define NN_SPV_GETTER(NAME) \ - [[nodiscard]] static const std::vector& get_##NAME##_spirv() \ - { \ - return NN_SPV_FN_##NAME(); \ - } +#define NN_SPV_GETTER(NAME) \ + [[nodiscard]] static const std::vector& get_##NAME##_spirv() \ + { return NN_SPV_FN_##NAME(); } NN_SPV_GETTER_LIST(NN_SPV_GETTER) #undef NN_SPV_GETTER @@ -715,8 +713,7 @@ class GpuBackend : public std::enable_shared_from_this // 释放帧环 fences（环内 command buffers 随 command pool 一起释放） for (auto& f : frames_) { - if (f.fence != VK_NULL_HANDLE) - vkDestroyFence(device_.device(), f.fence, nullptr); + if (f.fence != VK_NULL_HANDLE) vkDestroyFence(device_.device(), f.fence, nullptr); } // 独立模式复用 fence（solo command buffer 随 command pool 释放） @@ -728,8 +725,7 @@ class GpuBackend : public std::enable_shared_from_this if (gpu_tensor_pool_ != VK_NULL_HANDLE) vkDestroyDescriptorPool(device_.device(), gpu_tensor_pool_, nullptr); - if (ts_pool_ != VK_NULL_HANDLE) - vkDestroyQueryPool(device_.device(), ts_pool_, nullptr); + if (ts_pool_ != VK_NULL_HANDLE) vkDestroyQueryPool(device_.device(), ts_pool_, nullptr); if (command_pool_ != VK_NULL_HANDLE) vkDestroyCommandPool(device_.device(), command_pool_, nullptr); } @@ -754,10 +750,7 @@ class GpuBackend : public std::enable_shared_from_this // 设备域标识（生命周期 v3）：backend 对象地址。GpuBuffer 持 keepalive ⇒ // 张量存活期内该地址不会被新 backend 复用（P3 改 weak 缓存后依然成立）。 - [[nodiscard]] DomainId identity() const noexcept - { - return reinterpret_cast(this); - } + [[nodiscard]] DomainId identity() const noexcept { return reinterpret_cast(this); } // 进程共享 backend（生命周期 v3 / D5）：weak 缓存 + mutex。 // · 有活使用者（GpuEngine / GpuBuffer / 调用方句柄）时共享同一实例； @@ -770,8 +763,7 @@ class GpuBackend : public std::enable_shared_from_this static std::mutex mu; static std::weak_ptr weak; std::lock_guard lock(mu); - if (auto alive = weak.lock()) - return alive; + if (auto alive = weak.lock()) return alive; std::shared_ptr backend(new GpuBackend()); weak = backend; return backend; @@ -779,14 +771,10 @@ class GpuBackend : public std::enable_shared_from_this // 独立实例（测试/多设备）：由调用方决定生命周期。 [[nodiscard]] static std::shared_ptr create() - { - return std::shared_ptr(new GpuBackend()); - } + { return std::shared_ptr(new GpuBackend()); } [[nodiscard]] std::shared_ptr keepalive() noexcept - { - return weak_from_this().lock(); - } + { return weak_from_this().lock(); } // ── 成员定义按主题拆分为 4 段（RV-08 物理拆分；纯移动，顺序即成员序）── #include "compute_vk_backend_destroy.inc" // 延迟销毁 / reap / 收帧 / 显存回收 @@ -806,8 +794,7 @@ class GpuBackend : public std::enable_shared_from_this // ══════════════════════════════════════════════════════════════════════ inline GpuBuffer::~GpuBuffer() { - if (buffer_ == VK_NULL_HANDLE) - return; + if (buffer_ == VK_NULL_HANDLE) return; // 本 keepalive 保证 backend 在析构全程存活；不再经全局 instance() 取 // 一个可能已经销毁的单例。 @@ -833,8 +820,7 @@ inline GpuBuffer::~GpuBuffer() { // 无在飞帧引用：立即销毁 + 归还内存 vkDestroyBuffer(device_, buffer_, nullptr); - if (pool_ && alloc_.valid()) - pool_->free(alloc_); + if (pool_ && alloc_.valid()) pool_->free(alloc_); } buffer_ = VK_NULL_HANDLE; diff --git a/include/neuralnet.cpp/backend/compute_vk_backend_destroy.inc b/include/neuralnet.cpp/backend/compute_vk_backend_destroy.inc index 9f0bf3f7..2d8f7092 100644 --- a/include/neuralnet.cpp/backend/compute_vk_backend_destroy.inc +++ b/include/neuralnet.cpp/backend/compute_vk_backend_destroy.inc @@ -4,631 +4,588 @@ // ⚠ 本文件不是独立头：仅由 backend/compute_vk_backend.hpp 在 // `class GpuBackend { ... }` 体内**按序** #include（顺序即成员序）。 // ═══════════════════════════════════════════════════════════════════════════ - // ── 延迟销毁决策（GpuBuffer 析构时调用）──────────────────────────── - // 返回 true = 必须延迟（调用 defer_buffer_destroy）；false = 可立即销毁。 - // - batch 录制中：延迟，打当前帧标签（buffer 仍被本帧已录制命令引用） - // - 非 batch：最近一次提交的帧若仍在飞行则延迟（buffer 可能被该帧 - // 命令引用，帧 fence 信号即可证明安全）；否则立即销毁 - // 内存归还与 buffer 销毁同批：buffer 销毁前不得把区间还给池。 - [[nodiscard]] bool needs_deferred_destroy() const noexcept +// ── 延迟销毁决策（GpuBuffer 析构时调用）──────────────────────────── +// 返回 true = 必须延迟（调用 defer_buffer_destroy）；false = 可立即销毁。 +// - batch 录制中：延迟，打当前帧标签（buffer 仍被本帧已录制命令引用） +// - 非 batch：最近一次提交的帧若仍在飞行则延迟（buffer 可能被该帧 +// 命令引用，帧 fence 信号即可证明安全）；否则立即销毁 +// 内存归还与 buffer 销毁同批：buffer 销毁前不得把区间还给池。 +[[nodiscard]] bool needs_deferred_destroy() const noexcept +{ + if (batch_mode_) return true; + return last_active_frame_ < frames_.size() && frames_[last_active_frame_].in_flight; +} + +// ── 延迟销毁入口（GpuBuffer 析构时调用）────────────────────────── +// 打上帧标签：batch 录制中 → 当前帧；非 batch → 最近活跃帧。 +void defer_buffer_destroy(VkDevice device, VkBuffer buffer, const MemoryPool::Allocation& alloc, + observer_ptr pool) +{ + const std::size_t frame = batch_mode_ ? batch_frame_ : last_active_frame_; + std::lock_guard lock(pending_mutex_); + pending_destroys_.push_back({device, buffer, alloc, pool, frame}); +} + +// ── 统一执行延迟销毁（必须在对应帧的 fence 信号之后调用）───────── +// 先 vkDestroyBuffer 再归还内存到池：保证 buffer 对象销毁后 +// 其区间才可被新分配复用，杜绝存活 buffer 间的内存重叠。 +// flush 全部：用于"所有帧已完成"的 drain 点（wait_in_flight / 析构）。 +void flush_pending_destroys() +{ + std::lock_guard lock(pending_mutex_); + for (auto& pd : pending_destroys_) { - if (batch_mode_) - return true; - return last_active_frame_ < frames_.size() && - frames_[last_active_frame_].in_flight; + if (pd.buffer != VK_NULL_HANDLE && pd.device != VK_NULL_HANDLE) + vkDestroyBuffer(pd.device, pd.buffer, nullptr); + if (pd.pool && pd.alloc.valid()) pd.pool->free(pd.alloc); } - - // ── 延迟销毁入口（GpuBuffer 析构时调用）────────────────────────── - // 打上帧标签：batch 录制中 → 当前帧；非 batch → 最近活跃帧。 - void defer_buffer_destroy(VkDevice device, VkBuffer buffer, - const MemoryPool::Allocation& alloc, - observer_ptr pool) + pending_destroys_.clear(); +} + +// ── reduce scratch 的生命周期 ──────────────────────────────────────── +// scratch 由 backend 直接持有（不是 GpuBuffer），避免 +// backend → GpuTensor → GpuBuffer(shared_ptr) 的引用环。 + +// 析构期：device 已 idle，直接销毁并归还 pool。 +void destroy_reduce_partial_now_() +{ + if (reduce_partial_.buffer == VK_NULL_HANDLE) return; + vkDestroyBuffer(device_.device(), reduce_partial_.buffer, nullptr); + if (reduce_partial_.pool && reduce_partial_.alloc.valid()) + reduce_partial_.pool->free(reduce_partial_.alloc); + reduce_partial_ = {}; +} + +// 运行期形状变化：batch 录制中 / 帧在飞时走延迟队列，其余立即释放。 +void retire_reduce_partial_() +{ + if (reduce_partial_.buffer == VK_NULL_HANDLE) return; + if (needs_deferred_destroy()) { - const std::size_t frame = - batch_mode_ ? batch_frame_ : last_active_frame_; - std::lock_guard lock(pending_mutex_); - pending_destroys_.push_back({device, buffer, alloc, pool, frame}); + defer_buffer_destroy(device_.device(), reduce_partial_.buffer, reduce_partial_.alloc, + observer_ptr(reduce_partial_.pool)); } - - // ── 统一执行延迟销毁（必须在对应帧的 fence 信号之后调用）───────── - // 先 vkDestroyBuffer 再归还内存到池：保证 buffer 对象销毁后 - // 其区间才可被新分配复用，杜绝存活 buffer 间的内存重叠。 - // flush 全部：用于"所有帧已完成"的 drain 点（wait_in_flight / 析构）。 - void flush_pending_destroys() - { - std::lock_guard lock(pending_mutex_); - for (auto& pd : pending_destroys_) - { - if (pd.buffer != VK_NULL_HANDLE && pd.device != VK_NULL_HANDLE) - vkDestroyBuffer(pd.device, pd.buffer, nullptr); - if (pd.pool && pd.alloc.valid()) - pd.pool->free(pd.alloc); - } - pending_destroys_.clear(); - } - - // ── reduce scratch 的生命周期 ──────────────────────────────────────── - // scratch 由 backend 直接持有（不是 GpuBuffer），避免 - // backend → GpuTensor → GpuBuffer(shared_ptr) 的引用环。 - - // 析构期：device 已 idle，直接销毁并归还 pool。 - void destroy_reduce_partial_now_() + else { - if (reduce_partial_.buffer == VK_NULL_HANDLE) - return; vkDestroyBuffer(device_.device(), reduce_partial_.buffer, nullptr); if (reduce_partial_.pool && reduce_partial_.alloc.valid()) reduce_partial_.pool->free(reduce_partial_.alloc); - reduce_partial_ = {}; } + reduce_partial_ = {}; +} - // 运行期形状变化：batch 录制中 / 帧在飞时走延迟队列，其余立即释放。 - void retire_reduce_partial_() +[[nodiscard]] Result ensure_reduce_partial_(std::uint32_t rows, std::uint32_t cols) +{ + if (reduce_partial_.buffer != VK_NULL_HANDLE && reduce_partial_.rows == rows && + reduce_partial_.cols == cols) + return {}; + + retire_reduce_partial_(); + + MemoryPool& pool = alloc_pool(); + VkBufferCreateInfo info{}; + info.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO; + info.size = static_cast(rows) * cols * sizeof(float); + info.usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT | + VK_BUFFER_USAGE_TRANSFER_DST_BIT; + if (pool.device_address_alloc()) info.usage |= VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT; + info.sharingMode = VK_SHARING_MODE_EXCLUSIVE; + + VkBuffer buffer = VK_NULL_HANDLE; + const VkResult cr = vkCreateBuffer(device_.device(), &info, nullptr, &buffer); + if (cr != VK_SUCCESS) + NN_FAIL("reduce scratch vkCreateBuffer failed: " + std::to_string(static_cast(cr))); + + VkMemoryRequirements req{}; + vkGetBufferMemoryRequirements(device_.device(), buffer, &req); + auto alloc_r = pool.allocate(req, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT); + if (!alloc_r) { - if (reduce_partial_.buffer == VK_NULL_HANDLE) - return; - if (needs_deferred_destroy()) - { - defer_buffer_destroy(device_.device(), reduce_partial_.buffer, - reduce_partial_.alloc, - observer_ptr(reduce_partial_.pool)); - } - else - { - vkDestroyBuffer(device_.device(), reduce_partial_.buffer, nullptr); - if (reduce_partial_.pool && reduce_partial_.alloc.valid()) - reduce_partial_.pool->free(reduce_partial_.alloc); - } - reduce_partial_ = {}; + vkDestroyBuffer(device_.device(), buffer, nullptr); + NN_TRY_CHECK(alloc_r); } - - [[nodiscard]] Result ensure_reduce_partial_(std::uint32_t rows, - std::uint32_t cols) + const VkResult br = + vkBindBufferMemory(device_.device(), buffer, alloc_r->memory, alloc_r->offset); + if (br != VK_SUCCESS) { - if (reduce_partial_.buffer != VK_NULL_HANDLE && - reduce_partial_.rows == rows && reduce_partial_.cols == cols) - return {}; - - retire_reduce_partial_(); - - MemoryPool& pool = alloc_pool(); - VkBufferCreateInfo info{}; - info.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO; - info.size = static_cast(rows) * cols * sizeof(float); - info.usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | - VK_BUFFER_USAGE_TRANSFER_SRC_BIT | - VK_BUFFER_USAGE_TRANSFER_DST_BIT; - if (pool.device_address_alloc()) - info.usage |= VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT; - info.sharingMode = VK_SHARING_MODE_EXCLUSIVE; - - VkBuffer buffer = VK_NULL_HANDLE; - const VkResult cr = - vkCreateBuffer(device_.device(), &info, nullptr, &buffer); - if (cr != VK_SUCCESS) - NN_FAIL("reduce scratch vkCreateBuffer failed: " + - std::to_string(static_cast(cr))); - - VkMemoryRequirements req{}; - vkGetBufferMemoryRequirements(device_.device(), buffer, &req); - auto alloc_r = pool.allocate(req, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT); - if (!alloc_r) - { - vkDestroyBuffer(device_.device(), buffer, nullptr); - NN_TRY_CHECK(alloc_r); - } - const VkResult br = vkBindBufferMemory( - device_.device(), buffer, alloc_r->memory, alloc_r->offset); - if (br != VK_SUCCESS) - { - pool.free(*alloc_r); - vkDestroyBuffer(device_.device(), buffer, nullptr); - NN_FAIL("reduce scratch vkBindBufferMemory failed: " + - std::to_string(static_cast(br))); - } - - reduce_partial_ = ReduceScratch{buffer, *alloc_r, &pool, rows, cols}; - return {}; + pool.free(*alloc_r); + vkDestroyBuffer(device_.device(), buffer, nullptr); + NN_FAIL("reduce scratch vkBindBufferMemory failed: " + + std::to_string(static_cast(br))); } - // ── reap 单帧（多帧流水线）────────────────────────────────── - // 等待该帧 fence（wait=true；wait=false 要求调用方已用零超时 - // vkWaitForFences 确认 VK_SUCCESS）→ 销毁该帧的延迟 buffer 并归还 - // 内存 → 释放该帧描述符集（帧提交期间被已录制命令引用，fence 信号 - // 前不得释放）→ in_flight=false - [[nodiscard]] Result reap_frame(std::size_t i, bool wait) + reduce_partial_ = ReduceScratch{buffer, *alloc_r, &pool, rows, cols}; + return {}; +} + +// ── reap 单帧（多帧流水线）────────────────────────────────── +// 等待该帧 fence（wait=true；wait=false 要求调用方已用零超时 +// vkWaitForFences 确认 VK_SUCCESS）→ 销毁该帧的延迟 buffer 并归还 +// 内存 → 释放该帧描述符集（帧提交期间被已录制命令引用，fence 信号 +// 前不得释放）→ in_flight=false +[[nodiscard]] Result reap_frame(std::size_t i, bool wait) +{ + auto& f = frames_[i]; + if (f.in_flight && wait) { - auto& f = frames_[i]; - if (f.in_flight && wait) + constexpr uint64_t kFrameTimeoutNs = 60'000'000'000ULL; // 60s + const VkResult wr = + vkWaitForFences(device_.device(), 1, &f.fence, VK_TRUE, kFrameTimeoutNs); + if (wr == VK_ERROR_DEVICE_LOST) { - constexpr uint64_t kFrameTimeoutNs = 60'000'000'000ULL; // 60s - const VkResult wr = vkWaitForFences( - device_.device(), 1, &f.fence, VK_TRUE, kFrameTimeoutNs); - if (wr == VK_ERROR_DEVICE_LOST) - { - // 设备丢失（TDR 触发）：GPU 已死亡，无法恢复 - device_lost_ = true; - NN_FAIL("GPU 设备丢失 (VK_ERROR_DEVICE_LOST): Windows TDR 已重置 GPU 驱动。" "\n模型已自动保存，请使用 --resume 重启训练。" "\n建议：减小 --batch-size 或 --seq-len，或增大 Windows TDR 超时" " (注册表 TdrDelay)"); - } - if (wr != VK_SUCCESS) - NN_FAIL(std::string("GPU 帧等待失败: Vulkan error ") + std::to_string(static_cast(wr)) + "\n建议：减小 --batch-size 或 --seq-len，或增大 Windows TDR 超时" " (注册表 TdrDelay)"); + // 设备丢失（TDR 触发）：GPU 已死亡，无法恢复 + device_lost_ = true; + NN_FAIL("GPU 设备丢失 (VK_ERROR_DEVICE_LOST): Windows TDR 已重置 GPU 驱动。" + "\n模型已自动保存，请使用 --resume 重启训练。" + "\n建议：减小 --batch-size 或 --seq-len，或增大 Windows TDR 超时" + " (注册表 TdrDelay)"); } - ts_resolve(f.ts); // 设备侧计时：fence 已等（或调用方已确认），读回本帧 kernel 时间 - // 该帧的延迟销毁：先 vkDestroyBuffer 再归还内存 - { - std::lock_guard lock(pending_mutex_); - std::vector rest; - rest.reserve(pending_destroys_.size()); - for (auto& pd : pending_destroys_) - { - if (pd.frame == i) - { - if (pd.buffer != VK_NULL_HANDLE && pd.device != VK_NULL_HANDLE) - vkDestroyBuffer(pd.device, pd.buffer, nullptr); - if (pd.pool && pd.alloc.valid()) - pd.pool->free(pd.alloc); - } - else - { - rest.push_back(std::move(pd)); - } - } - pending_destroys_ = std::move(rest); - } - if (!f.desc_sets.empty()) - vkFreeDescriptorSets(device_.device(), gpu_tensor_pool_, - static_cast(f.desc_sets.size()), f.desc_sets.data()); - f.desc_sets.clear(); - f.in_flight = false; - return {}; + if (wr != VK_SUCCESS) + NN_FAIL(std::string("GPU 帧等待失败: Vulkan error ") + + std::to_string(static_cast(wr)) + + "\n建议：减小 --batch-size 或 --seq-len，或增大 Windows TDR 超时" + " (注册表 TdrDelay)"); } - - // 初始化 - // device_selector：手动指定计算设备（索引 "2" 或名称子串 "40HX"/"NVIDIA"）。 - // 空 = 自动选择（设备类型 + apiVersion 打分，见 VulkanDevice::initialize）。 - [[nodiscard]] Result initialize(std::string device_selector = {}) + ts_resolve(f.ts); // 设备侧计时：fence 已等（或调用方已确认），读回本帧 kernel 时间 + // 该帧的延迟销毁：先 vkDestroyBuffer 再归还内存 { - std::lock_guard lock(init_mutex_); - if (initialized_) - return {}; - - // 1. 检查 SPIR-V 是否可用 - const auto& spirv = get_matmul_spirv(); - if (spirv.empty()) - NN_FAIL("matmul SPIR-V bytecode not embedded"); - // 2. 初始化 Vulkan 设备（先登记设备选择器） - if (!device_selector.empty()) - device_.set_device_selector(std::move(device_selector)); - auto dev_r = device_.initialize(); - if (!dev_r) - return dev_r; - - // 显存峰值优化开关：默认在帧提交后非阻塞收割已完成帧，让步内已 - // 析构张量的内存尽早还池；NN_NO_EARLY_REAP=1 关闭（A/B 对照用）。 - early_reap_enabled_ = VulkanDevice::get_env("NN_NO_EARLY_REAP").empty(); - - // 3. 按 pipeline_specs() 表创建全部原语 pipeline。 - // 缺失（SPIR-V 未嵌入 / 设备能力不足 / 创建失败）= 空句柄， - // 由上层按 has_*() 回退；唯一硬要求是 matmul 本身。 - pipelines_.create_all(device_.device(), device_, pipeline_specs()); - if (!pipelines_.has(PipelineId::matmul)) - NN_FAIL("matmul pipeline 创建失败"); - - // 4. 创建 memory pool（持久 + 瞬态） - // 底材粒度可用 NN_POOL_BLOCK_MB 覆盖（显存峰值调参/探针用）： - // 128MB 大底材在混合尺寸+混合生命周期下会产生不可归还的内部碎片。 - VkDeviceSize block_bytes = MemoryPool::DEFAULT_BLOCK_SIZE; - { - const std::string v = VulkanDevice::get_env("NN_POOL_BLOCK_MB"); - if (!v.empty()) - { - const unsigned long long mb = std::strtoull(v.c_str(), nullptr, 10); - if (mb > 0) block_bytes = static_cast(mb) * 1024ull * 1024ull; - } - } - // 自适应阶梯上限（NN_POOL_LADDER_MAX_MB>0 启用）：中等分配按最小 - // 2 的幂类别分池，兼顾小负载抗碎片与大负载少分配次数。 - VkDeviceSize ladder_max = 0; + std::lock_guard lock(pending_mutex_); + std::vector rest; + rest.reserve(pending_destroys_.size()); + for (auto& pd : pending_destroys_) { - const std::string v = VulkanDevice::get_env("NN_POOL_LADDER_MAX_MB"); - if (!v.empty()) + if (pd.frame == i) { - const unsigned long long mb = std::strtoull(v.c_str(), nullptr, 10); - if (mb > 0) ladder_max = static_cast(mb) * 1024ull * 1024ull; + if (pd.buffer != VK_NULL_HANDLE && pd.device != VK_NULL_HANDLE) + vkDestroyBuffer(pd.device, pd.buffer, nullptr); + if (pd.pool && pd.alloc.valid()) pd.pool->free(pd.alloc); } - } - memory_pool_ = std::make_unique( - device_.device(), device_.physical_device(), block_bytes, - MemoryPool::DEFAULT_SMALL_BLOCK_SIZE, ladder_max); - transient_pool_ = std::make_unique( - device_.device(), device_.physical_device(), block_bytes, - MemoryPool::DEFAULT_SMALL_BLOCK_SIZE, ladder_max); - // P3（buffer device address）：设备支持即让**全部**池分配可取设备地址 - // （实例表引用张量缓冲；块内子分配同样覆盖），须在首次 allocate 前设置 - memory_pool_->set_device_address_alloc(device_.bda_available()); - transient_pool_->set_device_address_alloc(device_.bda_available()); - - // 5. 创建 command pool - VkCommandPoolCreateInfo pool_info{}; - pool_info.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO; - pool_info.queueFamilyIndex = device_.queue_family_index(); - pool_info.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT; - - auto r = detail::vk_check( - vkCreateCommandPool(device_.device(), &pool_info, nullptr, &command_pool_), - __FILE__, __LINE__); - if (!r) - return r; - - // 6. 创建 staging ring - // use_timeline = 设备是否支持时间线信号量（跨 submit 依赖的正确 - // 原语；不支持则 staging 环走 host 等 fence 回退，见 StagingRing） - staging_ring_ = std::make_unique(); - auto st_r = staging_ring_->initialize( - device_.device(), device_.physical_device(), command_pool_, *memory_pool_, - device_.has_timeline_semaphores()); - if (!st_r) - return st_r; - - // 6.5 异步标量回读槽：非阻塞取 loss 标量，避免每 step drain - auto rb_r = init_scalar_readback_slots(); - if (!rb_r) - return rb_r; - - // 7. 创建 descriptor pool for GPU-resident path - // elementwise_v2 需要 4 个描述符/次，按最大值计算 - // batch 模式下描述符集延迟到 end_batch 释放，需足够大以容纳整个 batch - // ViT forward+backward ≈ 230 ops/样本，batch_size=512 时累积 ~118K sets - // 增大到 262144 以留出余量，对应 desc 总量 1M - constexpr std::size_t TENSOR_POOL_SETS = 262144; - constexpr std::size_t TENSOR_POOL_DESCS = TENSOR_POOL_SETS * 4; // elementwise_v2 需要 4 - - VkDescriptorPoolSize tensor_pool_size{}; - tensor_pool_size.type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER; - tensor_pool_size.descriptorCount = static_cast(TENSOR_POOL_DESCS); - - VkDescriptorPoolCreateInfo tensor_pool_info{}; - tensor_pool_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO; - tensor_pool_info.maxSets = static_cast(TENSOR_POOL_SETS); - tensor_pool_info.poolSizeCount = 1; - tensor_pool_info.pPoolSizes = &tensor_pool_size; - tensor_pool_info.flags = VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT; - - r = detail::vk_check( - vkCreateDescriptorPool(device_.device(), &tensor_pool_info, nullptr, &gpu_tensor_pool_), - __FILE__, __LINE__); - if (!r) - return r; - - // 7b. 预分配流水线帧环：N 个 command buffer + N 个 fence， - // 环内轮转复用、全程零分配——录制/提交路径不触碰 - // vkAllocateCommandBuffers / vkCreateFence。 - frames_.resize(PIPELINE_FRAMES); - { - VkCommandBufferAllocateInfo cmd_alloc{}; - cmd_alloc.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO; - cmd_alloc.commandPool = command_pool_; - cmd_alloc.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; - cmd_alloc.commandBufferCount = static_cast(PIPELINE_FRAMES); - std::vector cmds(PIPELINE_FRAMES, VK_NULL_HANDLE); - r = detail::vk_check( - vkAllocateCommandBuffers(device_.device(), &cmd_alloc, cmds.data()), - __FILE__, __LINE__); - if (!r) - return r; - for (std::size_t i = 0; i < PIPELINE_FRAMES; ++i) - frames_[i].cmd = cmds[i]; - - VkFenceCreateInfo fence_info{}; - fence_info.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO; - // 初始 unsignaled：vkQueueSubmit 要求 fence 未 signal - for (std::size_t i = 0; i < PIPELINE_FRAMES; ++i) + else { - r = detail::vk_check( - vkCreateFence(device_.device(), &fence_info, nullptr, &frames_[i].fence), - __FILE__, __LINE__); - if (!r) - return r; + rest.push_back(std::move(pd)); } } - frame_next_ = 0; - last_active_frame_ = 0; - - // 7c. 独立模式复用 fence + command buffer（见成员注释；省去 - // 每算子 create/destroy + alloc/free 的 host/driver 开销） + pending_destroys_ = std::move(rest); + } + if (!f.desc_sets.empty()) + vkFreeDescriptorSets(device_.device(), gpu_tensor_pool_, + static_cast(f.desc_sets.size()), f.desc_sets.data()); + f.desc_sets.clear(); + f.in_flight = false; + return {}; +} + +// 初始化 +// device_selector：手动指定计算设备（索引 "2" 或名称子串 "40HX"/"NVIDIA"）。 +// 空 = 自动选择（设备类型 + apiVersion 打分，见 VulkanDevice::initialize）。 +[[nodiscard]] Result initialize(std::string device_selector = {}) +{ + std::lock_guard lock(init_mutex_); + if (initialized_) return {}; + + // 1. 检查 SPIR-V 是否可用 + const auto& spirv = get_matmul_spirv(); + if (spirv.empty()) NN_FAIL("matmul SPIR-V bytecode not embedded"); + // 2. 初始化 Vulkan 设备（先登记设备选择器） + if (!device_selector.empty()) device_.set_device_selector(std::move(device_selector)); + auto dev_r = device_.initialize(); + if (!dev_r) return dev_r; + + // 显存峰值优化开关：默认在帧提交后非阻塞收割已完成帧，让步内已 + // 析构张量的内存尽早还池；NN_NO_EARLY_REAP=1 关闭（A/B 对照用）。 + early_reap_enabled_ = VulkanDevice::get_env("NN_NO_EARLY_REAP").empty(); + + // 3. 按 pipeline_specs() 表创建全部原语 pipeline。 + // 缺失（SPIR-V 未嵌入 / 设备能力不足 / 创建失败）= 空句柄， + // 由上层按 has_*() 回退；唯一硬要求是 matmul 本身。 + pipelines_.create_all(device_.device(), device_, pipeline_specs()); + if (!pipelines_.has(PipelineId::matmul)) NN_FAIL("matmul pipeline 创建失败"); + + // 4. 创建 memory pool（持久 + 瞬态） + // 底材粒度可用 NN_POOL_BLOCK_MB 覆盖（显存峰值调参/探针用）： + // 128MB 大底材在混合尺寸+混合生命周期下会产生不可归还的内部碎片。 + VkDeviceSize block_bytes = MemoryPool::DEFAULT_BLOCK_SIZE; + { + const std::string v = VulkanDevice::get_env("NN_POOL_BLOCK_MB"); + if (!v.empty()) { - VkCommandBufferAllocateInfo cmd_alloc{}; - cmd_alloc.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO; - cmd_alloc.commandPool = command_pool_; - cmd_alloc.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; - cmd_alloc.commandBufferCount = 1; - r = detail::vk_check( - vkAllocateCommandBuffers(device_.device(), &cmd_alloc, &solo_cmd_), - __FILE__, __LINE__); - if (!r) - return r; - - VkFenceCreateInfo fence_info{}; - fence_info.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO; - // 初始 unsignaled：vkQueueSubmit 要求 fence 未 signal - r = detail::vk_check( - vkCreateFence(device_.device(), &fence_info, nullptr, &solo_fence_), - __FILE__, __LINE__); - if (!r) - return r; + const unsigned long long mb = std::strtoull(v.c_str(), nullptr, 10); + if (mb > 0) block_bytes = static_cast(mb) * 1024ull * 1024ull; } - - // 7d. timestamp query pool（NN_PROFILE 设备侧计时；按帧分块 + solo 块） - // 所选队列族 timestampValidBits=0（部分虚拟/软件设备）→ 整体 - // 降级 host 计时，ts_enabled_ 保持 false，零额外开销。 - ts_enabled_ = nn::prof::enabled() && - device_.queue_timestamp_valid_bits() > 0; - if (ts_enabled_) + } + // 自适应阶梯上限（NN_POOL_LADDER_MAX_MB>0 启用）：中等分配按最小 + // 2 的幂类别分池，兼顾小负载抗碎片与大负载少分配次数。 + VkDeviceSize ladder_max = 0; + { + const std::string v = VulkanDevice::get_env("NN_POOL_LADDER_MAX_MB"); + if (!v.empty()) { - ts_period_ns_ = device_.timestamp_period_ns(); - VkQueryPoolCreateInfo qp{}; - qp.sType = VK_STRUCTURE_TYPE_QUERY_POOL_CREATE_INFO; - qp.queryType = VK_QUERY_TYPE_TIMESTAMP; - qp.queryCount = (PIPELINE_FRAMES + 1) * kTsPairsPerBlock * 2; - const VkResult qres = - vkCreateQueryPool(device_.device(), &qp, nullptr, &ts_pool_); - if (qres != VK_SUCCESS) - { - ts_enabled_ = false; - ts_pool_ = VK_NULL_HANDLE; - std::fprintf(stderr, - "[profile] timestamp query pool 创建失败(%d)，设备侧计时降级\n", - static_cast(qres)); - } + const unsigned long long mb = std::strtoull(v.c_str(), nullptr, 10); + if (mb > 0) ladder_max = static_cast(mb) * 1024ull * 1024ull; } + } + memory_pool_ = + std::make_unique(device_.device(), device_.physical_device(), block_bytes, + MemoryPool::DEFAULT_SMALL_BLOCK_SIZE, ladder_max); + transient_pool_ = + std::make_unique(device_.device(), device_.physical_device(), block_bytes, + MemoryPool::DEFAULT_SMALL_BLOCK_SIZE, ladder_max); + // P3（buffer device address）：设备支持即让**全部**池分配可取设备地址 + // （实例表引用张量缓冲；块内子分配同样覆盖），须在首次 allocate 前设置 + memory_pool_->set_device_address_alloc(device_.bda_available()); + transient_pool_->set_device_address_alloc(device_.bda_available()); + + // 5. 创建 command pool + VkCommandPoolCreateInfo pool_info{}; + pool_info.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO; + pool_info.queueFamilyIndex = device_.queue_family_index(); + pool_info.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT; + + auto r = + detail::vk_check(vkCreateCommandPool(device_.device(), &pool_info, nullptr, &command_pool_), + __FILE__, __LINE__); + if (!r) return r; + + // 6. 创建 staging ring + // use_timeline = 设备是否支持时间线信号量（跨 submit 依赖的正确 + // 原语；不支持则 staging 环走 host 等 fence 回退，见 StagingRing） + staging_ring_ = std::make_unique(); + auto st_r = + staging_ring_->initialize(device_.device(), device_.physical_device(), command_pool_, + *memory_pool_, device_.has_timeline_semaphores()); + if (!st_r) return st_r; + + // 6.5 异步标量回读槽：非阻塞取 loss 标量，避免每 step drain + auto rb_r = init_scalar_readback_slots(); + if (!rb_r) return rb_r; + + // 7. 创建 descriptor pool for GPU-resident path + // elementwise_v2 需要 4 个描述符/次，按最大值计算 + // batch 模式下描述符集延迟到 end_batch 释放，需足够大以容纳整个 batch + // ViT forward+backward ≈ 230 ops/样本，batch_size=512 时累积 ~118K sets + // 增大到 262144 以留出余量，对应 desc 总量 1M + constexpr std::size_t TENSOR_POOL_SETS = 262144; + constexpr std::size_t TENSOR_POOL_DESCS = TENSOR_POOL_SETS * 4; // elementwise_v2 需要 4 + + VkDescriptorPoolSize tensor_pool_size{}; + tensor_pool_size.type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER; + tensor_pool_size.descriptorCount = static_cast(TENSOR_POOL_DESCS); + + VkDescriptorPoolCreateInfo tensor_pool_info{}; + tensor_pool_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO; + tensor_pool_info.maxSets = static_cast(TENSOR_POOL_SETS); + tensor_pool_info.poolSizeCount = 1; + tensor_pool_info.pPoolSizes = &tensor_pool_size; + tensor_pool_info.flags = VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT; + + r = detail::vk_check( + vkCreateDescriptorPool(device_.device(), &tensor_pool_info, nullptr, &gpu_tensor_pool_), + __FILE__, __LINE__); + if (!r) return r; + + // 7b. 预分配流水线帧环：N 个 command buffer + N 个 fence， + // 环内轮转复用、全程零分配——录制/提交路径不触碰 + // vkAllocateCommandBuffers / vkCreateFence。 + frames_.resize(PIPELINE_FRAMES); + { + VkCommandBufferAllocateInfo cmd_alloc{}; + cmd_alloc.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO; + cmd_alloc.commandPool = command_pool_; + cmd_alloc.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; + cmd_alloc.commandBufferCount = static_cast(PIPELINE_FRAMES); + std::vector cmds(PIPELINE_FRAMES, VK_NULL_HANDLE); + r = detail::vk_check(vkAllocateCommandBuffers(device_.device(), &cmd_alloc, cmds.data()), + __FILE__, __LINE__); + if (!r) return r; for (std::size_t i = 0; i < PIPELINE_FRAMES; ++i) - frames_[i].ts.base = - static_cast(i * kTsPairsPerBlock * 2); - ts_solo_.base = static_cast( - PIPELINE_FRAMES * kTsPairsPerBlock * 2); + frames_[i].cmd = cmds[i]; -#ifdef NN_FUSED_REGISTRY_EMBEDDED - // 16. 注册 AOT 融合 shader pipelines（构建期 scan_exprs 收集 + - // fused_generate.hpp 合成）。条目形态： - // · V0 / V2 / V1s / V1t（键 `key` / `key#a` / `key#s` / - // `key#xxxx` 特化签名）：N 输入 + 1 输出 binding； - // · V1 运行期精度分派（键 `key#x`）：2N + 2 binding（每个输入/输出 - // 各绑 f32 与 f16 两个视图）；PC 固定头末尾多一个 `prec` 槽。 - // push constants = 固定头（count/cols/…）+ vp 槽 + 常量池 + rp 槽。 - for (const auto& fs : nn::fused::kFusedShaders) + VkFenceCreateInfo fence_info{}; + fence_info.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO; + // 初始 unsignaled：vkQueueSubmit 要求 fence 未 signal + for (std::size_t i = 0; i < PIPELINE_FRAMES; ++i) { - if (!fs.spirv || fs.spirv_words == 0) - continue; // 空 SPIR-V：跳过（构建配置缺失） - // 运行期精度分派变体（键 "#x"）：输入/输出各双视图 + PC 多一个 prec 槽 - const bool is_disp = - std::string_view(fs.key).ends_with(nn::EXPR_PREC_DISPATCH_SUFFIX); - // P3 批量变体（键 "#b"）：单 binding（实例表），PC 与 V0 同布局 - const bool is_batch = - std::string_view(fs.key).ends_with(nn::EXPR_BATCH_SUFFIX); - // P0 细 tile 变体（键 "#f"）：V0 同布局（同绑定数、同 PC）， - // 只是另一套网格粒度（输出块 16×16）的 f32 shader。 - const bool is_fine = - std::string_view(fs.key).ends_with(nn::EXPR_MM_FINE_SUFFIX); - // K 切分变体（键 "#ks"，每 WG 一段 K → 部分和）与其伴生归约 - // （键 "#kr"）：f32 变体，PC/绑定见 EXPR_MM_KSPLIT_SUFFIX 注释。 - const bool is_ks = - std::string_view(fs.key).ends_with(nn::EXPR_MM_KSPLIT_SUFFIX); - const bool is_kr = - std::string_view(fs.key).ends_with(nn::EXPR_MM_KSPLIT_REDUCE_SUFFIX); - // 带 '#' 的 f16 变体（#x / #s / #xxxx 特化 / #a）需要 SSBO 16 位 - // 存储；设备未启用 → 跳过注册（运行时查不到 → 基类 NVI 入口回退 - // 边界 cast）。`#b` 与**纯 V0 的** `#f`/`#ks`/`#kr` 全 f32 无此 - // 要求（不得被这条误伤）。 - // ⚠ 后缀**可叠加**：`key#10003#ks` = f16 特化的 K 切分，其 shader - // 里有 float16_t 缓冲 → 仍需要 16 位存储。判据 = 去掉网格粒度 - // 尾后缀（#kr/#ks/#f）之后是否仍含 '#'。 - const std::string_view key_sv(fs.key); - std::string_view grid_stripped = key_sv; - for (std::string_view suf : - {std::string_view(nn::EXPR_BATCH_SUFFIX), - std::string_view(nn::EXPR_MM_KSPLIT_REDUCE_SUFFIX), - std::string_view(nn::EXPR_MM_KSPLIT_SUFFIX), - std::string_view(nn::EXPR_MM_FINE_SUFFIX)}) - { - if (grid_stripped.size() > suf.size() && - grid_stripped.ends_with(suf)) - { - grid_stripped.remove_suffix(suf.size()); - break; - } - } - const bool needs_16bit_storage = - key_sv.find('#') != std::string_view::npos && - !((is_batch || is_fine || is_ks || is_kr) && - grid_stripped.find('#') == std::string_view::npos); - if (needs_16bit_storage && !device_.has_16bit_storage()) - continue; - // `#b` 批量变体需要 buffer device address（实例表按 64 位地址解引用 - // 各实例缓冲）；未启用 → 跳过注册 → 运行期逐 kernel 降级（可观测）。 - if (is_batch && !device_.bda_available()) - continue; - // ALU 变体（key 以 "#a" 结尾，native16 原生 f16 算术）额外需要 - // shaderFloat16 特性；未启用 → 跳过 → 运行时回退 `#x`（其 pipeline - // 独立存在）或边界 cast——正确性不变。 - if (std::string_view(fs.key).ends_with(nn::EXPR_PREC_ALU_SUFFIX) && - !device_.has_shader_float16()) - continue; - const std::uint32_t n_in = - static_cast(fs.spec.views.size()); - // 分派变体：输入 i 的 f32 视图 binding=i、f16 视图 binding=n_in+i； - // 输出 f32/f16 视图 binding=2*n_in / 2*n_in+1。 - // P2 多输出（extras）：V0 追加 k 个输出 binding（bex，主输出之后）； - // `#x`/`#a` 变体不生成 extras（生成器跳过）⇒ 分派分支恒 2N+2。 - // `#b` 批量变体：只绑 binding=0 一个 storage buffer（实例表）。 - const std::uint32_t num_bindings = - is_batch ? 1u - : is_disp ? (2u * n_in + 2u) - : (n_in + 1u + static_cast(fs.spec.extras.size())); - // push constant **固定头**槽数（uint）：逐形态与生成器的 PC 声明 - // 逐槽一致—— - // 逐元素 2（count, cols）/ +matmul 6（…, rows, mm_k, mm_batch, - // mm_trans）/ 归约 4（…, rows, vector_out）/ 归约+matmul 7 / - // fold 5（…, vector_out, fold_k）/ fold+mm 7（…, fold_k, mm_k, - // mm_batch）+ 分派变体 1 个 `prec`。 - // ⚠ 该值**只在此处计算一次**并存进 FusedKernel::pc_base_fixed，写入侧 - // (run_fused_gpu) 只读：此前两侧各推导一遍，漂移会让 - // vkCmdPushConstants 超 range 而被驱动丢弃 → shader 读到未定义 - // 残留（随前序 op 漂移，表现为时对时错的假 PASS，"全绿"验不出）。 - const std::uint32_t pc_base_fixed = - is_ks ? 7u // K 切分：matmul 6 槽 + ks_span - : (fs.spec.fold && fs.spec.fold->matmul) ? 7u // fold+mm 7 槽 - : (fs.spec.fold) ? 5u - : (fs.reduce_axis >= 0 && fs.has_matmul) ? 7u - : (fs.reduce_axis >= 0) ? 4u - : (fs.has_matmul) ? 6u - : 2u; - // 分派变体在固定头末尾多一个 uint prec（生成器同位置声明） - const std::uint32_t pc_base = pc_base_fixed + (is_disp ? 1u : 0u); - const std::uint32_t pc_uints = pc_base + fs.view_param_count; - const std::uint32_t pc_size = - static_cast(pc_uints * sizeof(std::uint32_t) + - sizeof(Scalar) * fs.spec.consts.size() + - sizeof(Scalar) * fs.spec.rparams.size()); - auto fp_r = VulkanPipeline::create_generic( - device_.device(), std::span(fs.spirv, fs.spirv_words), - num_bindings, pc_size); - if (fp_r) - { - FusedKernel kernel; - kernel.pipeline = std::move(*fp_r); - kernel.reduce_axis = fs.reduce_axis; - kernel.has_matmul = fs.has_matmul != 0; - kernel.fold = FoldMeta{ - fs.spec.fold - ? (fs.spec.fold->vec_state_len > 0 ? fs.spec.fold->vec_state_len : 1u) - : 0u, - // per_thread_row 判据与 generate_glsl 分派同源（v2 = 有 - // vec 态或 mm 段；v1 = 纯标量） - static_cast(fs.spec.fold && - !(fs.spec.fold->vec_state_len > 0 || fs.spec.fold->matmul))}; - kernel.view_param_count = fs.view_param_count; - kernel.rparam_count = fs.rparam_count; - kernel.vec_width = fs.vec_width; - kernel.pc_base_fixed = pc_base_fixed; - // P3 批量变体：实例表 stride（uvec2 槽数），run_fused_gpu_batch - // 据此校验实例布局（生成器 S = 输入+主输出+extras+尾槽，同源） - if (is_batch) - kernel.batch_stride = static_cast( - n_in + 1u + fs.spec.extras.size() + 1u); - fused_kernels_.emplace(fs.key, std::move(kernel)); - } + r = detail::vk_check( + vkCreateFence(device_.device(), &fence_info, nullptr, &frames_[i].fence), __FILE__, + __LINE__); + if (!r) return r; } -#endif - - // ── 查询 GPU f16 硬件能力（§7.1）───────────────────────────── - // shaderFloat16：Vulkan 1.2 核心特性 / VK_KHR_shader_float16_int8 扩展。 - // 以**设备侧已启用**状态为准（VulkanDevice::initialize 查询 + 启用 + - // NN_VULKAN_NO_16BIT_ALU 逃生阀）——只查物理特性不启用 = pipeline 创建 - // 时校验失败，故此处不独立查询。 - has_shader_float16_ = device_.has_shader_float16(); - - initialized_ = true; - return {}; } + frame_next_ = 0; + last_active_frame_ = 0; - [[nodiscard]] bool is_initialized() const noexcept { return initialized_; } - [[nodiscard]] bool gpu_available() const noexcept { return initialized_; } - // ── Pipeline 变体可用性 ───────────────────────────────────────── - // 唯一查询入口：id 见 compute_vk_pipeline_catalog.hpp；缺失 = 该变体 - // 未创建（SPIR-V 未嵌入 / 设备能力不足 / 创建失败）。 - [[nodiscard]] bool has(PipelineId id) const noexcept { return pipelines_.has(id); } - - // 扫描族 f16 变体的整体可用性：四条路径齐备才算接通（缺任一 → 调用方 - // 回退边界 cast，避免"部分直读、部分 cast"的混合语义）。 - [[nodiscard]] bool has_scan_f16_pipelines() const noexcept + // 7c. 独立模式复用 fence + command buffer（见成员注释；省去 + // 每算子 create/destroy + alloc/free 的 host/driver 开销） { - return has(PipelineId::scan_prefix_outer_f16) && - has(PipelineId::scan_suffix_outer_f16) && - has(PipelineId::scan_prefix_outer_gen_f16) && - has(PipelineId::scan_suffix_outer_gen_f16); + VkCommandBufferAllocateInfo cmd_alloc{}; + cmd_alloc.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO; + cmd_alloc.commandPool = command_pool_; + cmd_alloc.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; + cmd_alloc.commandBufferCount = 1; + r = detail::vk_check(vkAllocateCommandBuffers(device_.device(), &cmd_alloc, &solo_cmd_), + __FILE__, __LINE__); + if (!r) return r; + + VkFenceCreateInfo fence_info{}; + fence_info.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO; + // 初始 unsignaled：vkQueueSubmit 要求 fence 未 signal + r = detail::vk_check(vkCreateFence(device_.device(), &fence_info, nullptr, &solo_fence_), + __FILE__, __LINE__); + if (!r) return r; } - // ── f16 硬件能力查询（§7.1）──────────────────────────────────── - [[nodiscard]] bool has_shader_float16() const noexcept { return has_shader_float16_; } - - [[nodiscard]] VulkanDevice& device() noexcept { return device_; } - [[nodiscard]] MemoryPool& memory_pool() noexcept { return *memory_pool_; } - [[nodiscard]] MemoryPool& transient_pool() noexcept { return *transient_pool_; } - // 延迟销毁队列总字节（探针归因用）：已析构但因帧未完成而尚未归还池的 - // 内存。解释"Tensor 已置空但 live 不降"的账目缺口。 - [[nodiscard]] VkDeviceSize pending_destroy_bytes() + // 7d. timestamp query pool（NN_PROFILE 设备侧计时；按帧分块 + solo 块） + // 所选队列族 timestampValidBits=0（部分虚拟/软件设备）→ 整体 + // 降级 host 计时，ts_enabled_ 保持 false，零额外开销。 + ts_enabled_ = nn::prof::enabled() && device_.queue_timestamp_valid_bits() > 0; + if (ts_enabled_) { - std::lock_guard lock(pending_mutex_); - VkDeviceSize total = 0; - for (const auto& pd : pending_destroys_) total += pd.alloc.size; - return total; - } - // 瞬态/持久分池选择器：batch 录制期→瞬态池，否则→持久池（纯组织性）。 - // 供 GpuTensorT 分配（create_empty / from_matrix / create_host_visible_empty） - // 使用，使批量内的临时/激活与构建期的参数/权重分池驻留。 - [[nodiscard]] MemoryPool& alloc_pool() noexcept - { - return batch_mode_ ? *transient_pool_ : *memory_pool_; - } - // 异步标量回读槽位数：供 GpuEngine 暴露给调用方做环形复用。 - [[nodiscard]] static constexpr std::size_t scalar_readback_slot_count() noexcept - { - return SCALAR_READBACK_SLOTS; + ts_period_ns_ = device_.timestamp_period_ns(); + VkQueryPoolCreateInfo qp{}; + qp.sType = VK_STRUCTURE_TYPE_QUERY_POOL_CREATE_INFO; + qp.queryType = VK_QUERY_TYPE_TIMESTAMP; + qp.queryCount = (PIPELINE_FRAMES + 1) * kTsPairsPerBlock * 2; + const VkResult qres = vkCreateQueryPool(device_.device(), &qp, nullptr, &ts_pool_); + if (qres != VK_SUCCESS) + { + ts_enabled_ = false; + ts_pool_ = VK_NULL_HANDLE; + std::fprintf(stderr, "[profile] timestamp query pool 创建失败(%d)，设备侧计时降级\n", + static_cast(qres)); + } } - [[nodiscard]] StagingRing& staging_ring() noexcept { return *staging_ring_; } - [[nodiscard]] VkCommandPool command_pool() const noexcept { return command_pool_; } - [[nodiscard]] VkDescriptorPool gpu_tensor_pool() const noexcept { return gpu_tensor_pool_; } - - // ── 非阻塞收割所有已完成帧（不含空闲块归还）──────────────────── - // 显存峰值关键项：收割必须紧随每个帧提交点（submit_frame_no_wait）。 - // 收割若推迟到 step 边界，步内已析构中间张量的延迟销毁会一直挂在 - // "已完成但未被收割"的帧上、内存不还池 → backward 只能继续申请新底材 - // （实测 bench 配置 transient live 3837MB 里 pending（已析构未还池） - // 高达 3000MB，步内峰值 ≈ 稳态 2.8×）。 - [[nodiscard]] Result reap_completed_frames() + for (std::size_t i = 0; i < PIPELINE_FRAMES; ++i) + frames_[i].ts.base = static_cast(i * kTsPairsPerBlock * 2); + ts_solo_.base = static_cast(PIPELINE_FRAMES * kTsPairsPerBlock * 2); + +#ifdef NN_FUSED_REGISTRY_EMBEDDED + // 16. 注册 AOT 融合 shader pipelines（构建期 scan_exprs 收集 + + // fused_generate.hpp 合成）。条目形态： + // · V0 / V2 / V1s / V1t（键 `key` / `key#a` / `key#s` / + // `key#xxxx` 特化签名）：N 输入 + 1 输出 binding； + // · V1 运行期精度分派（键 `key#x`）：2N + 2 binding（每个输入/输出 + // 各绑 f32 与 f16 两个视图）；PC 固定头末尾多一个 `prec` 槽。 + // push constants = 固定头（count/cols/…）+ vp 槽 + 常量池 + rp 槽。 + for (const auto& fs : nn::fused::kFusedShaders) { - if (!initialized_) - return {}; - for (std::size_t i = 0; i < frames_.size(); ++i) + if (!fs.spirv || fs.spirv_words == 0) continue; // 空 SPIR-V：跳过（构建配置缺失） + // 运行期精度分派变体（键 "#x"）：输入/输出各双视图 + PC 多一个 prec 槽 + const bool is_disp = std::string_view(fs.key).ends_with(nn::EXPR_PREC_DISPATCH_SUFFIX); + // P3 批量变体（键 "#b"）：单 binding（实例表），PC 与 V0 同布局 + const bool is_batch = std::string_view(fs.key).ends_with(nn::EXPR_BATCH_SUFFIX); + // P0 细 tile 变体（键 "#f"）：V0 同布局（同绑定数、同 PC）， + // 只是另一套网格粒度（输出块 16×16）的 f32 shader。 + const bool is_fine = std::string_view(fs.key).ends_with(nn::EXPR_MM_FINE_SUFFIX); + // K 切分变体（键 "#ks"，每 WG 一段 K → 部分和）与其伴生归约 + // （键 "#kr"）：f32 变体，PC/绑定见 EXPR_MM_KSPLIT_SUFFIX 注释。 + const bool is_ks = std::string_view(fs.key).ends_with(nn::EXPR_MM_KSPLIT_SUFFIX); + const bool is_kr = std::string_view(fs.key).ends_with(nn::EXPR_MM_KSPLIT_REDUCE_SUFFIX); + // 带 '#' 的 f16 变体（#x / #s / #xxxx 特化 / #a）需要 SSBO 16 位 + // 存储；设备未启用 → 跳过注册（运行时查不到 → 基类 NVI 入口回退 + // 边界 cast）。`#b` 与**纯 V0 的** `#f`/`#ks`/`#kr` 全 f32 无此 + // 要求（不得被这条误伤）。 + // ⚠ 后缀**可叠加**：`key#10003#ks` = f16 特化的 K 切分，其 shader + // 里有 float16_t 缓冲 → 仍需要 16 位存储。判据 = 去掉网格粒度 + // 尾后缀（#kr/#ks/#f）之后是否仍含 '#'。 + const std::string_view key_sv(fs.key); + std::string_view grid_stripped = key_sv; + for (std::string_view suf : {std::string_view(nn::EXPR_BATCH_SUFFIX), + std::string_view(nn::EXPR_MM_KSPLIT_REDUCE_SUFFIX), + std::string_view(nn::EXPR_MM_KSPLIT_SUFFIX), + std::string_view(nn::EXPR_MM_FINE_SUFFIX)}) { - if (!frames_[i].in_flight) - continue; - // 非阻塞查询：vkWaitForFences(timeout=0) 返回 VK_SUCCESS - // （已就绪）或 VK_TIMEOUT（未就绪）。VK_NOT_READY 是 - // vkGetFenceStatus 的返回值，vkWaitForFences 不使用它。 - const VkResult rs = - vkWaitForFences(device_.device(), 1, &frames_[i].fence, - VK_TRUE, 0); - if (rs == VK_ERROR_DEVICE_LOST) - { - // 设备丢失（TDR 触发）：GPU 已死亡，无法恢复 - device_lost_ = true; - NN_FAIL("GPU 设备丢失 (VK_ERROR_DEVICE_LOST): Windows TDR 已重置 GPU 驱动。" "\n模型已自动保存，请使用 --resume 重启训练。" "\n建议：减小 --batch-size 或 --seq-len，或增大 Windows TDR 超时" " (注册表 TdrDelay)"); - } - if (rs != VK_TIMEOUT && rs != VK_SUCCESS) - NN_FAIL(std::string("GPU fence 非阻塞查询失败: Vulkan error ") + std::to_string(static_cast(rs))); - if (rs == VK_SUCCESS) + if (grid_stripped.size() > suf.size() && grid_stripped.ends_with(suf)) { - auto rr = reap_frame(i, /*wait=*/false); - NN_TRY_CHECK(rr); + grid_stripped.remove_suffix(suf.size()); + break; } } - return {}; + const bool needs_16bit_storage = key_sv.find('#') != std::string_view::npos && + !((is_batch || is_fine || is_ks || is_kr) && + grid_stripped.find('#') == std::string_view::npos); + if (needs_16bit_storage && !device_.has_16bit_storage()) continue; + // `#b` 批量变体需要 buffer device address（实例表按 64 位地址解引用 + // 各实例缓冲）；未启用 → 跳过注册 → 运行期逐 kernel 降级（可观测）。 + if (is_batch && !device_.bda_available()) continue; + // ALU 变体（key 以 "#a" 结尾，native16 原生 f16 算术）额外需要 + // shaderFloat16 特性；未启用 → 跳过 → 运行时回退 `#x`（其 pipeline + // 独立存在）或边界 cast——正确性不变。 + if (std::string_view(fs.key).ends_with(nn::EXPR_PREC_ALU_SUFFIX) && + !device_.has_shader_float16()) + continue; + const std::uint32_t n_in = static_cast(fs.spec.views.size()); + // 分派变体：输入 i 的 f32 视图 binding=i、f16 视图 binding=n_in+i； + // 输出 f32/f16 视图 binding=2*n_in / 2*n_in+1。 + // P2 多输出（extras）：V0 追加 k 个输出 binding（bex，主输出之后）； + // `#x`/`#a` 变体不生成 extras（生成器跳过）⇒ 分派分支恒 2N+2。 + // `#b` 批量变体：只绑 binding=0 一个 storage buffer（实例表）。 + const std::uint32_t num_bindings = + is_batch ? 1u + : is_disp ? (2u * n_in + 2u) + : (n_in + 1u + static_cast(fs.spec.extras.size())); + // push constant **固定头**槽数（uint）：逐形态与生成器的 PC 声明 + // 逐槽一致—— + // 逐元素 2（count, cols）/ +matmul 6（…, rows, mm_k, mm_batch, + // mm_trans）/ 归约 4（…, rows, vector_out）/ 归约+matmul 7 / + // fold 5（…, vector_out, fold_k）/ fold+mm 7（…, fold_k, mm_k, + // mm_batch）+ 分派变体 1 个 `prec`。 + // ⚠ 该值**只在此处计算一次**并存进 FusedKernel::pc_base_fixed，写入侧 + // (run_fused_gpu) 只读：此前两侧各推导一遍，漂移会让 + // vkCmdPushConstants 超 range 而被驱动丢弃 → shader 读到未定义 + // 残留（随前序 op 漂移，表现为时对时错的假 PASS，"全绿"验不出）。 + const std::uint32_t pc_base_fixed = is_ks ? 7u // K 切分：matmul 6 槽 + ks_span + : (fs.spec.fold && fs.spec.fold->matmul) + ? 7u // fold+mm 7 槽 + : (fs.spec.fold) ? 5u + : (fs.reduce_axis >= 0 && fs.has_matmul) ? 7u + : (fs.reduce_axis >= 0) ? 4u + : (fs.has_matmul) ? 6u + : 2u; + // 分派变体在固定头末尾多一个 uint prec（生成器同位置声明） + const std::uint32_t pc_base = pc_base_fixed + (is_disp ? 1u : 0u); + const std::uint32_t pc_uints = pc_base + fs.view_param_count; + const std::uint32_t pc_size = static_cast( + pc_uints * sizeof(std::uint32_t) + sizeof(Scalar) * fs.spec.consts.size() + + sizeof(Scalar) * fs.spec.rparams.size()); + auto fp_r = VulkanPipeline::create_generic( + device_.device(), std::span(fs.spirv, fs.spirv_words), + num_bindings, pc_size); + if (fp_r) + { + FusedKernel kernel; + kernel.pipeline = std::move(*fp_r); + kernel.reduce_axis = fs.reduce_axis; + kernel.has_matmul = fs.has_matmul != 0; + kernel.fold = FoldMeta{ + fs.spec.fold ? (fs.spec.fold->vec_state_len > 0 ? fs.spec.fold->vec_state_len : 1u) + : 0u, + // per_thread_row 判据与 generate_glsl 分派同源（v2 = 有 + // vec 态或 mm 段；v1 = 纯标量） + static_cast( + fs.spec.fold && !(fs.spec.fold->vec_state_len > 0 || fs.spec.fold->matmul))}; + kernel.view_param_count = fs.view_param_count; + kernel.rparam_count = fs.rparam_count; + kernel.vec_width = fs.vec_width; + kernel.pc_base_fixed = pc_base_fixed; + // P3 批量变体：实例表 stride（uvec2 槽数），run_fused_gpu_batch + // 据此校验实例布局（生成器 S = 输入+主输出+extras+尾槽，同源） + if (is_batch) + kernel.batch_stride = + static_cast(n_in + 1u + fs.spec.extras.size() + 1u); + fused_kernels_.emplace(fs.key, std::move(kernel)); + } } +#endif - // ── 显存回收（L2）：先非阻塞收割已完成帧，再归还完全空闲的池底材 ── - // 调用时机：step 边界（end_batch 提交完成、延迟销毁已收割之后）。 - // 未完成的帧留给环槽复用 / 下次 reap——不在 step 边界阻塞流水线。 - [[nodiscard]] Result release_idle_pool_blocks() + // ── 查询 GPU f16 硬件能力（§7.1）───────────────────────────── + // shaderFloat16：Vulkan 1.2 核心特性 / VK_KHR_shader_float16_int8 扩展。 + // 以**设备侧已启用**状态为准（VulkanDevice::initialize 查询 + 启用 + + // NN_VULKAN_NO_16BIT_ALU 逃生阀）——只查物理特性不启用 = pipeline 创建 + // 时校验失败，故此处不独立查询。 + has_shader_float16_ = device_.has_shader_float16(); + + initialized_ = true; + return {}; +} + +[[nodiscard]] bool is_initialized() const noexcept +{ return initialized_; } +[[nodiscard]] bool gpu_available() const noexcept +{ return initialized_; } +// ── Pipeline 变体可用性 ───────────────────────────────────────── +// 唯一查询入口：id 见 compute_vk_pipeline_catalog.hpp；缺失 = 该变体 +// 未创建（SPIR-V 未嵌入 / 设备能力不足 / 创建失败）。 +[[nodiscard]] bool has(PipelineId id) const noexcept +{ return pipelines_.has(id); } + +// 扫描族 f16 变体的整体可用性：四条路径齐备才算接通（缺任一 → 调用方 +// 回退边界 cast，避免"部分直读、部分 cast"的混合语义）。 +[[nodiscard]] bool has_scan_f16_pipelines() const noexcept +{ + return has(PipelineId::scan_prefix_outer_f16) && has(PipelineId::scan_suffix_outer_f16) && + has(PipelineId::scan_prefix_outer_gen_f16) && has(PipelineId::scan_suffix_outer_gen_f16); +} + +// ── f16 硬件能力查询（§7.1）──────────────────────────────────── +[[nodiscard]] bool has_shader_float16() const noexcept +{ return has_shader_float16_; } + +[[nodiscard]] VulkanDevice& device() noexcept +{ return device_; } +[[nodiscard]] MemoryPool& memory_pool() noexcept +{ return *memory_pool_; } +[[nodiscard]] MemoryPool& transient_pool() noexcept +{ return *transient_pool_; } +// 延迟销毁队列总字节（探针归因用）：已析构但因帧未完成而尚未归还池的 +// 内存。解释"Tensor 已置空但 live 不降"的账目缺口。 +[[nodiscard]] VkDeviceSize pending_destroy_bytes() +{ + std::lock_guard lock(pending_mutex_); + VkDeviceSize total = 0; + for (const auto& pd : pending_destroys_) + total += pd.alloc.size; + return total; +} +// 瞬态/持久分池选择器：batch 录制期→瞬态池，否则→持久池（纯组织性）。 +// 供 GpuTensorT 分配（create_empty / from_matrix / create_host_visible_empty） +// 使用，使批量内的临时/激活与构建期的参数/权重分池驻留。 +[[nodiscard]] MemoryPool& alloc_pool() noexcept +{ return batch_mode_ ? *transient_pool_ : *memory_pool_; } +// 异步标量回读槽位数：供 GpuEngine 暴露给调用方做环形复用。 +[[nodiscard]] static constexpr std::size_t scalar_readback_slot_count() noexcept +{ return SCALAR_READBACK_SLOTS; } +[[nodiscard]] StagingRing& staging_ring() noexcept +{ return *staging_ring_; } +[[nodiscard]] VkCommandPool command_pool() const noexcept +{ return command_pool_; } +[[nodiscard]] VkDescriptorPool gpu_tensor_pool() const noexcept +{ return gpu_tensor_pool_; } + +// ── 非阻塞收割所有已完成帧（不含空闲块归还）──────────────────── +// 显存峰值关键项：收割必须紧随每个帧提交点（submit_frame_no_wait）。 +// 收割若推迟到 step 边界，步内已析构中间张量的延迟销毁会一直挂在 +// "已完成但未被收割"的帧上、内存不还池 → backward 只能继续申请新底材 +// （实测 bench 配置 transient live 3837MB 里 pending（已析构未还池） +// 高达 3000MB，步内峰值 ≈ 稳态 2.8×）。 +[[nodiscard]] Result reap_completed_frames() +{ + if (!initialized_) return {}; + for (std::size_t i = 0; i < frames_.size(); ++i) { - auto r = reap_completed_frames(); - if (!r) - return r; - if (memory_pool_) - memory_pool_->release_idle_blocks(); - if (transient_pool_) - transient_pool_->release_idle_blocks(); - return {}; + if (!frames_[i].in_flight) continue; + // 非阻塞查询：vkWaitForFences(timeout=0) 返回 VK_SUCCESS + // （已就绪）或 VK_TIMEOUT（未就绪）。VK_NOT_READY 是 + // vkGetFenceStatus 的返回值，vkWaitForFences 不使用它。 + const VkResult rs = vkWaitForFences(device_.device(), 1, &frames_[i].fence, VK_TRUE, 0); + if (rs == VK_ERROR_DEVICE_LOST) + { + // 设备丢失（TDR 触发）：GPU 已死亡，无法恢复 + device_lost_ = true; + NN_FAIL("GPU 设备丢失 (VK_ERROR_DEVICE_LOST): Windows TDR 已重置 GPU 驱动。" + "\n模型已自动保存，请使用 --resume 重启训练。" + "\n建议：减小 --batch-size 或 --seq-len，或增大 Windows TDR 超时" + " (注册表 TdrDelay)"); + } + if (rs != VK_TIMEOUT && rs != VK_SUCCESS) + NN_FAIL(std::string("GPU fence 非阻塞查询失败: Vulkan error ") + + std::to_string(static_cast(rs))); + if (rs == VK_SUCCESS) + { + auto rr = reap_frame(i, /*wait=*/false); + NN_TRY_CHECK(rr); + } } - - + return {}; +} + +// ── 显存回收（L2）：先非阻塞收割已完成帧，再归还完全空闲的池底材 ── +// 调用时机：step 边界（end_batch 提交完成、延迟销毁已收割之后）。 +// 未完成的帧留给环槽复用 / 下次 reap——不在 step 边界阻塞流水线。 +[[nodiscard]] Result release_idle_pool_blocks() +{ + auto r = reap_completed_frames(); + if (!r) return r; + if (memory_pool_) memory_pool_->release_idle_blocks(); + if (transient_pool_) transient_pool_->release_idle_blocks(); + return {}; +} diff --git a/include/neuralnet.cpp/backend/compute_vk_backend_frame.inc b/include/neuralnet.cpp/backend/compute_vk_backend_frame.inc index 96882eb7..44fabf79 100644 --- a/include/neuralnet.cpp/backend/compute_vk_backend_frame.inc +++ b/include/neuralnet.cpp/backend/compute_vk_backend_frame.inc @@ -4,769 +4,764 @@ // ⚠ 本文件不是独立头：仅由 backend/compute_vk_backend.hpp 在 // `class GpuBackend { ... }` 体内**按序** #include（顺序即成员序）。 // ═══════════════════════════════════════════════════════════════════════════ - // ══════════════════════════════════════════════════════════════════ - // Command Buffer Batching API（多帧流水线） - // ══════════════════════════════════════════════════════════════════ - // 用法： - // backend.begin_batch(); - // // ... 多次 matmul_gpu / elementwise_gpu / layernorm_gpu 调用 ... - // backend.end_batch(); - // 效果：所有操作录制到同一个 command buffer，一次提交（**不等待**）； - // host 继续录制下一段，GPU 在队列上先行执行（host 录制与 GPU 执行 - // 重叠）。真正要读 GPU 数据时调 wait_in_flight()（to_matrix 内部）。 - // ══════════════════════════════════════════════════════════════════ - [[nodiscard]] bool in_batch() const noexcept { return batch_mode_; } - - // ── 取下一帧并开始录制（begin_batch / flush_batch 共用）────────── - // 环槽复用前等待：该帧若仍在飞行，等其 fence（这是唯一的常规等待点； - // 环深度 N 限制 host 至多超前 N 帧）。 - [[nodiscard]] Result start_frame() +// ══════════════════════════════════════════════════════════════════ +// Command Buffer Batching API（多帧流水线） +// ══════════════════════════════════════════════════════════════════ +// 用法： +// backend.begin_batch(); +// // ... 多次 matmul_gpu / elementwise_gpu / layernorm_gpu 调用 ... +// backend.end_batch(); +// 效果：所有操作录制到同一个 command buffer，一次提交（**不等待**）； +// host 继续录制下一段，GPU 在队列上先行执行（host 录制与 GPU 执行 +// 重叠）。真正要读 GPU 数据时调 wait_in_flight()（to_matrix 内部）。 +// ══════════════════════════════════════════════════════════════════ +[[nodiscard]] bool in_batch() const noexcept +{ return batch_mode_; } + +// ── 取下一帧并开始录制（begin_batch / flush_batch 共用）────────── +// 环槽复用前等待：该帧若仍在飞行，等其 fence（这是唯一的常规等待点； +// 环深度 N 限制 host 至多超前 N 帧）。 +[[nodiscard]] Result start_frame() +{ + const std::size_t i = frame_next_; + frame_next_ = (i + 1) % PIPELINE_FRAMES; + + auto rr = reap_frame(i, /*wait=*/true); + if (!rr) { - const std::size_t i = frame_next_; - frame_next_ = (i + 1) % PIPELINE_FRAMES; + batch_mode_ = false; + NN_TRY_CHECK(rr); + } - auto rr = reap_frame(i, /*wait=*/true); - if (!rr) - { - batch_mode_ = false; - NN_TRY_CHECK(rr); - } + auto& f = frames_[i]; + auto r = detail::vk_check(vkResetCommandBuffer(f.cmd, 0), __FILE__, __LINE__); + NN_TRY_CHECK(r); + VkCommandBufferBeginInfo begin_info{}; + begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO; + begin_info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT; + r = detail::vk_check(vkBeginCommandBuffer(f.cmd, &begin_info), __FILE__, __LINE__); + NN_TRY_CHECK(r); + batch_cmd_ = f.cmd; + batch_frame_ = i; + last_active_frame_ = i; + batch_has_ops_ = false; + batch_mode_ = true; + ts_begin_block(f.cmd, f.ts); // 设备侧计时：本帧块开始 + return {}; +} + +// ── 提交当前帧（不等待）──────────────────────────────────────────── +// 空帧不提交（无 op 的帧跳过，省一次 vkQueueSubmit + 一个环槽）。 +[[nodiscard]] Result submit_frame_no_wait() +{ + auto& f = frames_[batch_frame_]; + if (!batch_has_ops_) return {}; // 空帧：不提交 + + // reap 已确保 fence 处于 unsignaled（或从未提交）状态 + auto r = detail::vk_check(vkResetFences(device_.device(), 1, &f.fence), __FILE__, __LINE__); + NN_TRY_CHECK(r); + // 跨 submit 数据依赖：帧命令的输入可能刚由 batch 内的 from_matrix + // 独立上传提交写入——等 in-flight region 信号量（GPU 队列内等待， + // host 不阻塞） + const auto sw = collect_staging_waits(); - auto& f = frames_[i]; - auto r = detail::vk_check(vkResetCommandBuffer(f.cmd, 0), __FILE__, __LINE__); - NN_TRY_CHECK(r); - VkCommandBufferBeginInfo begin_info{}; - begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO; - begin_info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT; - r = detail::vk_check(vkBeginCommandBuffer(f.cmd, &begin_info), __FILE__, __LINE__); - NN_TRY_CHECK(r); - batch_cmd_ = f.cmd; - batch_frame_ = i; - last_active_frame_ = i; - batch_has_ops_ = false; - batch_mode_ = true; - ts_begin_block(f.cmd, f.ts); // 设备侧计时：本帧块开始 - return {}; + { + std::lock_guard lock(queue_mutex_); + VkSubmitInfo submit_info{}; + submit_info.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO; + submit_info.commandBufferCount = 1; + submit_info.pCommandBuffers = &f.cmd; + VkTimelineSemaphoreSubmitInfo timeline_info{}; + attach_staging_waits(submit_info, sw, timeline_info); + r = detail::vk_check(vkQueueSubmit(device_.compute_queue(), 1, &submit_info, f.fence), + __FILE__, __LINE__); } - - // ── 提交当前帧（不等待）──────────────────────────────────────────── - // 空帧不提交（无 op 的帧跳过，省一次 vkQueueSubmit + 一个环槽）。 - [[nodiscard]] Result submit_frame_no_wait() + NN_TRY_CHECK(r); + f.in_flight = true; + last_active_frame_ = batch_frame_; + // 提交后立即非阻塞收割其它已完成帧：让步内已析构张量的内存尽早 + // 还池，被后续分配复用（否则要等到环槽复用或 step 边界）。 + // NN_NO_EARLY_REAP=1 可关闭（A/B 对照用，见 bench/run_mem_ab.ps1）。 + if (early_reap_enabled_) { - auto& f = frames_[batch_frame_]; - if (!batch_has_ops_) - return {}; // 空帧：不提交 - - // reap 已确保 fence 处于 unsignaled（或从未提交）状态 - auto r = detail::vk_check( - vkResetFences(device_.device(), 1, &f.fence), __FILE__, __LINE__); - NN_TRY_CHECK(r); - // 跨 submit 数据依赖：帧命令的输入可能刚由 batch 内的 from_matrix - // 独立上传提交写入——等 in-flight region 信号量（GPU 队列内等待， - // host 不阻塞） - const auto sw = collect_staging_waits(); - - { - std::lock_guard lock(queue_mutex_); - VkSubmitInfo submit_info{}; - submit_info.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO; - submit_info.commandBufferCount = 1; - submit_info.pCommandBuffers = &f.cmd; - VkTimelineSemaphoreSubmitInfo timeline_info{}; - attach_staging_waits(submit_info, sw, timeline_info); - r = detail::vk_check( - vkQueueSubmit(device_.compute_queue(), 1, &submit_info, f.fence), - __FILE__, __LINE__); - } + auto rr = reap_completed_frames(); + NN_TRY_CHECK(rr); + } + return {}; +} + +[[nodiscard]] Result begin_batch() +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + if (batch_mode_) NN_FAIL("Already in batch mode"); + return start_frame(); +} + +// ── 结束 batch：提交当前帧，**不等待**───────────────────── +// 帧在队列上执行；其延迟销毁 / 描述符集在 reap 时释放。 +// 设备丢失等错误在 wait_in_flight / 环槽复用 reap 时浮出。 +[[nodiscard]] Result end_batch() +{ + if (!batch_mode_) return {}; // 容错：非 batch 模式时 no-op + + // 1. 结束录制 + auto r = detail::vk_check(vkEndCommandBuffer(batch_cmd_), __FILE__, __LINE__); + if (!r) + { + batch_mode_ = false; NN_TRY_CHECK(r); - f.in_flight = true; - last_active_frame_ = batch_frame_; - // 提交后立即非阻塞收割其它已完成帧：让步内已析构张量的内存尽早 - // 还池，被后续分配复用（否则要等到环槽复用或 step 边界）。 - // NN_NO_EARLY_REAP=1 可关闭（A/B 对照用，见 bench/run_mem_ab.ps1）。 - if (early_reap_enabled_) - { - auto rr = reap_completed_frames(); - NN_TRY_CHECK(rr); - } - return {}; } - [[nodiscard]] Result begin_batch() + // 2. 提交（不等待） + r = submit_frame_no_wait(); + batch_mode_ = false; + return r; +} + +// ── 批处理中点刷新（防 TDR，提交不等待）──────────────────────── +// 提交当前帧（不等待），然后开始录制下一帧。用于拆分大 batch（如 +// forward 与 backward 之间），避免单次提交时间过长触发 Windows TDR。 +// 当前帧为空（无 op 录制）时 no-op——不提交、不换帧（避免浪费环槽 +// 与 submit，如 to_matrix 的 drain 已换过帧的场景）。 +[[nodiscard]] Result flush_batch() +{ + if (!batch_mode_) return {}; // 非 batch 模式时 no-op + if (!batch_has_ops_) return {}; // 空帧：无需提交，继续在同一帧录制 + + auto r = detail::vk_check(vkEndCommandBuffer(batch_cmd_), __FILE__, __LINE__); + if (!r) { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - if (batch_mode_) - NN_FAIL("Already in batch mode"); - return start_frame(); + batch_mode_ = false; + NN_TRY_CHECK(r); } - // ── 结束 batch：提交当前帧，**不等待**───────────────────── - // 帧在队列上执行；其延迟销毁 / 描述符集在 reap 时释放。 - // 设备丢失等错误在 wait_in_flight / 环槽复用 reap 时浮出。 - [[nodiscard]] Result end_batch() + r = submit_frame_no_wait(); // 不等待 + if (!r) { - if (!batch_mode_) - return {}; // 容错：非 batch 模式时 no-op - - // 1. 结束录制 - auto r = detail::vk_check(vkEndCommandBuffer(batch_cmd_), __FILE__, __LINE__); - if (!r) - { - batch_mode_ = false; - NN_TRY_CHECK(r); - } - - // 2. 提交（不等待） - r = submit_frame_no_wait(); batch_mode_ = false; - return r; + NN_TRY_CHECK(r); } - // ── 批处理中点刷新（防 TDR，提交不等待）──────────────────────── - // 提交当前帧（不等待），然后开始录制下一帧。用于拆分大 batch（如 - // forward 与 backward 之间），避免单次提交时间过长触发 Windows TDR。 - // 当前帧为空（无 op 录制）时 no-op——不提交、不换帧（避免浪费环槽 - // 与 submit，如 to_matrix 的 drain 已换过帧的场景）。 - [[nodiscard]] Result flush_batch() + // 仍处于 batch 模式：开始录制下一帧 + return start_frame(); +} + +// ── "真正要结果"的阻塞点─────────────────────────────────── +// 等待所有在飞帧完成并 reap（延迟销毁 + 描述符集释放）。 +// to_matrix / copy_from 读 GPU 内存前必须调用：要读的数据可能刚由 +// 在飞帧写入，队列 FIFO 只保证顺序、不保证完成。 +[[nodiscard]] Result wait_in_flight() +{ + for (std::size_t i = 0; i < frames_.size(); ++i) { - if (!batch_mode_) - return {}; // 非 batch 模式时 no-op - if (!batch_has_ops_) - return {}; // 空帧：无需提交，继续在同一帧录制 - - auto r = detail::vk_check(vkEndCommandBuffer(batch_cmd_), __FILE__, __LINE__); - if (!r) + if (frames_[i].in_flight) { - batch_mode_ = false; + auto r = reap_frame(i, /*wait=*/true); NN_TRY_CHECK(r); } - - r = submit_frame_no_wait(); // 不等待 - if (!r) - { - batch_mode_ = false; - NN_TRY_CHECK(r); - } - - // 仍处于 batch 模式：开始录制下一帧 - return start_frame(); } - - // ── "真正要结果"的阻塞点─────────────────────────────────── - // 等待所有在飞帧完成并 reap（延迟销毁 + 描述符集释放）。 - // to_matrix / copy_from 读 GPU 内存前必须调用：要读的数据可能刚由 - // 在飞帧写入，队列 FIFO 只保证顺序、不保证完成。 - [[nodiscard]] Result wait_in_flight() + // 所有帧已完成：兜底清理剩余延迟销毁（正常应为空） + flush_pending_destroys(); + return {}; +} + +// ── 阻塞式上传：CPU → GPU（支持分块传输大矩阵）───────────────────── +// 当数据超过 staging region 大小时，自动分块上传。 +// 每块大小不超过 staging region，通过多次 memcpy + vkCmdCopyBuffer 完成。 +// 元素类型由 P 决定（f32=4B / f16=2B），memcpy 为字节级操作，§6.5。 +template +[[nodiscard]] Result upload_blocking(GpuTensorT

& dst, std::span> cpu_data) +{ + using ElemType = elem

; + if (!initialized_) NN_FAIL("GPU backend not initialized"); + if (!dst.valid()) NN_FAIL("Invalid destination GpuTensor"); + const std::size_t elem_count = dst.rows() * dst.cols(); + if (cpu_data.size() != elem_count) NN_FAIL("Upload size mismatch"); + const std::size_t total_bytes = elem_count * sizeof(ElemType); + const std::size_t staging_cap = staging_ring_->region_size(); + + // ── 小矩阵快速路径：单次传输 ────────────────────────────────── + // region 专属 command buffer：上传 copy 录在 region 自己的 + // cmd 上，提交后既不 free 也不立刻复用——VUID- + // vkFreeCommandBuffers-pCommandBuffers-00058 禁止释放 pending + // （已提交、fence 未 signal）的 command buffer；"submit 后立即 + // free"是规范违规（实测导致驱动通道排序失效、fence 提前 + // signal、跨 submit 乱序）。acquire 已等在飞 fence → 该 cmd 的 + // 上一次使用完成（invalid 状态）→ vkResetCommandBuffer 复用合法。 + if (total_bytes <= staging_cap) { - for (std::size_t i = 0; i < frames_.size(); ++i) - { - if (frames_[i].in_flight) - { - auto r = reap_frame(i, /*wait=*/true); - NN_TRY_CHECK(r); - } - } - // 所有帧已完成：兜底清理剩余延迟销毁（正常应为空） - flush_pending_destroys(); - return {}; - } + auto ri = staging_ring_->acquire(); - // ── 阻塞式上传：CPU → GPU（支持分块传输大矩阵）───────────────────── - // 当数据超过 staging region 大小时，自动分块上传。 - // 每块大小不超过 staging region，通过多次 memcpy + vkCmdCopyBuffer 完成。 - // 元素类型由 P 决定（f32=4B / f16=2B），memcpy 为字节级操作，§6.5。 - template - [[nodiscard]] Result upload_blocking( - GpuTensorT

& dst, std::span> cpu_data) - { - using ElemType = elem

; - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - if (!dst.valid()) - NN_FAIL("Invalid destination GpuTensor"); - const std::size_t elem_count = dst.rows() * dst.cols(); - if (cpu_data.size() != elem_count) - NN_FAIL("Upload size mismatch"); - const std::size_t total_bytes = elem_count * sizeof(ElemType); - const std::size_t staging_cap = staging_ring_->region_size(); - - // ── 小矩阵快速路径：单次传输 ────────────────────────────────── - // region 专属 command buffer：上传 copy 录在 region 自己的 - // cmd 上，提交后既不 free 也不立刻复用——VUID- - // vkFreeCommandBuffers-pCommandBuffers-00058 禁止释放 pending - // （已提交、fence 未 signal）的 command buffer；"submit 后立即 - // free"是规范违规（实测导致驱动通道排序失效、fence 提前 - // signal、跨 submit 乱序）。acquire 已等在飞 fence → 该 cmd 的 - // 上一次使用完成（invalid 状态）→ vkResetCommandBuffer 复用合法。 - if (total_bytes <= staging_cap) - { - auto ri = staging_ring_->acquire(); + auto r = staging_ring_->upload(ri, cpu_data, 0); + if (!r) return r; - auto r = staging_ring_->upload(ri, cpu_data, 0); - if (!r) return r; + auto cmd = staging_ring_->command_buffer(ri); - auto cmd = staging_ring_->command_buffer(ri); + r = detail::vk_check(vkResetCommandBuffer(cmd, 0), __FILE__, __LINE__); + if (!r) return r; - r = detail::vk_check(vkResetCommandBuffer(cmd, 0), __FILE__, __LINE__); - if (!r) return r; + VkCommandBufferBeginInfo begin_info{}; + begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO; + begin_info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT; - VkCommandBufferBeginInfo begin_info{}; - begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO; - begin_info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT; + r = detail::vk_check(vkBeginCommandBuffer(cmd, &begin_info), __FILE__, __LINE__); + if (!r) return r; - r = detail::vk_check(vkBeginCommandBuffer(cmd, &begin_info), __FILE__, __LINE__); - if (!r) return r; + VkBufferCopy cp{0, 0, total_bytes}; + vkCmdCopyBuffer(cmd, staging_ring_->buffer(ri), dst.buffer().impl(), 1, &cp); - VkBufferCopy cp{0, 0, total_bytes}; - vkCmdCopyBuffer(cmd, staging_ring_->buffer(ri), dst.buffer().impl(), 1, &cp); + r = detail::vk_check(vkEndCommandBuffer(cmd), __FILE__, __LINE__); + if (!r) return r; - r = detail::vk_check(vkEndCommandBuffer(cmd), __FILE__, __LINE__); - if (!r) return r; + auto fence = staging_ring_->fence(ri); + vkResetFences(device_.device(), 1, &fence); - auto fence = staging_ring_->fence(ri); - vkResetFences(device_.device(), 1, &fence); + { + std::lock_guard lock(queue_mutex_); + VkSubmitInfo submit_info{}; + submit_info.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO; + submit_info.commandBufferCount = 1; + submit_info.pCommandBuffers = &cmd; + // 跨 submit 数据依赖：时间线信号量 signal 本次 + // 上传的 value。读取"本 copy 写入的 buffer"的后续 submit 必须 + // wait 该 value（见 collect_staging_waits）。 + VkTimelineSemaphoreSubmitInfo timeline_info{}; + VkSemaphore sem = VK_NULL_HANDLE; + std::uint64_t signal_value = 0; + attach_upload_signal(submit_info, *staging_ring_, ri, timeline_info, sem, signal_value); - { - std::lock_guard lock(queue_mutex_); - VkSubmitInfo submit_info{}; - submit_info.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO; - submit_info.commandBufferCount = 1; - submit_info.pCommandBuffers = &cmd; - // 跨 submit 数据依赖：时间线信号量 signal 本次 - // 上传的 value。读取"本 copy 写入的 buffer"的后续 submit 必须 - // wait 该 value（见 collect_staging_waits）。 - VkTimelineSemaphoreSubmitInfo timeline_info{}; - VkSemaphore sem = VK_NULL_HANDLE; - std::uint64_t signal_value = 0; - attach_upload_signal(submit_info, *staging_ring_, ri, - timeline_info, sem, signal_value); - - r = detail::vk_check( - vkQueueSubmit(device_.compute_queue(), 1, &submit_info, fence), - __FILE__, __LINE__); - } - if (!r) return r; - - // 提交后不等待——host 可立即继续录制。标记 region - // in use：下次 acquire 到该 region 时等其 fence（staging 安全 - // + 等待后该 cmd 离开 pending，可 reset 复用）。**不释放 cmd** - // （pending 状态，见上方 VUID 注释）。 - staging_ring_->mark_in_flight(ri); - return {}; + r = detail::vk_check(vkQueueSubmit(device_.compute_queue(), 1, &submit_info, fence), + __FILE__, __LINE__); } + if (!r) return r; - // ── 大矩阵分块上传 ─────────────────────────────────────────── - // 每块大小对齐到 sizeof(ElemType)，确保元素边界对齐 - const std::size_t chunk_bytes = (staging_cap / sizeof(ElemType)) * sizeof(ElemType); - std::size_t dst_offset = 0; + // 提交后不等待——host 可立即继续录制。标记 region + // in use：下次 acquire 到该 region 时等其 fence（staging 安全 + // + 等待后该 cmd 离开 pending，可 reset 复用）。**不释放 cmd** + // （pending 状态，见上方 VUID 注释）。 + staging_ring_->mark_in_flight(ri); + return {}; + } - while (dst_offset < total_bytes) - { - const std::size_t this_chunk = std::min(chunk_bytes, total_bytes - dst_offset); - const std::size_t this_elems = this_chunk / sizeof(ElemType); - const std::size_t src_elem_offset = dst_offset / sizeof(ElemType); + // ── 大矩阵分块上传 ─────────────────────────────────────────── + // 每块大小对齐到 sizeof(ElemType)，确保元素边界对齐 + const std::size_t chunk_bytes = (staging_cap / sizeof(ElemType)) * sizeof(ElemType); + std::size_t dst_offset = 0; - // 1. 获取 staging region 并上传当前块 - auto ri = staging_ring_->acquire(); - auto r = staging_ring_->upload( - ri, - std::span(cpu_data.data() + src_elem_offset, this_elems), - 0); - if (!r) return r; + while (dst_offset < total_bytes) + { + const std::size_t this_chunk = std::min(chunk_bytes, total_bytes - dst_offset); + const std::size_t this_elems = this_chunk / sizeof(ElemType); + const std::size_t src_elem_offset = dst_offset / sizeof(ElemType); - // 2. 录制 copy 命令（region 专属 cmd buffer：禁止 pending 释放， - // 见小矩阵路径的 VUID 注释；acquire 已等在飞 fence，reset 合法） - auto cmd = staging_ring_->command_buffer(ri); + // 1. 获取 staging region 并上传当前块 + auto ri = staging_ring_->acquire(); + auto r = staging_ring_->upload( + ri, std::span(cpu_data.data() + src_elem_offset, this_elems), 0); + if (!r) return r; - r = detail::vk_check(vkResetCommandBuffer(cmd, 0), __FILE__, __LINE__); - if (!r) return r; + // 2. 录制 copy 命令（region 专属 cmd buffer：禁止 pending 释放， + // 见小矩阵路径的 VUID 注释；acquire 已等在飞 fence，reset 合法） + auto cmd = staging_ring_->command_buffer(ri); - VkCommandBufferBeginInfo begin_info{}; - begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO; - begin_info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT; + r = detail::vk_check(vkResetCommandBuffer(cmd, 0), __FILE__, __LINE__); + if (!r) return r; - r = detail::vk_check(vkBeginCommandBuffer(cmd, &begin_info), __FILE__, __LINE__); - if (!r) return r; + VkCommandBufferBeginInfo begin_info{}; + begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO; + begin_info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT; - VkBufferCopy cp{0, dst_offset, this_chunk}; - vkCmdCopyBuffer(cmd, staging_ring_->buffer(ri), dst.buffer().impl(), 1, &cp); + r = detail::vk_check(vkBeginCommandBuffer(cmd, &begin_info), __FILE__, __LINE__); + if (!r) return r; - r = detail::vk_check(vkEndCommandBuffer(cmd), __FILE__, __LINE__); - if (!r) return r; + VkBufferCopy cp{0, dst_offset, this_chunk}; + vkCmdCopyBuffer(cmd, staging_ring_->buffer(ri), dst.buffer().impl(), 1, &cp); - // 3. 提交（不等待；region 标记 in use，下次 acquire 到该 - // region 时等其 fence；不释放 cmd——pending 状态）+ 信号本 - // region 信号量（跨 submit 数据依赖，见 collect_staging_waits） - auto fence = staging_ring_->fence(ri); - vkResetFences(device_.device(), 1, &fence); + r = detail::vk_check(vkEndCommandBuffer(cmd), __FILE__, __LINE__); + if (!r) return r; - { - std::lock_guard lock(queue_mutex_); - VkSubmitInfo submit_info{}; - submit_info.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO; - submit_info.commandBufferCount = 1; - submit_info.pCommandBuffers = &cmd; - // 跨 submit 数据依赖：时间线信号量 signal 本次上传的 value - // （读取本 copy 写入 buffer 的后续 submit 会 wait 该 value） - VkTimelineSemaphoreSubmitInfo timeline_info{}; - VkSemaphore sem = VK_NULL_HANDLE; - std::uint64_t signal_value = 0; - attach_upload_signal(submit_info, *staging_ring_, ri, - timeline_info, sem, signal_value); - - r = detail::vk_check( - vkQueueSubmit(device_.compute_queue(), 1, &submit_info, fence), - __FILE__, __LINE__); - } - if (!r) return r; + // 3. 提交（不等待；region 标记 in use，下次 acquire 到该 + // region 时等其 fence；不释放 cmd——pending 状态）+ 信号本 + // region 信号量（跨 submit 数据依赖，见 collect_staging_waits） + auto fence = staging_ring_->fence(ri); + vkResetFences(device_.device(), 1, &fence); + + { + std::lock_guard lock(queue_mutex_); + VkSubmitInfo submit_info{}; + submit_info.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO; + submit_info.commandBufferCount = 1; + submit_info.pCommandBuffers = &cmd; + // 跨 submit 数据依赖：时间线信号量 signal 本次上传的 value + // （读取本 copy 写入 buffer 的后续 submit 会 wait 该 value） + VkTimelineSemaphoreSubmitInfo timeline_info{}; + VkSemaphore sem = VK_NULL_HANDLE; + std::uint64_t signal_value = 0; + attach_upload_signal(submit_info, *staging_ring_, ri, timeline_info, sem, signal_value); - staging_ring_->mark_in_flight(ri); - dst_offset += this_chunk; + r = detail::vk_check(vkQueueSubmit(device_.compute_queue(), 1, &submit_info, fence), + __FILE__, __LINE__); } + if (!r) return r; - return {}; + staging_ring_->mark_in_flight(ri); + dst_offset += this_chunk; } - // ── 阻塞式下载：GPU → CPU（支持分块传输大矩阵）───────────────────── - // 当数据超过 staging region 大小时，自动分块下载。 - // 元素类型由 P 决定（f32=4B / f16=2B），memcpy 为字节级操作，§6.5。 - template - [[nodiscard]] Result download_blocking( - const GpuTensorT

& src, std::span> cpu_data) + return {}; +} + +// ── 阻塞式下载：GPU → CPU（支持分块传输大矩阵）───────────────────── +// 当数据超过 staging region 大小时，自动分块下载。 +// 元素类型由 P 决定（f32=4B / f16=2B），memcpy 为字节级操作，§6.5。 +template +[[nodiscard]] Result download_blocking(const GpuTensorT

& src, std::span> cpu_data) +{ + using ElemType = elem

; + if (!initialized_) NN_FAIL("GPU backend not initialized"); + if (!src.valid()) NN_FAIL("Invalid source GpuTensor"); + const std::size_t elem_count = src.rows() * src.cols(); + if (cpu_data.size() != elem_count) NN_FAIL("Download size mismatch"); + const std::size_t total_bytes = elem_count * sizeof(ElemType); + const std::size_t staging_cap = staging_ring_->region_size(); + + // ── 小矩阵快速路径：单次传输 ────────────────────────────────── + if (total_bytes <= staging_cap) { - using ElemType = elem

; - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - if (!src.valid()) - NN_FAIL("Invalid source GpuTensor"); - const std::size_t elem_count = src.rows() * src.cols(); - if (cpu_data.size() != elem_count) - NN_FAIL("Download size mismatch"); - const std::size_t total_bytes = elem_count * sizeof(ElemType); - const std::size_t staging_cap = staging_ring_->region_size(); - - // ── 小矩阵快速路径：单次传输 ────────────────────────────────── - if (total_bytes <= staging_cap) - { - auto ri = staging_ring_->acquire(); + auto ri = staging_ring_->acquire(); - VkCommandBufferAllocateInfo cmd_alloc{}; - cmd_alloc.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO; - cmd_alloc.commandPool = command_pool_; - cmd_alloc.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; - cmd_alloc.commandBufferCount = 1; + VkCommandBufferAllocateInfo cmd_alloc{}; + cmd_alloc.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO; + cmd_alloc.commandPool = command_pool_; + cmd_alloc.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; + cmd_alloc.commandBufferCount = 1; - VkCommandBuffer cmd = VK_NULL_HANDLE; - auto r = detail::vk_check( - vkAllocateCommandBuffers(device_.device(), &cmd_alloc, &cmd), - __FILE__, __LINE__); - if (!r) return r; + VkCommandBuffer cmd = VK_NULL_HANDLE; + auto r = detail::vk_check(vkAllocateCommandBuffers(device_.device(), &cmd_alloc, &cmd), + __FILE__, __LINE__); + if (!r) return r; - VkCommandBufferBeginInfo begin_info{}; - begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO; - begin_info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT; + VkCommandBufferBeginInfo begin_info{}; + begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO; + begin_info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT; - r = detail::vk_check(vkBeginCommandBuffer(cmd, &begin_info), __FILE__, __LINE__); - if (!r) { vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); return r; } + r = detail::vk_check(vkBeginCommandBuffer(cmd, &begin_info), __FILE__, __LINE__); + if (!r) + { + vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); + return r; + } - VkBufferCopy cp{0, 0, total_bytes}; - vkCmdCopyBuffer(cmd, src.buffer().impl(), staging_ring_->buffer(ri), 1, &cp); + VkBufferCopy cp{0, 0, total_bytes}; + vkCmdCopyBuffer(cmd, src.buffer().impl(), staging_ring_->buffer(ri), 1, &cp); - r = detail::vk_check(vkEndCommandBuffer(cmd), __FILE__, __LINE__); - if (!r) { vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); return r; } + r = detail::vk_check(vkEndCommandBuffer(cmd), __FILE__, __LINE__); + if (!r) + { + vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); + return r; + } + + auto fence = staging_ring_->fence(ri); + vkResetFences(device_.device(), 1, &fence); - auto fence = staging_ring_->fence(ri); - vkResetFences(device_.device(), 1, &fence); + // 跨 submit 数据依赖：源 buffer 可能刚由在飞上传的 copy 写入 + // （from_matrix → to_matrix 无中间 op 即此场景）——等 in-flight + // region 信号量（本下载自用的 region 已在 acquire 中等待并重置， + // 不在 in-flight 集合内）。时间线信号量允许同一 value 被多个 + // submit 等待，故这里与前面的 matmul 帧重复等待也合法。 + const auto sw = collect_staging_waits(); - // 跨 submit 数据依赖：源 buffer 可能刚由在飞上传的 copy 写入 - // （from_matrix → to_matrix 无中间 op 即此场景）——等 in-flight - // region 信号量（本下载自用的 region 已在 acquire 中等待并重置， - // 不在 in-flight 集合内）。时间线信号量允许同一 value 被多个 - // submit 等待，故这里与前面的 matmul 帧重复等待也合法。 - const auto sw = collect_staging_waits(); + { + std::lock_guard lock(queue_mutex_); + VkSubmitInfo submit_info{}; + submit_info.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO; + submit_info.commandBufferCount = 1; + submit_info.pCommandBuffers = &cmd; + VkTimelineSemaphoreSubmitInfo timeline_info{}; + attach_staging_waits(submit_info, sw, timeline_info); + r = detail::vk_check(vkQueueSubmit(device_.compute_queue(), 1, &submit_info, fence), + __FILE__, __LINE__); + if (!r) { - std::lock_guard lock(queue_mutex_); - VkSubmitInfo submit_info{}; - submit_info.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO; - submit_info.commandBufferCount = 1; - submit_info.pCommandBuffers = &cmd; - VkTimelineSemaphoreSubmitInfo timeline_info{}; - attach_staging_waits(submit_info, sw, timeline_info); - - r = detail::vk_check( - vkQueueSubmit(device_.compute_queue(), 1, &submit_info, fence), - __FILE__, __LINE__); - if (!r) { vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); return r; } + vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); + return r; } + } - r = detail::vk_check( - vkWaitForFences(device_.device(), 1, &fence, VK_TRUE, 10'000'000'000ULL), - __FILE__, __LINE__); - if (!r) { vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); return r; } - - r = staging_ring_->download(ri, cpu_data, 0); - + r = detail::vk_check( + vkWaitForFences(device_.device(), 1, &fence, VK_TRUE, 10'000'000'000ULL), __FILE__, + __LINE__); + if (!r) + { vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); return r; } - // ── 大矩阵分块下载 ─────────────────────────────────────────── - const std::size_t chunk_bytes = (staging_cap / sizeof(ElemType)) * sizeof(ElemType); - std::size_t src_offset = 0; + r = staging_ring_->download(ri, cpu_data, 0); - while (src_offset < total_bytes) - { - const std::size_t this_chunk = std::min(chunk_bytes, total_bytes - src_offset); - const std::size_t this_elems = this_chunk / sizeof(ElemType); - const std::size_t dst_elem_offset = src_offset / sizeof(ElemType); + vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); + return r; + } - // 1. 获取 staging region - auto ri = staging_ring_->acquire(); + // ── 大矩阵分块下载 ─────────────────────────────────────────── + const std::size_t chunk_bytes = (staging_cap / sizeof(ElemType)) * sizeof(ElemType); + std::size_t src_offset = 0; - // 2. 录制 copy 命令：GPU → staging - VkCommandBufferAllocateInfo cmd_alloc{}; - cmd_alloc.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO; - cmd_alloc.commandPool = command_pool_; - cmd_alloc.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; - cmd_alloc.commandBufferCount = 1; + while (src_offset < total_bytes) + { + const std::size_t this_chunk = std::min(chunk_bytes, total_bytes - src_offset); + const std::size_t this_elems = this_chunk / sizeof(ElemType); + const std::size_t dst_elem_offset = src_offset / sizeof(ElemType); - VkCommandBuffer cmd = VK_NULL_HANDLE; - auto r = detail::vk_check( - vkAllocateCommandBuffers(device_.device(), &cmd_alloc, &cmd), - __FILE__, __LINE__); - if (!r) return r; + // 1. 获取 staging region + auto ri = staging_ring_->acquire(); - VkCommandBufferBeginInfo begin_info{}; - begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO; - begin_info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT; + // 2. 录制 copy 命令：GPU → staging + VkCommandBufferAllocateInfo cmd_alloc{}; + cmd_alloc.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO; + cmd_alloc.commandPool = command_pool_; + cmd_alloc.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; + cmd_alloc.commandBufferCount = 1; - r = detail::vk_check(vkBeginCommandBuffer(cmd, &begin_info), __FILE__, __LINE__); - if (!r) { vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); return r; } + VkCommandBuffer cmd = VK_NULL_HANDLE; + auto r = detail::vk_check(vkAllocateCommandBuffers(device_.device(), &cmd_alloc, &cmd), + __FILE__, __LINE__); + if (!r) return r; - VkBufferCopy cp{src_offset, 0, this_chunk}; - vkCmdCopyBuffer(cmd, src.buffer().impl(), staging_ring_->buffer(ri), 1, &cp); + VkCommandBufferBeginInfo begin_info{}; + begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO; + begin_info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT; - r = detail::vk_check(vkEndCommandBuffer(cmd), __FILE__, __LINE__); - if (!r) { vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); return r; } + r = detail::vk_check(vkBeginCommandBuffer(cmd, &begin_info), __FILE__, __LINE__); + if (!r) + { + vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); + return r; + } - // 3. 提交并等待（+ 跨 submit 数据依赖：等 in-flight region - // 信号量，源 buffer 可能刚由在飞上传的 copy 写入） - auto fence = staging_ring_->fence(ri); - vkResetFences(device_.device(), 1, &fence); + VkBufferCopy cp{src_offset, 0, this_chunk}; + vkCmdCopyBuffer(cmd, src.buffer().impl(), staging_ring_->buffer(ri), 1, &cp); - const auto sw = collect_staging_waits(); + r = detail::vk_check(vkEndCommandBuffer(cmd), __FILE__, __LINE__); + if (!r) + { + vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); + return r; + } - { - std::lock_guard lock(queue_mutex_); - VkSubmitInfo submit_info{}; - submit_info.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO; - submit_info.commandBufferCount = 1; - submit_info.pCommandBuffers = &cmd; - VkTimelineSemaphoreSubmitInfo timeline_info{}; - attach_staging_waits(submit_info, sw, timeline_info); - - r = detail::vk_check( - vkQueueSubmit(device_.compute_queue(), 1, &submit_info, fence), - __FILE__, __LINE__); - if (!r) { vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); return r; } - } + // 3. 提交并等待（+ 跨 submit 数据依赖：等 in-flight region + // 信号量，源 buffer 可能刚由在飞上传的 copy 写入） + auto fence = staging_ring_->fence(ri); + vkResetFences(device_.device(), 1, &fence); - r = detail::vk_check( - vkWaitForFences(device_.device(), 1, &fence, VK_TRUE, 10'000'000'000ULL), - __FILE__, __LINE__); - if (!r) { vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); return r; } + const auto sw = collect_staging_waits(); - // 4. 从 staging 下载当前块到 CPU - r = staging_ring_->download( - ri, - std::span(cpu_data.data() + dst_elem_offset, this_elems), - 0); - if (!r) { vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); return r; } + { + std::lock_guard lock(queue_mutex_); + VkSubmitInfo submit_info{}; + submit_info.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO; + submit_info.commandBufferCount = 1; + submit_info.pCommandBuffers = &cmd; + VkTimelineSemaphoreSubmitInfo timeline_info{}; + attach_staging_waits(submit_info, sw, timeline_info); - vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); - src_offset += this_chunk; + r = detail::vk_check(vkQueueSubmit(device_.compute_queue(), 1, &submit_info, fence), + __FILE__, __LINE__); + if (!r) + { + vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); + return r; + } } - return {}; - } - - // ── 通用 compute dispatch（matmul 去重核心）────────────────── - // 统一的 descriptor/cmd/barrier/dispatch/submit 样板，供 matmul / - // batched_matmul / rearrange 等所有"N 输入 + 1 输出"的 compute kernel 复用。 - // - // 参数： - // pipeline : 目标 compute pipeline - // inputs : 输入张量（写入 binding 0..n-1） - // output : 输出张量（写入最后一个 binding） - // pc : push constants 字节流（按 pipeline_layout 布局） - // wg_x/y/z : dispatch 工作组数（三维） - [[nodiscard]] Result dispatch_compute( - const VulkanPipeline& pipeline, std::span inputs, - const GpuTensor& output, const std::vector& pc, - std::uint32_t wg_x, std::uint32_t wg_y, std::uint32_t wg_z, - const char* ts_label = "dispatch", // 设备侧计时标号（NN_PROFILE） - // per-key 归因 detail（NN_PROFILE；变体名如 "tiled"/"gemv"/"fine"， - // 进表 5 的 "detail@label@stage" 行；空 = 只进表 4） - std::string_view ts_detail = {}) - { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - const auto t_disp_start = std::chrono::steady_clock::now(); - const std::size_t n = inputs.size() + 1; - std::vector bindings(n); - for (std::size_t i = 0; i < inputs.size(); ++i) - bindings[i] = inputs[i].buffer().impl(); - bindings[n - 1] = output.buffer().impl(); - auto ds_r = alloc_desc_set_for(pipeline.descriptor_layout(), bindings); - NN_TRY_CHECK(ds_r); - VkDescriptorSet desc_set = *ds_r; - - auto cmd_r = acquire_cmd(); - if (!cmd_r) + r = detail::vk_check( + vkWaitForFences(device_.device(), 1, &fence, VK_TRUE, 10'000'000'000ULL), __FILE__, + __LINE__); + if (!r) { - release_desc_set_now(desc_set); - NN_TRY_CHECK(cmd_r); + vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); + return r; } - auto [cmd, owns_cmd] = *cmd_r; - - std::vector in_bufs; - in_bufs.reserve(inputs.size()); - for (const auto& t : inputs) - in_bufs.push_back(t.buffer().impl()); - record_input_barriers(cmd, in_bufs); - - vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline.handle()); - vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipeline.pipeline_layout(), 0, 1, &desc_set, 0, nullptr); - if (!pc.empty()) - vkCmdPushConstants(cmd, pipeline.pipeline_layout(), - VK_SHADER_STAGE_COMPUTE_BIT, 0, - static_cast(pc.size()), pc.data()); - ts_dispatch(cmd, ts_label, wg_x, wg_y, wg_z, ts_detail); - record_output_barrier(cmd, output.buffer().impl()); - - if (owns_cmd) + + // 4. 从 staging 下载当前块到 CPU + r = staging_ring_->download( + ri, std::span(cpu_data.data() + dst_elem_offset, this_elems), 0); + if (!r) { - const bool prof = profile_ops_enabled(); - const auto t_rec = std::chrono::steady_clock::now(); - auto r = submit_and_wait(cmd, desc_set); - if (prof) - std::fprintf(stderr, - "[gpu-profile] record+setup=%lldus（提交/等待分段见上一条 solo_op）\n", - static_cast(std::chrono::duration_cast( - t_rec - t_disp_start).count())); - NN_TRY_CHECK(r); + vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); + return r; } - return {}; + + vkFreeCommandBuffers(device_.device(), command_pool_, 1, &cmd); + src_offset += this_chunk; } - // ── 纯 GPU 矩阵乘法（阻塞等待完成）────────────────────────────── - // batch_mode_ 时：录制到 batch_cmd_，不提交不等待 - // transA: A 存储为 (K,M)，按 A^T 使用 - // transB: B 存储为 (N,K)，按 B^T 使用 - [[nodiscard]] Result matmul_gpu( - const GpuTensor& A, const GpuTensor& B, - uint32_t transA = 0, uint32_t transB = 0, - bool f16_io = false) + return {}; +} + +// ── 通用 compute dispatch（matmul 去重核心）────────────────── +// 统一的 descriptor/cmd/barrier/dispatch/submit 样板，供 matmul / +// batched_matmul / rearrange 等所有"N 输入 + 1 输出"的 compute kernel 复用。 +// +// 参数： +// pipeline : 目标 compute pipeline +// inputs : 输入张量（写入 binding 0..n-1） +// output : 输出张量（写入最后一个 binding） +// pc : push constants 字节流（按 pipeline_layout 布局） +// wg_x/y/z : dispatch 工作组数（三维） +[[nodiscard]] Result +dispatch_compute(const VulkanPipeline& pipeline, std::span inputs, + const GpuTensor& output, const std::vector& pc, std::uint32_t wg_x, + std::uint32_t wg_y, std::uint32_t wg_z, + const char* ts_label = "dispatch", // 设备侧计时标号（NN_PROFILE） + // per-key 归因 detail（NN_PROFILE；变体名如 "tiled"/"gemv"/"fine"， + // 进表 5 的 "detail@label@stage" 行；空 = 只进表 4） + std::string_view ts_detail = {}) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + const auto t_disp_start = std::chrono::steady_clock::now(); + const std::size_t n = inputs.size() + 1; + std::vector bindings(n); + for (std::size_t i = 0; i < inputs.size(); ++i) + bindings[i] = inputs[i].buffer().impl(); + bindings[n - 1] = output.buffer().impl(); + auto ds_r = alloc_desc_set_for(pipeline.descriptor_layout(), bindings); + NN_TRY_CHECK(ds_r); + VkDescriptorSet desc_set = *ds_r; + + auto cmd_r = acquire_cmd(); + if (!cmd_r) { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - // 计算有效维度（考虑转置） - const auto M = static_cast(transA ? A.cols() : A.rows()); - const auto K = static_cast(transA ? A.rows() : A.cols()); - const auto N = static_cast(transB ? B.rows() : B.cols()); - const auto K_B = static_cast(transB ? B.cols() : B.rows()); - if (K != K_B) - NN_FAIL("Dimension mismatch"); - // 1. 分配输出 Tensor（f16_io：按 2B/元素分配 + 纯绑定视图，见 - // batched_matmul_gpu 的同款做法） - const bool use_f16 = f16_io; - // f16 可用性：N≤8 走 GEMV（C2e，仅需 gemv_f16），否则 tiled_f16 - if (use_f16 && N <= 8u && !has(PipelineId::matmul_tiled_f16) && - has(PipelineId::matmul_gemv_f16)) - { - // gemv_f16 可用 → 下方分派走 GEMV，无需 tiled_f16 - } - else if (use_f16 && !has(PipelineId::matmul_tiled_f16)) - NN_FAIL("matmul_gpu: f16 pipeline 不可用（设备无 SSBO 16 位存储？）"); - std::optional owned; - if (use_f16) - { - auto f16r = GpuTensorF16::create_empty(M, N, *this); - NN_TRY_CHECK(f16r); - owned.emplace(f16r->shared_buffer(), f16r->rows(), f16r->cols()); - } - else - { - auto C_res = GpuTensor::create_empty(M, N, *this); - NN_TRY_CHECK(C_res); - owned.emplace(std::move(*C_res)); - } - GpuTensor C = *owned; - - // 2. 选择 Pipeline（优先级：f16 存储版 > 小 N GEMV > 粗化分块 > naive） - // GEMV：N ≤ GEMV_MAX_N（shader ROWS=4 行/WG，grid.x=ceil(M/4)）。 - // tiled 在小 N 下 64×64 块 (64-N)/64 列空转，GEMV 每线程只算 - // 真实 N 列（借鉴 ggml mul_mat_vec 分派）。 - constexpr uint32_t GEMV_MAX_N = 8; - constexpr uint32_t GEMV_ROWS = 4; // 与 matmul_gemv.comp 的 ROWS 一致 - // subgroup≥4 门禁：shader red[..][64] 的容量假设（256/4 = 64 槽）—— - // 更小 subgroup 会溢出 64 槽上限、部分和不落表 → 静默错值 - // （"只有 GPU 错"一类，见 docs/todo/08-pitfalls-and-lessons.md - // §4.10；实测桌面卡恒 ≥8，门禁是保险丝） - // Phase C2e：f16 也可走 GEMV（f16 变体 pipeline 存在时）——小 N - // （n_out≤8）f16 GEMM 直读直写；变体 pipeline 缺失时引擎回退 - // 边界 cast。 - const bool use_gemv = has(PipelineId::matmul_gemv) && N <= GEMV_MAX_N - && device_.subgroup_size() >= 4u - && (!use_f16 || has(PipelineId::matmul_gemv_f16)); - const bool use_tiled = (use_f16 || has(PipelineId::matmul_tiled)) && !use_gemv; - // P0 细 tile（wgrad 网格饥饿）：64-tile 网格只有 - // ceil(N/64)×ceil(M/64) 个工作组——反传 wgrad 形状（M,N=64..256 × - // K=16384）只派 1~4 个，40HX（34 SM）大面积空转（实测 208~278 - // GFLOPS vs 同卡方形 2489）。低于阈值换 16×16 输出块变体：网格细 - // 16×、每输出元素 K 累加序不变（逐位一致）。仅 f32（fine 变体无 - // f16 版），GEMV（N≤8）与 naive 兜底不受影响。 - bool use_fine = false; - bool use_naive_fine = false; - if (use_tiled && !use_f16 && !mm_fine_disabled()) - { - const std::uint64_t wg64 = - static_cast((N + 63u) / 64u) * - ((M + 63u) / 64u); - if (wg64 < nn::EXPR_MATMUL_FINE_MIN_WGS) - { - // NN_MM_FINE_NAIVE=1 → naive（每线程 1 输出：线程数 = M×N， - // 是细 tiled 的 16×；无 shared/barrier，靠 warp 切换吃延迟） - use_naive_fine = mm_fine_naive() && - pipelines_.get(PipelineId::matmul).handle() != VK_NULL_HANDLE; - use_fine = !use_naive_fine && has(PipelineId::matmul_tiled_fine); - } - } - const auto& pipeline = use_gemv ? (use_f16 ? pipelines_.get(PipelineId::matmul_gemv_f16) - : pipelines_.get(PipelineId::matmul_gemv)) - : use_f16 ? pipelines_.get(PipelineId::matmul_tiled_f16) - : use_naive_fine ? pipelines_.get(PipelineId::matmul) - : use_fine ? pipelines_.get(PipelineId::matmul_tiled_fine) - : use_tiled ? pipelines_.get(PipelineId::matmul_tiled) - : pipelines_.get(PipelineId::matmul); - - // 3. 复用通用 dispatch：2 输入 + 1 输出，push {M,N,K,transA,transB} - // Dispatch：按 shader 输出块尺寸计算工作组数 - // matmul.comp（naive）：16×16 线程网格 = 16×16 输出块 - // matmul_tiled.comp：64×64 输出块（BM/BN） - // matmul_gemv.comp：每 WG ROWS=4 行 × N 列（单维 grid.x） - // ⚠ dispatch 必须按各 shader 的实际输出块尺寸计算（naive 16×16 / - // tiled 64×64 / GEMV ROWS=4 行），不可按 WORKGROUP_SIZE=16 统一算： - // 那会让 tiled 版多派 16 倍冗余工作组（每 16×16 一个组而非 64×64）， - // GPU 做 16 倍无效计算，实测 matmul 峰值只剩 ~4%（0.7/15.7 TFLOPS）。 - const uint32_t push_data[5] = {M, N, K, transA, transB}; - std::vector pc(sizeof(push_data)); - std::memcpy(pc.data(), push_data, sizeof(push_data)); - - std::vector inputs{A, B}; - // 设备侧计时 detail：把 matmul 的 pipeline 变体带进 per-key 表 - // （gemv / tiled / tiled_f16 / fine / naive_fine / naive） - const char* mm_variant = use_gemv ? (use_f16 ? "gemv_f16" : "gemv") - : use_naive_fine ? "naive_fine" - : use_fine ? "fine" - : use_tiled ? (use_f16 ? "tiled_f16" : "tiled") - : "naive"; - if (use_gemv) - { - auto r = dispatch_compute(pipeline, inputs, C, pc, - (M + GEMV_ROWS - 1u) / GEMV_ROWS, 1u, 1u, "matmul", mm_variant); - NN_TRY_CHECK(r); - return C; - } - const uint32_t tile = (use_fine || use_naive_fine) ? 16u - : use_tiled ? 64u : 16u; - auto r = dispatch_compute(pipeline, inputs, C, pc, - (N + tile - 1u) / tile, (M + tile - 1u) / tile, 1u, - "matmul", mm_variant); + release_desc_set_now(desc_set); + NN_TRY_CHECK(cmd_r); + } + auto [cmd, owns_cmd] = *cmd_r; + + std::vector in_bufs; + in_bufs.reserve(inputs.size()); + for (const auto& t : inputs) + in_bufs.push_back(t.buffer().impl()); + record_input_barriers(cmd, in_bufs); + + vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline.handle()); + vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline.pipeline_layout(), 0, 1, + &desc_set, 0, nullptr); + if (!pc.empty()) + vkCmdPushConstants(cmd, pipeline.pipeline_layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, + static_cast(pc.size()), pc.data()); + ts_dispatch(cmd, ts_label, wg_x, wg_y, wg_z, ts_detail); + record_output_barrier(cmd, output.buffer().impl()); + + if (owns_cmd) + { + const bool prof = profile_ops_enabled(); + const auto t_rec = std::chrono::steady_clock::now(); + auto r = submit_and_wait(cmd, desc_set); + if (prof) + std::fprintf( + stderr, "[gpu-profile] record+setup=%lldus（提交/等待分段见上一条 solo_op）\n", + static_cast( + std::chrono::duration_cast(t_rec - t_disp_start) + .count())); NN_TRY_CHECK(r); - return C; } - - // ── 纯 GPU 批量矩阵乘法 ────────────────────────────────────────── - // 对每个 batch b 计算 C_b = alpha * op(A_b, B_b)，结果垂直堆叠为 (batch*M, N) - // A: (batch * A_rows_per_batch, A_cols)，B: (batch * B_rows_per_batch, B_cols) - // batch 步长由 shader 从 M*K / K*N / M*N 推导，无需额外传入 - // alpha: 输出缩放系数（cuBLAS sgemm 语义），在 shader 写出时一次完成 - // - // f16_io（Phase 2 in-kernel f16）：A/B 是 **f16 缓冲**、C 也按 f16 分配 - // （2B/元素）→ 用 f16 存储版 pipeline（f16 载入 + f32 累加 + f16 写出）， - // 彻底消除"每个操作数抬 f32 的整份副本"。此时 A/B 只作**缓冲视图**使用 - // （GpuTensorT 包裹 f16 buffer，字节布局由 pipeline 决定 —— 与 - // run_fused_gpu 的 out_f16 同一套做法）。 - [[nodiscard]] Result batched_matmul_gpu( - const GpuTensor& A, const GpuTensor& B, - uint32_t batch, - uint32_t transA = 0, uint32_t transB = 0, - float alpha = 1.0f, - bool f16_io = false) + return {}; +} + +// ── 纯 GPU 矩阵乘法（阻塞等待完成）────────────────────────────── +// batch_mode_ 时：录制到 batch_cmd_，不提交不等待 +// transA: A 存储为 (K,M)，按 A^T 使用 +// transB: B 存储为 (N,K)，按 B^T 使用 +[[nodiscard]] Result matmul_gpu(const GpuTensor& A, const GpuTensor& B, + uint32_t transA = 0, uint32_t transB = 0, + bool f16_io = false) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + // 计算有效维度（考虑转置） + const auto M = static_cast(transA ? A.cols() : A.rows()); + const auto K = static_cast(transA ? A.rows() : A.cols()); + const auto N = static_cast(transB ? B.rows() : B.cols()); + const auto K_B = static_cast(transB ? B.cols() : B.rows()); + if (K != K_B) NN_FAIL("Dimension mismatch"); + // 1. 分配输出 Tensor（f16_io：按 2B/元素分配 + 纯绑定视图，见 + // batched_matmul_gpu 的同款做法） + const bool use_f16 = f16_io; + // f16 可用性：N≤8 走 GEMV（C2e，仅需 gemv_f16），否则 tiled_f16 + if (use_f16 && N <= 8u && !has(PipelineId::matmul_tiled_f16) && + has(PipelineId::matmul_gemv_f16)) { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - if (batch == 0) - NN_FAIL("batched_matmul_gpu: batch must be > 0"); - const bool use_f16 = f16_io && has(PipelineId::batched_matmul_f16); - if (f16_io && !use_f16) - NN_FAIL("batched_matmul_gpu: f16 pipeline 不可用（设备无 SSBO 16 位存储？）"); - if (!use_f16 && !has(PipelineId::batched_matmul)) - NN_FAIL("batched_matmul_gpu: pipeline not available"); - // 校验：A.rows() 必须能被 batch 整除 - if (A.rows() % batch != 0 || B.rows() % batch != 0) - NN_FAIL("batched_matmul_gpu: rows not divisible by batch"); - // 计算每个 batch 的逻辑维度 - const auto a_rows_per = static_cast(A.rows() / batch); - const auto b_rows_per = static_cast(B.rows() / batch); - const auto M = transA ? static_cast(A.cols()) : a_rows_per; - const auto K = transA ? a_rows_per : static_cast(A.cols()); - const auto K_B = transB ? static_cast(B.cols()) : b_rows_per; - const auto N = transB ? b_rows_per : static_cast(B.cols()); - if (K != K_B) - NN_FAIL("batched_matmul_gpu: K dimension mismatch"); - // 1. 分配输出 Tensor: (batch * M, N) - // f16：按 2B/元素分配，再用 GpuTensor(shared_buffer,...) 包一层纯 - // 绑定视图交给 dispatch（调用方按 f16_io 重贴 GpuTensorF16）。 - std::optional owned; - if (use_f16) - { - auto f16r = GpuTensorF16::create_empty( - static_cast(batch) * M, N, *this); - NN_TRY_CHECK(f16r); - owned.emplace(f16r->shared_buffer(), f16r->rows(), f16r->cols()); - } - else + // gemv_f16 可用 → 下方分派走 GEMV，无需 tiled_f16 + } + else if (use_f16 && !has(PipelineId::matmul_tiled_f16)) + NN_FAIL("matmul_gpu: f16 pipeline 不可用（设备无 SSBO 16 位存储？）"); + std::optional owned; + if (use_f16) + { + auto f16r = GpuTensorF16::create_empty(M, N, *this); + NN_TRY_CHECK(f16r); + owned.emplace(f16r->shared_buffer(), f16r->rows(), f16r->cols()); + } + else + { + auto C_res = GpuTensor::create_empty(M, N, *this); + NN_TRY_CHECK(C_res); + owned.emplace(std::move(*C_res)); + } + GpuTensor C = *owned; + + // 2. 选择 Pipeline（优先级：f16 存储版 > 小 N GEMV > 粗化分块 > naive） + // GEMV：N ≤ GEMV_MAX_N（shader ROWS=4 行/WG，grid.x=ceil(M/4)）。 + // tiled 在小 N 下 64×64 块 (64-N)/64 列空转，GEMV 每线程只算 + // 真实 N 列（借鉴 ggml mul_mat_vec 分派）。 + constexpr uint32_t GEMV_MAX_N = 8; + constexpr uint32_t GEMV_ROWS = 4; // 与 matmul_gemv.comp 的 ROWS 一致 + // subgroup≥4 门禁：shader red[..][64] 的容量假设（256/4 = 64 槽）—— + // 更小 subgroup 会溢出 64 槽上限、部分和不落表 → 静默错值 + // （"只有 GPU 错"一类，见 docs/todo/08-pitfalls-and-lessons.md + // §4.10；实测桌面卡恒 ≥8，门禁是保险丝） + // Phase C2e：f16 也可走 GEMV（f16 变体 pipeline 存在时）——小 N + // （n_out≤8）f16 GEMM 直读直写；变体 pipeline 缺失时引擎回退 + // 边界 cast。 + const bool use_gemv = has(PipelineId::matmul_gemv) && N <= GEMV_MAX_N && + device_.subgroup_size() >= 4u && + (!use_f16 || has(PipelineId::matmul_gemv_f16)); + const bool use_tiled = (use_f16 || has(PipelineId::matmul_tiled)) && !use_gemv; + // P0 细 tile（wgrad 网格饥饿）：64-tile 网格只有 + // ceil(N/64)×ceil(M/64) 个工作组——反传 wgrad 形状（M,N=64..256 × + // K=16384）只派 1~4 个，40HX（34 SM）大面积空转（实测 208~278 + // GFLOPS vs 同卡方形 2489）。低于阈值换 16×16 输出块变体：网格细 + // 16×、每输出元素 K 累加序不变（逐位一致）。仅 f32（fine 变体无 + // f16 版），GEMV（N≤8）与 naive 兜底不受影响。 + bool use_fine = false; + bool use_naive_fine = false; + if (use_tiled && !use_f16 && !mm_fine_disabled()) + { + const std::uint64_t wg64 = static_cast((N + 63u) / 64u) * ((M + 63u) / 64u); + if (wg64 < nn::EXPR_MATMUL_FINE_MIN_WGS) { - auto C_res = GpuTensor::create_empty(static_cast(batch) * M, N, *this); - NN_TRY_CHECK(C_res); - owned.emplace(std::move(*C_res)); + // NN_MM_FINE_NAIVE=1 → naive（每线程 1 输出：线程数 = M×N， + // 是细 tiled 的 16×；无 shared/barrier，靠 warp 切换吃延迟） + use_naive_fine = + mm_fine_naive() && pipelines_.get(PipelineId::matmul).handle() != VK_NULL_HANDLE; + use_fine = !use_naive_fine && has(PipelineId::matmul_tiled_fine); } - GpuTensor C = *owned; - - // 2. 复用通用 dispatch：2 输入 + 1 输出 - // Push constants: M, N, K, transA, transB, alpha（24B，含输出缩放系数） - // shader 采用 64×64 寄存器分块（与 matmul_tiled 相同）， - // 每个 WorkGroup 计算 64×64 输出块，Z 维度索引 batch - struct PushDataBmm { - uint32_t M, N, K, transA, transB; - float alpha; - } push{M, N, K, transA, transB, alpha}; - std::vector pc(sizeof(push)); - std::memcpy(pc.data(), &push, sizeof(push)); - - // Dispatch: X/Y 覆盖 64×64 输出块，Z 维度 = batch - constexpr uint32_t BM = 64, BN = 64; - std::vector inputs{A, B}; - auto r = dispatch_compute( - use_f16 ? pipelines_.get(PipelineId::batched_matmul_f16) : pipelines_.get(PipelineId::batched_matmul), - inputs, C, pc, (N + BN - 1) / BN, (M + BM - 1) / BM, batch, - "batched_matmul", use_f16 ? "f16" : "f32"); + } + const auto& pipeline = use_gemv ? (use_f16 ? pipelines_.get(PipelineId::matmul_gemv_f16) + : pipelines_.get(PipelineId::matmul_gemv)) + : use_f16 ? pipelines_.get(PipelineId::matmul_tiled_f16) + : use_naive_fine ? pipelines_.get(PipelineId::matmul) + : use_fine ? pipelines_.get(PipelineId::matmul_tiled_fine) + : use_tiled ? pipelines_.get(PipelineId::matmul_tiled) + : pipelines_.get(PipelineId::matmul); + + // 3. 复用通用 dispatch：2 输入 + 1 输出，push {M,N,K,transA,transB} + // Dispatch：按 shader 输出块尺寸计算工作组数 + // matmul.comp（naive）：16×16 线程网格 = 16×16 输出块 + // matmul_tiled.comp：64×64 输出块（BM/BN） + // matmul_gemv.comp：每 WG ROWS=4 行 × N 列（单维 grid.x） + // ⚠ dispatch 必须按各 shader 的实际输出块尺寸计算（naive 16×16 / + // tiled 64×64 / GEMV ROWS=4 行），不可按 WORKGROUP_SIZE=16 统一算： + // 那会让 tiled 版多派 16 倍冗余工作组（每 16×16 一个组而非 64×64）， + // GPU 做 16 倍无效计算，实测 matmul 峰值只剩 ~4%（0.7/15.7 TFLOPS）。 + const uint32_t push_data[5] = {M, N, K, transA, transB}; + std::vector pc(sizeof(push_data)); + std::memcpy(pc.data(), push_data, sizeof(push_data)); + + std::vector inputs{A, B}; + // 设备侧计时 detail：把 matmul 的 pipeline 变体带进 per-key 表 + // （gemv / tiled / tiled_f16 / fine / naive_fine / naive） + const char* mm_variant = use_gemv ? (use_f16 ? "gemv_f16" : "gemv") + : use_naive_fine ? "naive_fine" + : use_fine ? "fine" + : use_tiled ? (use_f16 ? "tiled_f16" : "tiled") + : "naive"; + if (use_gemv) + { + auto r = dispatch_compute(pipeline, inputs, C, pc, (M + GEMV_ROWS - 1u) / GEMV_ROWS, 1u, 1u, + "matmul", mm_variant); NN_TRY_CHECK(r); return C; } + const uint32_t tile = (use_fine || use_naive_fine) ? 16u : use_tiled ? 64u : 16u; + auto r = dispatch_compute(pipeline, inputs, C, pc, (N + tile - 1u) / tile, + (M + tile - 1u) / tile, 1u, "matmul", mm_variant); + NN_TRY_CHECK(r); + return C; +} + +// ── 纯 GPU 批量矩阵乘法 ────────────────────────────────────────── +// 对每个 batch b 计算 C_b = alpha * op(A_b, B_b)，结果垂直堆叠为 (batch*M, N) +// A: (batch * A_rows_per_batch, A_cols)，B: (batch * B_rows_per_batch, B_cols) +// batch 步长由 shader 从 M*K / K*N / M*N 推导，无需额外传入 +// alpha: 输出缩放系数（cuBLAS sgemm 语义），在 shader 写出时一次完成 +// +// f16_io（Phase 2 in-kernel f16）：A/B 是 **f16 缓冲**、C 也按 f16 分配 +// （2B/元素）→ 用 f16 存储版 pipeline（f16 载入 + f32 累加 + f16 写出）， +// 彻底消除"每个操作数抬 f32 的整份副本"。此时 A/B 只作**缓冲视图**使用 +// （GpuTensorT 包裹 f16 buffer，字节布局由 pipeline 决定 —— 与 +// run_fused_gpu 的 out_f16 同一套做法）。 +[[nodiscard]] Result batched_matmul_gpu(const GpuTensor& A, const GpuTensor& B, + uint32_t batch, uint32_t transA = 0, + uint32_t transB = 0, float alpha = 1.0f, + bool f16_io = false) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + if (batch == 0) NN_FAIL("batched_matmul_gpu: batch must be > 0"); + const bool use_f16 = f16_io && has(PipelineId::batched_matmul_f16); + if (f16_io && !use_f16) + NN_FAIL("batched_matmul_gpu: f16 pipeline 不可用（设备无 SSBO 16 位存储？）"); + if (!use_f16 && !has(PipelineId::batched_matmul)) + NN_FAIL("batched_matmul_gpu: pipeline not available"); + // 校验：A.rows() 必须能被 batch 整除 + if (A.rows() % batch != 0 || B.rows() % batch != 0) + NN_FAIL("batched_matmul_gpu: rows not divisible by batch"); + // 计算每个 batch 的逻辑维度 + const auto a_rows_per = static_cast(A.rows() / batch); + const auto b_rows_per = static_cast(B.rows() / batch); + const auto M = transA ? static_cast(A.cols()) : a_rows_per; + const auto K = transA ? a_rows_per : static_cast(A.cols()); + const auto K_B = transB ? static_cast(B.cols()) : b_rows_per; + const auto N = transB ? b_rows_per : static_cast(B.cols()); + if (K != K_B) NN_FAIL("batched_matmul_gpu: K dimension mismatch"); + // 1. 分配输出 Tensor: (batch * M, N) + // f16：按 2B/元素分配，再用 GpuTensor(shared_buffer,...) 包一层纯 + // 绑定视图交给 dispatch（调用方按 f16_io 重贴 GpuTensorF16）。 + std::optional owned; + if (use_f16) + { + auto f16r = GpuTensorF16::create_empty(static_cast(batch) * M, N, *this); + NN_TRY_CHECK(f16r); + owned.emplace(f16r->shared_buffer(), f16r->rows(), f16r->cols()); + } + else + { + auto C_res = GpuTensor::create_empty(static_cast(batch) * M, N, *this); + NN_TRY_CHECK(C_res); + owned.emplace(std::move(*C_res)); + } + GpuTensor C = *owned; + // 2. 复用通用 dispatch：2 输入 + 1 输出 + // Push constants: M, N, K, transA, transB, alpha（24B，含输出缩放系数） + // shader 采用 64×64 寄存器分块（与 matmul_tiled 相同）， + // 每个 WorkGroup 计算 64×64 输出块，Z 维度索引 batch + struct PushDataBmm + { + uint32_t M, N, K, transA, transB; + float alpha; + } push{M, N, K, transA, transB, alpha}; + std::vector pc(sizeof(push)); + std::memcpy(pc.data(), &push, sizeof(push)); + + // Dispatch: X/Y 覆盖 64×64 输出块，Z 维度 = batch + constexpr uint32_t BM = 64, BN = 64; + std::vector inputs{A, B}; + auto r = dispatch_compute(use_f16 ? pipelines_.get(PipelineId::batched_matmul_f16) + : pipelines_.get(PipelineId::batched_matmul), + inputs, C, pc, (N + BN - 1) / BN, (M + BM - 1) / BM, batch, + "batched_matmul", use_f16 ? "f16" : "f32"); + NN_TRY_CHECK(r); + return C; +} diff --git a/include/neuralnet.cpp/backend/compute_vk_backend_ops.inc b/include/neuralnet.cpp/backend/compute_vk_backend_ops.inc index 46970f3a..74f4c53b 100644 --- a/include/neuralnet.cpp/backend/compute_vk_backend_ops.inc +++ b/include/neuralnet.cpp/backend/compute_vk_backend_ops.inc @@ -4,2408 +4,2226 @@ // ⚠ 本文件不是独立头：仅由 backend/compute_vk_backend.hpp 在 // `class GpuBackend { ... }` 体内**按序** #include（顺序即成员序）。 // ═══════════════════════════════════════════════════════════════════════════ - // ══════════════════════════════════════════════════════════════════ - // elementwise_v2_gpu — 逐元素原语（unary/binary/select/axpy） - // - // Push Constants (32 bytes): - // count, mode, op, cmp_op, flags, scalar_b, scalar_then, scalar_else - // Bindings: A(0), B(1), C(2), OUT(3) - // - // 对未使用的 binding（如 unary 模式的 B/C），绑定 A 的 buffer（无害占位）。 - // - // out != nullptr 时为原地模式：OUT 直接绑定 out 的 buffer（通常即 A）， - // 免分配 + 免全量写出。逐元素 kernel 每线程只读写自己下标一次， - // read-before-write 天然成立，别名安全。 - // ══════════════════════════════════════════════════════════════════ - [[nodiscard]] Result elementwise_v2_gpu( - const GpuTensor& A, const GpuTensor* B, const GpuTensor* C, - uint32_t count, uint32_t mode, uint32_t op, uint32_t cmp_op, - uint32_t flags, float scalar_b, float scalar_then, float scalar_else, - const GpuTensor* out = nullptr, bool f16_io = false) - { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - if (!has(PipelineId::elementwise_v2)) - NN_FAIL("elementwise_v2 pipeline not available"); - // Phase C1b f16 变体（仅需 16 位存储门控）；请求但无 pipeline → - // 报错由引擎层回退边界 cast - if (f16_io && !has(PipelineId::elementwise_v2_f16)) - NN_FAIL("elementwise_v2_gpu: f16 pipeline 不可用（设备无16位存储）"); - const auto& pipeline = f16_io ? pipelines_.get(PipelineId::elementwise_v2_f16) - : pipelines_.get(PipelineId::elementwise_v2); - - // 1. 输出 Tensor：原地模式直接复用 out 的 buffer，否则新分配 - // （f16_io：按 2B/元素分配 + 纯绑定视图，同 matmul_gpu） - GpuTensor output; - if (out) - { - output = GpuTensor(*out); - } - else if (f16_io) - { - auto f16r = GpuTensorF16::create_empty(A.rows(), A.cols(), *this); - NN_TRY_CHECK(f16r); - output = GpuTensor(f16r->shared_buffer(), f16r->rows(), f16r->cols()); - } - else - { - auto output_res = GpuTensor::create_empty(A.rows(), A.cols(), *this); - NN_TRY_CHECK(output_res); - output = std::move(*output_res); - } - - // 2. 分配并写入描述符集（未使用的 binding 绑定 A 作为占位） - VkBuffer a_buf = A.buffer().impl(); - const VkBuffer bindings[4]{ - a_buf, - B ? B->buffer().impl() : a_buf, - C ? C->buffer().impl() : a_buf, - output.buffer().impl(), - }; - auto ds_r = alloc_desc_set_for(pipeline.descriptor_layout(), bindings); - NN_TRY_CHECK(ds_r); - VkDescriptorSet desc_set = *ds_r; +// ══════════════════════════════════════════════════════════════════ +// elementwise_v2_gpu — 逐元素原语（unary/binary/select/axpy） +// +// Push Constants (32 bytes): +// count, mode, op, cmp_op, flags, scalar_b, scalar_then, scalar_else +// Bindings: A(0), B(1), C(2), OUT(3) +// +// 对未使用的 binding（如 unary 模式的 B/C），绑定 A 的 buffer（无害占位）。 +// +// out != nullptr 时为原地模式：OUT 直接绑定 out 的 buffer（通常即 A）， +// 免分配 + 免全量写出。逐元素 kernel 每线程只读写自己下标一次， +// read-before-write 天然成立，别名安全。 +// ══════════════════════════════════════════════════════════════════ +[[nodiscard]] Result +elementwise_v2_gpu(const GpuTensor& A, const GpuTensor* B, const GpuTensor* C, uint32_t count, + uint32_t mode, uint32_t op, uint32_t cmp_op, uint32_t flags, float scalar_b, + float scalar_then, float scalar_else, const GpuTensor* out = nullptr, + bool f16_io = false) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + if (!has(PipelineId::elementwise_v2)) NN_FAIL("elementwise_v2 pipeline not available"); + // Phase C1b f16 变体（仅需 16 位存储门控）；请求但无 pipeline → + // 报错由引擎层回退边界 cast + if (f16_io && !has(PipelineId::elementwise_v2_f16)) + NN_FAIL("elementwise_v2_gpu: f16 pipeline 不可用（设备无16位存储）"); + const auto& pipeline = f16_io ? pipelines_.get(PipelineId::elementwise_v2_f16) + : pipelines_.get(PipelineId::elementwise_v2); + + // 1. 输出 Tensor：原地模式直接复用 out 的 buffer，否则新分配 + // （f16_io：按 2B/元素分配 + 纯绑定视图，同 matmul_gpu） + GpuTensor output; + if (out) { output = GpuTensor(*out); } + else if (f16_io) + { + auto f16r = GpuTensorF16::create_empty(A.rows(), A.cols(), *this); + NN_TRY_CHECK(f16r); + output = GpuTensor(f16r->shared_buffer(), f16r->rows(), f16r->cols()); + } + else + { + auto output_res = GpuTensor::create_empty(A.rows(), A.cols(), *this); + NN_TRY_CHECK(output_res); + output = std::move(*output_res); + } - // 3. 获取 command buffer - auto cmd_r = acquire_cmd(); - if (!cmd_r) - { - release_desc_set_now(desc_set); - NN_TRY_CHECK(cmd_r); - } - auto [cmd, owns_cmd] = *cmd_r; + // 2. 分配并写入描述符集（未使用的 binding 绑定 A 作为占位） + VkBuffer a_buf = A.buffer().impl(); + const VkBuffer bindings[4]{ + a_buf, + B ? B->buffer().impl() : a_buf, + C ? C->buffer().impl() : a_buf, + output.buffer().impl(), + }; + auto ds_r = alloc_desc_set_for(pipeline.descriptor_layout(), bindings); + NN_TRY_CHECK(ds_r); + VkDescriptorSet desc_set = *ds_r; + + // 3. 获取 command buffer + auto cmd_r = acquire_cmd(); + if (!cmd_r) + { + release_desc_set_now(desc_set); + NN_TRY_CHECK(cmd_r); + } + auto [cmd, owns_cmd] = *cmd_r; - // 4. 录制 - record_input_barriers(cmd, {a_buf}); - vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipeline.handle()); - vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipeline.pipeline_layout(), 0, 1, &desc_set, 0, nullptr); - - // Push constants: count, mode, op, cmp_op, flags (5×uint32) + scalar_b, scalar_then, scalar_else (3×float) - struct PushData { - uint32_t count, mode, op, cmp_op, flags; - float scalar_b, scalar_then, scalar_else; - } push{count, mode, op, cmp_op, flags, scalar_b, scalar_then, scalar_else}; - vkCmdPushConstants(cmd, pipeline.pipeline_layout(), - VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(push), &push); - - // vec4 kernel（与 DSL count/(256*vec_width) 同口径）：每线程 4 元素 - const uint32_t wg_count = (count + 256u * 4u - 1u) / (256u * 4u); - ts_dispatch(cmd, "elementwise_v2", wg_count, 1, 1); - record_output_barrier(cmd, output.buffer().impl()); + // 4. 录制 + record_input_barriers(cmd, {a_buf}); + vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline.handle()); + vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline.pipeline_layout(), 0, 1, + &desc_set, 0, nullptr); - // 6. 独立模式提交 - if (owns_cmd) - { - auto r = submit_and_wait(cmd, desc_set); - NN_TRY_CHECK(r); - } - return output; + // Push constants: count, mode, op, cmp_op, flags (5×uint32) + scalar_b, scalar_then, scalar_else (3×float) + struct PushData + { + uint32_t count, mode, op, cmp_op, flags; + float scalar_b, scalar_then, scalar_else; + } push{count, mode, op, cmp_op, flags, scalar_b, scalar_then, scalar_else}; + vkCmdPushConstants(cmd, pipeline.pipeline_layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, + sizeof(push), &push); + + // vec4 kernel（与 DSL count/(256*vec_width) 同口径）：每线程 4 元素 + const uint32_t wg_count = (count + 256u * 4u - 1u) / (256u * 4u); + ts_dispatch(cmd, "elementwise_v2", wg_count, 1, 1); + record_output_barrier(cmd, output.buffer().impl()); + + // 6. 独立模式提交 + if (owns_cmd) + { + auto r = submit_and_wait(cmd, desc_set); + NN_TRY_CHECK(r); } - - // ══════════════════════════════════════════════════════════════════ - // run_fused_gpu — AOT 融合 shader 通用执行入口 - // - // 与 glsl_gen.hpp 生成的 shader 布局严格对应： - // Bindings: 输入 0..N-1，输出 N（全部 storage buffer） - // Push Constants（逐元素）: uint count, uint cols, float c0..（常量池） - // Push Constants（归约）: uint count, uint cols, uint rows, uint vector_out, ... - // Push Constants（matmul 融合）: uint count, uint cols, uint rows, uint mm_k, ... - // local_size_x = 256 - // - // shader_name 必须是已注册的融合 shader 的 key（= expr_spec_key，见 - // 构建期 fused_registry.hpp）。 - // 所有输入同形状 (rows, cols)（matmul 段除外：A/B 按 matmul 形状解释， - // 但 Vulkan buffer 无形状概念，绑定即用）；同一 buffer 可绑定到多个输入。 - // matmul_k：前置 matmul 段的求和维度（形状参数，运行时填充 mm_k 槽； - // 非 matmul shader 传 nullopt）。 - // matmul_batch：matmul 段批量数（形状参数，填充 mm_batch 槽 + dispatch z）。 - // ══════════════════════════════════════════════════════════════════ - [[nodiscard]] Result run_fused_gpu( - const std::string& shader_name, - // 输入按**底层 buffer** 绑定（元素类型由 shader 声明决定）→ 同一结构可在 - // f32 / f16 两种存储上运行（Phase 2 in-kernel f16 的带类型变体）。 - std::span inputs, - std::span consts, - std::size_t rows, std::size_t cols, - bool vector_out = false, - std::span view_params = {}, - std::span rparams = {}, - GpuTensor* output_override = nullptr, - std::optional matmul_k = std::nullopt, - std::uint32_t matmul_batch = 1, - // matmul 段转置布局（bit0=transA，bit1=transB）：**运行期 operand layout， - // 不进 expr_spec_key** —— 同一结构覆盖 4 种转置组合。仅 matmul 形态 - // （matmul / matmul+归约）使用；fold 形态的 mm 段转置仍在 key 里 - // （生成期定死），此处忽略。 - std::uint32_t matmul_trans = 0u, - // fold 收缩轴长度（fold shader 传入填 fold_k 槽；非 fold 传 nullopt） - std::optional fold_k = std::nullopt, - // 输出存储精度（Phase 2 in-kernel f16）：true → 分配 f16 字节布局的 - // 输出缓冲（返回的 GpuTensor 仅是**占位标签**，调用方按此标志重贴 - // GpuTensorF16 —— 见 GpuEngine::eval_expr）。 - bool out_f16 = false, - // 运行期精度分派（键 "#x"）：真实 (输入精度位图, 输出精度)。仅在 - // 分派变体上写入 PC `prec` 槽；其余变体忽略（其精度已是身份）。 - nn::ExprPrecSig prec_sig = 0, - // P2 多输出（extras，IR-C 写穿物化）：主输出之外的**额外输出存储** - // （槽 1..k，与生成器的 `bex` 绑定同序）。仅 V0（全 f32）逐元素 - // shader 支持（`#x`/typed/native16 变体不生成 extras，见 expr_glsl_gen）。 - std::span extra_outs = {}) - { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - const auto prof_t0 = std::chrono::steady_clock::now(); - const auto it = fused_kernels_.find(shader_name); - if (it == fused_kernels_.end()) - NN_FAIL("fused shader not registered: " + shader_name); - const FusedKernel& kernel = it->second; - // （pipeline 绑定推迟到细 tile 选择之后——基础键必须存在，细变体可选） - // 运行期精度分派变体（键后缀 "#x"）：描述符 = 2N 输入视图 + 2 输出视图， - // 且 PC 固定头末尾多一个 `prec` 槽。 - const bool is_disp = - std::string_view(shader_name).ends_with(nn::EXPR_PREC_DISPATCH_SUFFIX); - // PC 固定头槽数：**注册时算定，此处只读**（唯一来源 = - // FusedKernel::pc_base_fixed；本地重推会让两侧漂移 → - // push constant 超 range 被驱动丢弃 → 静默读未定义残留）。 - const std::uint32_t pc_base_fixed = kernel.pc_base_fixed; - // 归约轴：-1=逐元素, 0=行归约, 1=列归约 - const int raxis = kernel.reduce_axis; - // 是否含前置 matmul 段（push constants 多 rows + mm_k） - const bool has_mm = kernel.has_matmul; - // fold 形态（注册时从 FusedShader.spec.fold 取，与生成器同源） - const FoldMeta fmeta = kernel.fold; - const bool is_fold = kernel.is_fold(); - // 运行时视图参数个数（RowMod/RotateHalf 的 vp 槽） - const std::uint32_t n_vp = kernel.view_param_count; - // 运行时标量参数个数（优化器 lr/eps/β 的 rp 槽） - const std::uint32_t n_rp = kernel.rparam_count; - - // ── P0 细 tile 选择（wgrad 网格饥饿）──────────────────────────────── - // 纯 matmul 形态（与下方 dispatch 分支同一判据）且 64-tile 网格的 - // 工作组数（含 batch 维）低于阈值 → 派发同结构的 16×16 输出块变体 - // （键 = shader_name + "#f"，注册表元数据/PC 布局/绑定数与 V0 完全 - // 同源，只是网格粒度不同）。`#f` 未登记（f16 精度变体键、设备未注册） - // 或 NN_MM_NO_FINE=1 → 留在 64-tile 路径（降级，不是错误）。 - // 两变体逐输出元素的 K 累加序一致（BK=16 逐块升序 + 同一 vec4 点积 - // 序）→ 结果逐位一致，字节锚不移动。 - bool use_fine = false; - // ── K 切分选择（wgrad 二次归约，键 #ks+#kr）──────────────────────── - // 细 tile 网格（m_per/16 × cols/16 × batch 个 WG）仍低于 - // EXPR_MATMUL_KSPLIT_MIN_WGS → 把 K 切成 S 段：#ks 按列块序号取段 - // [k_lo, k_lo+span)（dispatch x × S）写 (S·rows, cols) 部分和， - // #kr 跨段升序求和 + 跑尾链落最终输出。S = clamp(ceil(阈值/网格), - // 2, MAX_S)（NN_MM_KSPLIT_S 可强制，A/B 调参）；span 上取 16 对齐 - // （生成器 vec4 快路径的段界 k 连续判据依赖它）。 - // 门槛：V0（#ks/#kr 只注册基础键）、纯 matmul、无 extras、f32、 - // #ks+#kr 同时注册（配对）、NN_MM_NO_FINE 与 NN_MM_NO_KSPLIT 均未关。 - // ⚠ 跨段 K 累加序 ≠ 逐 k 串联 → **移动 GPU 字节锚**（已裁定放行， - // 与 P2-2 同处置：重测 + 更新 §12 + history 记新旧值）。 - bool use_ks = false; - std::uint32_t ks_s = 0; // 段数 S - std::uint32_t ks_span = 0; // 段宽（16 对齐） - if (has_mm && raxis < 0 && !is_fold && !mm_fine_disabled()) - { - const std::uint32_t rows_u = static_cast(rows); - const std::uint32_t m_per = rows_u / matmul_batch; // rows 按 batch 整除（契约） - constexpr std::uint32_t B64 = nn::EXPR_MATMUL_BLOCK; - const std::uint64_t wg64 = - static_cast( - (static_cast(cols) + B64 - 1u) / B64) * - ((m_per + B64 - 1u) / B64) * matmul_batch; - if (wg64 < nn::EXPR_MATMUL_FINE_MIN_WGS && - fused_kernels_.count(shader_name + nn::EXPR_MM_FINE_SUFFIX) != 0) - use_fine = true; - if (wg64 < nn::EXPR_MATMUL_FINE_MIN_WGS && !is_disp && - !mm_ksplit_disabled() && extra_outs.empty() && - matmul_k && *matmul_k > 0 && - fused_kernels_.count( - shader_name + nn::EXPR_MM_KSPLIT_SUFFIX) != 0 && - fused_kernels_.count( - shader_name + nn::EXPR_MM_KSPLIT_REDUCE_SUFFIX) != 0) + return output; +} + +// ══════════════════════════════════════════════════════════════════ +// run_fused_gpu — AOT 融合 shader 通用执行入口 +// +// 与 glsl_gen.hpp 生成的 shader 布局严格对应： +// Bindings: 输入 0..N-1，输出 N（全部 storage buffer） +// Push Constants（逐元素）: uint count, uint cols, float c0..（常量池） +// Push Constants（归约）: uint count, uint cols, uint rows, uint vector_out, ... +// Push Constants（matmul 融合）: uint count, uint cols, uint rows, uint mm_k, ... +// local_size_x = 256 +// +// shader_name 必须是已注册的融合 shader 的 key（= expr_spec_key，见 +// 构建期 fused_registry.hpp）。 +// 所有输入同形状 (rows, cols)（matmul 段除外：A/B 按 matmul 形状解释， +// 但 Vulkan buffer 无形状概念，绑定即用）；同一 buffer 可绑定到多个输入。 +// matmul_k：前置 matmul 段的求和维度（形状参数，运行时填充 mm_k 槽； +// 非 matmul shader 传 nullopt）。 +// matmul_batch：matmul 段批量数（形状参数，填充 mm_batch 槽 + dispatch z）。 +// ══════════════════════════════════════════════════════════════════ +[[nodiscard]] Result +run_fused_gpu(const std::string& shader_name, + // 输入按**底层 buffer** 绑定（元素类型由 shader 声明决定）→ 同一结构可在 + // f32 / f16 两种存储上运行（Phase 2 in-kernel f16 的带类型变体）。 + std::span inputs, std::span consts, + std::size_t rows, std::size_t cols, bool vector_out = false, + std::span view_params = {}, std::span rparams = {}, + GpuTensor* output_override = nullptr, + std::optional matmul_k = std::nullopt, std::uint32_t matmul_batch = 1, + // matmul 段转置布局（bit0=transA，bit1=transB）：**运行期 operand layout， + // 不进 expr_spec_key** —— 同一结构覆盖 4 种转置组合。仅 matmul 形态 + // （matmul / matmul+归约）使用；fold 形态的 mm 段转置仍在 key 里 + // （生成期定死），此处忽略。 + std::uint32_t matmul_trans = 0u, + // fold 收缩轴长度（fold shader 传入填 fold_k 槽；非 fold 传 nullopt） + std::optional fold_k = std::nullopt, + // 输出存储精度（Phase 2 in-kernel f16）：true → 分配 f16 字节布局的 + // 输出缓冲（返回的 GpuTensor 仅是**占位标签**，调用方按此标志重贴 + // GpuTensorF16 —— 见 GpuEngine::eval_expr）。 + bool out_f16 = false, + // 运行期精度分派（键 "#x"）：真实 (输入精度位图, 输出精度)。仅在 + // 分派变体上写入 PC `prec` 槽；其余变体忽略（其精度已是身份）。 + nn::ExprPrecSig prec_sig = 0, + // P2 多输出（extras，IR-C 写穿物化）：主输出之外的**额外输出存储** + // （槽 1..k，与生成器的 `bex` 绑定同序）。仅 V0（全 f32）逐元素 + // shader 支持（`#x`/typed/native16 变体不生成 extras，见 expr_glsl_gen）。 + std::span extra_outs = {}) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + const auto prof_t0 = std::chrono::steady_clock::now(); + const auto it = fused_kernels_.find(shader_name); + if (it == fused_kernels_.end()) NN_FAIL("fused shader not registered: " + shader_name); + const FusedKernel& kernel = it->second; + // （pipeline 绑定推迟到细 tile 选择之后——基础键必须存在，细变体可选） + // 运行期精度分派变体（键后缀 "#x"）：描述符 = 2N 输入视图 + 2 输出视图， + // 且 PC 固定头末尾多一个 `prec` 槽。 + const bool is_disp = std::string_view(shader_name).ends_with(nn::EXPR_PREC_DISPATCH_SUFFIX); + // PC 固定头槽数：**注册时算定，此处只读**（唯一来源 = + // FusedKernel::pc_base_fixed；本地重推会让两侧漂移 → + // push constant 超 range 被驱动丢弃 → 静默读未定义残留）。 + const std::uint32_t pc_base_fixed = kernel.pc_base_fixed; + // 归约轴：-1=逐元素, 0=行归约, 1=列归约 + const int raxis = kernel.reduce_axis; + // 是否含前置 matmul 段（push constants 多 rows + mm_k） + const bool has_mm = kernel.has_matmul; + // fold 形态（注册时从 FusedShader.spec.fold 取，与生成器同源） + const FoldMeta fmeta = kernel.fold; + const bool is_fold = kernel.is_fold(); + // 运行时视图参数个数（RowMod/RotateHalf 的 vp 槽） + const std::uint32_t n_vp = kernel.view_param_count; + // 运行时标量参数个数（优化器 lr/eps/β 的 rp 槽） + const std::uint32_t n_rp = kernel.rparam_count; + + // ── P0 细 tile 选择（wgrad 网格饥饿）──────────────────────────────── + // 纯 matmul 形态（与下方 dispatch 分支同一判据）且 64-tile 网格的 + // 工作组数（含 batch 维）低于阈值 → 派发同结构的 16×16 输出块变体 + // （键 = shader_name + "#f"，注册表元数据/PC 布局/绑定数与 V0 完全 + // 同源，只是网格粒度不同）。`#f` 未登记（f16 精度变体键、设备未注册） + // 或 NN_MM_NO_FINE=1 → 留在 64-tile 路径（降级，不是错误）。 + // 两变体逐输出元素的 K 累加序一致（BK=16 逐块升序 + 同一 vec4 点积 + // 序）→ 结果逐位一致，字节锚不移动。 + bool use_fine = false; + // ── K 切分选择（wgrad 二次归约，键 #ks+#kr）──────────────────────── + // 细 tile 网格（m_per/16 × cols/16 × batch 个 WG）仍低于 + // EXPR_MATMUL_KSPLIT_MIN_WGS → 把 K 切成 S 段：#ks 按列块序号取段 + // [k_lo, k_lo+span)（dispatch x × S）写 (S·rows, cols) 部分和， + // #kr 跨段升序求和 + 跑尾链落最终输出。S = clamp(ceil(阈值/网格), + // 2, MAX_S)（NN_MM_KSPLIT_S 可强制，A/B 调参）；span 上取 16 对齐 + // （生成器 vec4 快路径的段界 k 连续判据依赖它）。 + // 门槛：V0（#ks/#kr 只注册基础键）、纯 matmul、无 extras、f32、 + // #ks+#kr 同时注册（配对）、NN_MM_NO_FINE 与 NN_MM_NO_KSPLIT 均未关。 + // ⚠ 跨段 K 累加序 ≠ 逐 k 串联 → **移动 GPU 字节锚**（已裁定放行， + // 与 P2-2 同处置：重测 + 更新 §12 + history 记新旧值）。 + bool use_ks = false; + std::uint32_t ks_s = 0; // 段数 S + std::uint32_t ks_span = 0; // 段宽（16 对齐） + if (has_mm && raxis < 0 && !is_fold && !mm_fine_disabled()) + { + const std::uint32_t rows_u = static_cast(rows); + const std::uint32_t m_per = rows_u / matmul_batch; // rows 按 batch 整除（契约） + constexpr std::uint32_t B64 = nn::EXPR_MATMUL_BLOCK; + const std::uint64_t wg64 = + static_cast((static_cast(cols) + B64 - 1u) / B64) * + ((m_per + B64 - 1u) / B64) * matmul_batch; + if (wg64 < nn::EXPR_MATMUL_FINE_MIN_WGS && + fused_kernels_.count(shader_name + nn::EXPR_MM_FINE_SUFFIX) != 0) + use_fine = true; + if (wg64 < nn::EXPR_MATMUL_FINE_MIN_WGS && !is_disp && !mm_ksplit_disabled() && + extra_outs.empty() && matmul_k && *matmul_k > 0 && + fused_kernels_.count(shader_name + nn::EXPR_MM_KSPLIT_SUFFIX) != 0 && + fused_kernels_.count(shader_name + nn::EXPR_MM_KSPLIT_REDUCE_SUFFIX) != 0) + { + constexpr std::uint32_t Bf = nn::EXPR_MATMUL_BLOCK_FINE; + const std::uint64_t wg_fine = + static_cast((static_cast(cols) + Bf - 1u) / Bf) * + ((m_per + Bf - 1u) / Bf) * matmul_batch; + std::uint32_t s = mm_ksplit_forced_s(); + if (s == 0u) { - constexpr std::uint32_t Bf = nn::EXPR_MATMUL_BLOCK_FINE; - const std::uint64_t wg_fine = - static_cast( - (static_cast(cols) + Bf - 1u) / Bf) * - ((m_per + Bf - 1u) / Bf) * matmul_batch; - std::uint32_t s = mm_ksplit_forced_s(); - if (s == 0u) - { - if (wg_fine == 0 || wg_fine >= nn::EXPR_MATMUL_KSPLIT_MIN_WGS) - s = 0u; // 网格够密 → 不切 - else - s = static_cast( - (nn::EXPR_MATMUL_KSPLIT_MIN_WGS + wg_fine - 1u) - / wg_fine); - } - if (s >= 2u) + if (wg_fine == 0 || wg_fine >= nn::EXPR_MATMUL_KSPLIT_MIN_WGS) + s = 0u; // 网格够密 → 不切 + else + s = static_cast((nn::EXPR_MATMUL_KSPLIT_MIN_WGS + wg_fine - 1u) / + wg_fine); + } + if (s >= 2u) + { + if (s > nn::EXPR_MATMUL_KSPLIT_MAX_S) s = nn::EXPR_MATMUL_KSPLIT_MAX_S; + const std::uint32_t k = *matmul_k; + const std::uint32_t span = ((k + s - 1u) / s + 15u) & ~15u; // 上取 16 对齐 + // 部分和缓冲 = S·rows·cols 个 float，shader 线性索引是 u32 + const std::uint64_t part = + static_cast(s) * rows * static_cast(cols); + if (span > 0u && part <= 0xFFFFFFFFull) { - if (s > nn::EXPR_MATMUL_KSPLIT_MAX_S) - s = nn::EXPR_MATMUL_KSPLIT_MAX_S; - const std::uint32_t k = *matmul_k; - const std::uint32_t span = - ((k + s - 1u) / s + 15u) & ~15u; // 上取 16 对齐 - // 部分和缓冲 = S·rows·cols 个 float，shader 线性索引是 u32 - const std::uint64_t part = - static_cast(s) * rows * - static_cast(cols); - if (span > 0u && part <= 0xFFFFFFFFull) - { - use_ks = true; - use_fine = false; // K 切分独占派发（不走基础/细 tile） - ks_s = s; - ks_span = span; - } + use_ks = true; + use_fine = false; // K 切分独占派发（不走基础/细 tile） + ks_s = s; + ks_span = span; } } } - const VulkanPipeline& pipeline = - use_fine ? fused_kernels_.at(shader_name + nn::EXPR_MM_FINE_SUFFIX).pipeline - : kernel.pipeline; - - if (inputs.size() + 1 > EXPR_MAX_INPUTS + 1) - NN_FAIL("run_fused_gpu: too many inputs"); - if (consts.size() > EXPR_MAX_CONSTS) - NN_FAIL("run_fused_gpu: too many constants"); - if (view_params.size() != n_vp) - NN_FAIL("run_fused_gpu: view_params count mismatch for " + shader_name); - if (rparams.size() != n_rp) - NN_FAIL("run_fused_gpu: rparams count mismatch for " + shader_name); - if (has_mm && !matmul_k) - NN_FAIL("run_fused_gpu: matmul shader 缺少 matmul_k（求和维度）"); - if (!has_mm && matmul_k && !is_fold) - NN_FAIL("run_fused_gpu: 非 matmul shader 收到了 matmul_k"); - // fold 形态校验（与 matmul_k 同款对称防呆 + 调用约定）： - // fold 输出网格 = (rows, out_cols)，out_cols = vec_state_len（或 1）， - // cols 参数必须等于它（GpuEngine/层调用约定）；PC cols 槽填 fold_k - // （视图/块轴列数=键长，见下方 cols32），不填输出列。 - // fold+matmul 时 matmul_k/matmul_batch 同传（PC slot5/6）。 - if (is_fold && !fold_k) - NN_FAIL("run_fused_gpu: fold shader 缺少 fold_k（收缩轴长度）"); - if (!is_fold && fold_k) - NN_FAIL("run_fused_gpu: 非 fold shader 收到了 fold_k"); - if (is_fold && vector_out) - NN_FAIL("run_fused_gpu: fold shader 以 vector_out 调度（应走 eval_expr）"); - if (is_fold && cols == 0) - NN_FAIL("run_fused_gpu: fold 调用约定 cols（=输出列数 out_cols）必须 > 0"); - // （形状无关：out_cols 不做 shader 侧比对——veclen 不进 key，同一 - // shader 服务任意 d_k，PC vector_out 槽按调用方 cols 运行时填充） - // matmul+归约：raxis >= 0 时 mm_k 填入 PC 第 5 槽（见下方填充） - - // count 以 uint32 传入 shader（gl_GlobalInvocationID / push constant）， - // 必须保证 rows*cols 不溢出 uint32，否则分派与索引会静默截断。 - if (rows > 0 && cols > UINT32_MAX / rows) - NN_FAIL("run_fused_gpu: rows*cols exceeds uint32 range"); - const std::uint32_t count = static_cast(rows * cols); - - // 1. 分配输出 Tensor（vector_out：归约向量原生形状 (rows,1)/(1,cols)） - // 若调用方指定 output_override（dsl::compute_into 的原地目标）， - // 则复用其 buffer（形状必须匹配），结果直接写入目标 → 不额外分配、 - // 不额外拷贝。 - 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; - std::optional owned_output; - GpuTensor* output_ptr = output_override; - if (output_override) - { - if (output_override->rows() != out_rows || output_override->cols() != out_cols) - NN_FAIL("run_fused_gpu: output_override shape mismatch (" + std::to_string(out_rows) + "x" + std::to_string(out_cols) + ")"); - } - else if (out_f16) - { - // f16 输出：按 f16 字节布局分配（2B/元素，偶数槽位对齐见 create_f16_tensor） - auto f16r = GpuTensorF16::create_empty(out_rows, out_cols, *this); - NN_TRY_CHECK(f16r); - owned_output.emplace(f16r->shared_buffer(), out_rows, out_cols); - output_ptr = &*owned_output; - } - else - { - auto output_res = GpuTensor::create_empty(out_rows, out_cols, *this); - NN_TRY_CHECK(output_res); - owned_output.emplace(std::move(*output_res)); - output_ptr = &*owned_output; - } - GpuTensor output = *output_ptr; - - // P2 多输出（extras）：额外输出全部由调用方提供既有存储（写穿物化）， - // 主输出亦须是 output_override（全输出目标显式，无分配路径）。 - if (!extra_outs.empty()) - { - if (is_disp) - NN_FAIL("run_fused_gpu: 多输出（extras）仅支持 V0 变体"); - if (!output_override) - NN_FAIL("run_fused_gpu: 多输出（extras）要求主输出为 output_override"); - for (const GpuTensor* eo : extra_outs) - if (eo == nullptr) - NN_FAIL("run_fused_gpu: 多输出（extras）含空输出指针"); - } - - // ── K 切分两段派发（use_ks，选择判据见上方）──────────────────────── - // #ks：细块网格 × S 段（x = 列块序号决定 K 段）→ (S·rows, cols) - // 部分和；PC 7 槽 = matmul 头 + ks_span（pc_base_fixed=7）。 - // #kr：ceil(count/256) 每线程一输出，段升序求和 + 尾链 → 最终输出； - // PC 6 槽 = matmul 同布局（第 4 槽填段数 S，#kr 不读 mm_trans）。 - // 屏障：inputs 统一就绪 → #ks → partial 写转读（record_output_barrier - // 的 src SHADER_WRITE→dst MEMORY_READ 覆盖 compute→compute）→ #kr → - // 输出写转读。部分和缓冲局部析构 = 归还池；批量录制模式下块可能被 - // **录制序更晚**的算子复用——执行序 = 录制序，#kr 的读排在复用写 - // 之前（同一命令缓冲），故安全（与 owned_output 同一账目）。 - if (use_ks) - { - const VulkanPipeline& ks_pipe = - fused_kernels_.at(shader_name + nn::EXPR_MM_KSPLIT_SUFFIX).pipeline; - const VulkanPipeline& kr_pipe = - fused_kernels_.at(shader_name - + nn::EXPR_MM_KSPLIT_REDUCE_SUFFIX).pipeline; - const std::uint32_t rows32 = static_cast(rows); - const std::uint32_t cols32 = static_cast(cols); - const std::uint32_t k32 = *matmul_k; - - // 1. 部分和缓冲 (S·rows, cols) f32（选择处已验证 u32 索引上界） - auto part_r = GpuTensor::create_empty( - static_cast(ks_s) * rows, cols, *this); - NN_TRY_CHECK(part_r); - GpuTensor partial = std::move(*part_r); - - // 2. 两份 PC（固定头 + vp + consts + rparams，槽位与生成器逐字段同源） - const std::size_t tail_bytes = - sizeof(Scalar) * consts.size() + sizeof(Scalar) * rparams.size(); - std::vector pc_ks( - (7u + n_vp) * sizeof(std::uint32_t) + tail_bytes); - std::vector pc_kr( - (6u + n_vp) * sizeof(std::uint32_t) + tail_bytes); - const auto fill_head = [&](std::vector& pc, int slots, - std::uint32_t slot3, std::uint32_t slot4, - std::uint32_t slot5) - { - std::memcpy(pc.data(), &count, sizeof(std::uint32_t)); - std::memcpy(pc.data() + 4, &cols32, sizeof(std::uint32_t)); - std::memcpy(pc.data() + 8, &rows32, sizeof(std::uint32_t)); - std::memcpy(pc.data() + 12, &slot3, sizeof(std::uint32_t)); - std::memcpy(pc.data() + 16, &slot4, sizeof(std::uint32_t)); - std::memcpy(pc.data() + 20, &slot5, sizeof(std::uint32_t)); - const std::size_t vp_off = - static_cast(slots) * sizeof(std::uint32_t); - for (std::uint32_t i = 0; i < n_vp; ++i) - std::memcpy(pc.data() + vp_off + i * sizeof(std::uint32_t), - &view_params[i], sizeof(std::uint32_t)); - if (!consts.empty()) - std::memcpy(pc.data() + vp_off - + n_vp * sizeof(std::uint32_t), - consts.data(), sizeof(Scalar) * consts.size()); - if (!rparams.empty()) - std::memcpy(pc.data() + vp_off - + n_vp * sizeof(std::uint32_t) - + sizeof(Scalar) * consts.size(), - rparams.data(), - sizeof(Scalar) * rparams.size()); - }; - fill_head(pc_ks, 7, k32, matmul_batch, matmul_trans); - // ks_span 在固定头第 7 槽（offset 24） - std::memcpy(pc_ks.data() + 24, &ks_span, sizeof(std::uint32_t)); - // **有效段数** V = ceil(mm_k / span)（span 16 对齐上取 → 尾部可能 - // 存在 k_lo ≥ mm_k 的空段：#ks 早退不写部分和，#kr 必须只加这 V - // 段，否则读到未初始化缓冲的垃圾——ctest 小形状全挂的教训）。 - const std::uint32_t ks_valid = - (k32 + ks_span - 1u) / ks_span; - fill_head(pc_kr, 6, ks_valid, matmul_batch, 0u); - - // 3. 描述符集：#ks = N 输入 + 1 部分和；#kr = N 输入 + 部分和 + 输出 - const std::size_t n_in = inputs.size(); - std::vector ks_bindings(n_in + 1); - std::vector kr_bindings(n_in + 2); - for (std::size_t i = 0; i < n_in; ++i) - { - ks_bindings[i] = inputs[i]->impl(); - kr_bindings[i] = inputs[i]->impl(); - } - ks_bindings[n_in] = partial.buffer().impl(); - kr_bindings[n_in] = partial.buffer().impl(); - kr_bindings[n_in + 1] = output.buffer().impl(); - - auto ds1_r = alloc_desc_set_for(ks_pipe.descriptor_layout(), ks_bindings); - NN_TRY_CHECK(ds1_r); - VkDescriptorSet ds_ks = *ds1_r; - auto ds2_r = alloc_desc_set_for(kr_pipe.descriptor_layout(), kr_bindings); - if (!ds2_r) - { - release_desc_set_now(ds_ks); - NN_TRY_CHECK(ds2_r); - } - VkDescriptorSet ds_kr = *ds2_r; + } + const VulkanPipeline& pipeline = + use_fine ? fused_kernels_.at(shader_name + nn::EXPR_MM_FINE_SUFFIX).pipeline + : kernel.pipeline; + + if (inputs.size() + 1 > EXPR_MAX_INPUTS + 1) NN_FAIL("run_fused_gpu: too many inputs"); + if (consts.size() > EXPR_MAX_CONSTS) NN_FAIL("run_fused_gpu: too many constants"); + if (view_params.size() != n_vp) + NN_FAIL("run_fused_gpu: view_params count mismatch for " + shader_name); + if (rparams.size() != n_rp) NN_FAIL("run_fused_gpu: rparams count mismatch for " + shader_name); + if (has_mm && !matmul_k) NN_FAIL("run_fused_gpu: matmul shader 缺少 matmul_k（求和维度）"); + if (!has_mm && matmul_k && !is_fold) NN_FAIL("run_fused_gpu: 非 matmul shader 收到了 matmul_k"); + // fold 形态校验（与 matmul_k 同款对称防呆 + 调用约定）： + // fold 输出网格 = (rows, out_cols)，out_cols = vec_state_len（或 1）， + // cols 参数必须等于它（GpuEngine/层调用约定）；PC cols 槽填 fold_k + // （视图/块轴列数=键长，见下方 cols32），不填输出列。 + // fold+matmul 时 matmul_k/matmul_batch 同传（PC slot5/6）。 + if (is_fold && !fold_k) NN_FAIL("run_fused_gpu: fold shader 缺少 fold_k（收缩轴长度）"); + if (!is_fold && fold_k) NN_FAIL("run_fused_gpu: 非 fold shader 收到了 fold_k"); + if (is_fold && vector_out) + NN_FAIL("run_fused_gpu: fold shader 以 vector_out 调度（应走 eval_expr）"); + if (is_fold && cols == 0) + NN_FAIL("run_fused_gpu: fold 调用约定 cols（=输出列数 out_cols）必须 > 0"); + // （形状无关：out_cols 不做 shader 侧比对——veclen 不进 key，同一 + // shader 服务任意 d_k，PC vector_out 槽按调用方 cols 运行时填充） + // matmul+归约：raxis >= 0 时 mm_k 填入 PC 第 5 槽（见下方填充） + + // count 以 uint32 传入 shader（gl_GlobalInvocationID / push constant）， + // 必须保证 rows*cols 不溢出 uint32，否则分派与索引会静默截断。 + if (rows > 0 && cols > UINT32_MAX / rows) + NN_FAIL("run_fused_gpu: rows*cols exceeds uint32 range"); + const std::uint32_t count = static_cast(rows * cols); + + // 1. 分配输出 Tensor（vector_out：归约向量原生形状 (rows,1)/(1,cols)） + // 若调用方指定 output_override（dsl::compute_into 的原地目标）， + // 则复用其 buffer（形状必须匹配），结果直接写入目标 → 不额外分配、 + // 不额外拷贝。 + 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; + std::optional owned_output; + GpuTensor* output_ptr = output_override; + if (output_override) + { + if (output_override->rows() != out_rows || output_override->cols() != out_cols) + NN_FAIL("run_fused_gpu: output_override shape mismatch (" + std::to_string(out_rows) + + "x" + std::to_string(out_cols) + ")"); + } + else if (out_f16) + { + // f16 输出：按 f16 字节布局分配（2B/元素，偶数槽位对齐见 create_f16_tensor） + auto f16r = GpuTensorF16::create_empty(out_rows, out_cols, *this); + NN_TRY_CHECK(f16r); + owned_output.emplace(f16r->shared_buffer(), out_rows, out_cols); + output_ptr = &*owned_output; + } + else + { + auto output_res = GpuTensor::create_empty(out_rows, out_cols, *this); + NN_TRY_CHECK(output_res); + owned_output.emplace(std::move(*output_res)); + output_ptr = &*owned_output; + } + GpuTensor output = *output_ptr; - // 4. 录制：inputs 屏障 → #ks → partial 屏障 → #kr → 输出屏障 - auto cmd_r = acquire_cmd(); - if (!cmd_r) - { - release_desc_set_now(ds_kr); - release_desc_set_now(ds_ks); - NN_TRY_CHECK(cmd_r); - } - auto [cmd, owns_cmd] = *cmd_r; - std::vector in_bufs; - in_bufs.reserve(inputs.size()); - for (const auto* b : inputs) - in_bufs.push_back(b->impl()); - record_input_barriers(cmd, in_bufs); - - vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, ks_pipe.handle()); - vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - ks_pipe.pipeline_layout(), 0, 1, &ds_ks, 0, nullptr); - vkCmdPushConstants(cmd, ks_pipe.pipeline_layout(), - VK_SHADER_STAGE_COMPUTE_BIT, 0, - static_cast(pc_ks.size()), pc_ks.data()); - { - constexpr std::uint32_t Bf = nn::EXPR_MATMUL_BLOCK_FINE; - const std::uint32_t m_per = - rows32 / matmul_batch; - const std::uint32_t ks_valid = - (k32 + ks_span - 1u) / ks_span; - // 只派发**有效段**（空段的 WG 会早退且不写部分和 → 派了也是 - // 白跑）；x = 列块序号决定段号，与生成器 slice = wg.x / nx 同源 - const std::uint32_t gx = - ((cols32 + Bf - 1u) / Bf) * ks_valid; - const std::uint32_t gy = (m_per + Bf - 1u) / Bf; - // 段变体计时进同一 "fused_mm" 桶（#ks/#kr 两行与原单发同维对账） - ts_dispatch(cmd, "fused_mm", gx, gy, matmul_batch, shader_name, - nn::EXPR_MM_KSPLIT_SUFFIX); - } - record_output_barrier(cmd, partial.buffer().impl()); - - vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, kr_pipe.handle()); - vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - kr_pipe.pipeline_layout(), 0, 1, &ds_kr, 0, nullptr); - vkCmdPushConstants(cmd, kr_pipe.pipeline_layout(), - VK_SHADER_STAGE_COMPUTE_BIT, 0, - static_cast(pc_kr.size()), pc_kr.data()); - ts_dispatch(cmd, "fused_mm", (count + 256u - 1u) / 256u, 1u, 1u, - shader_name, nn::EXPR_MM_KSPLIT_REDUCE_SUFFIX); - record_output_barrier(cmd, output.buffer().impl()); - - // 5. 提交（独立模式等待后两个描述符集均可回收；批量模式由帧池 - // reset 统一回收，与既有路径一致）——本块只在 `owns_cmd` - // （⇒ solo）下可达，故这里的直接释放不违反 RV-15 的登记规则 - if (owns_cmd) - { - auto r = submit_and_wait(cmd, ds_ks); - vkFreeDescriptorSets(device_.device(), gpu_tensor_pool_, 1, &ds_kr); - NN_TRY_CHECK(r); - } - if (profile_ops_enabled()) - { - const long long dt = std::chrono::duration_cast< - std::chrono::microseconds>( - std::chrono::steady_clock::now() - prof_t0).count(); - std::fprintf(stderr, - "[gpu-profile] fused fold=0 out=%u total=%lldus key=%.16s ks\n", - static_cast(cols), dt, shader_name.c_str()); - } - return output; - } + // P2 多输出（extras）：额外输出全部由调用方提供既有存储（写穿物化）， + // 主输出亦须是 output_override（全输出目标显式，无分配路径）。 + if (!extra_outs.empty()) + { + if (is_disp) NN_FAIL("run_fused_gpu: 多输出（extras）仅支持 V0 变体"); + if (!output_override) + NN_FAIL("run_fused_gpu: 多输出（extras）要求主输出为 output_override"); + for (const GpuTensor* eo : extra_outs) + if (eo == nullptr) NN_FAIL("run_fused_gpu: 多输出（extras）含空输出指针"); + } - // 2. 分配描述符集（N 输入 + 1 输出；分派变体 = 2N + 2；extras 追加 k 个输出） + // ── K 切分两段派发（use_ks，选择判据见上方）──────────────────────── + // #ks：细块网格 × S 段（x = 列块序号决定 K 段）→ (S·rows, cols) + // 部分和；PC 7 槽 = matmul 头 + ks_span（pc_base_fixed=7）。 + // #kr：ceil(count/256) 每线程一输出，段升序求和 + 尾链 → 最终输出； + // PC 6 槽 = matmul 同布局（第 4 槽填段数 S，#kr 不读 mm_trans）。 + // 屏障：inputs 统一就绪 → #ks → partial 写转读（record_output_barrier + // 的 src SHADER_WRITE→dst MEMORY_READ 覆盖 compute→compute）→ #kr → + // 输出写转读。部分和缓冲局部析构 = 归还池；批量录制模式下块可能被 + // **录制序更晚**的算子复用——执行序 = 录制序，#kr 的读排在复用写 + // 之前（同一命令缓冲），故安全（与 owned_output 同一账目）。 + if (use_ks) + { + const VulkanPipeline& ks_pipe = + fused_kernels_.at(shader_name + nn::EXPR_MM_KSPLIT_SUFFIX).pipeline; + const VulkanPipeline& kr_pipe = + fused_kernels_.at(shader_name + nn::EXPR_MM_KSPLIT_REDUCE_SUFFIX).pipeline; + const std::uint32_t rows32 = static_cast(rows); + const std::uint32_t cols32 = static_cast(cols); + const std::uint32_t k32 = *matmul_k; + + // 1. 部分和缓冲 (S·rows, cols) f32（选择处已验证 u32 索引上界） + auto part_r = GpuTensor::create_empty(static_cast(ks_s) * rows, cols, *this); + NN_TRY_CHECK(part_r); + GpuTensor partial = std::move(*part_r); + + // 2. 两份 PC（固定头 + vp + consts + rparams，槽位与生成器逐字段同源） + const std::size_t tail_bytes = + sizeof(Scalar) * consts.size() + sizeof(Scalar) * rparams.size(); + std::vector pc_ks((7u + n_vp) * sizeof(std::uint32_t) + tail_bytes); + std::vector pc_kr((6u + n_vp) * sizeof(std::uint32_t) + tail_bytes); + const auto fill_head = [&](std::vector& pc, int slots, std::uint32_t slot3, + std::uint32_t slot4, std::uint32_t slot5) + { + std::memcpy(pc.data(), &count, sizeof(std::uint32_t)); + std::memcpy(pc.data() + 4, &cols32, sizeof(std::uint32_t)); + std::memcpy(pc.data() + 8, &rows32, sizeof(std::uint32_t)); + std::memcpy(pc.data() + 12, &slot3, sizeof(std::uint32_t)); + std::memcpy(pc.data() + 16, &slot4, sizeof(std::uint32_t)); + std::memcpy(pc.data() + 20, &slot5, sizeof(std::uint32_t)); + const std::size_t vp_off = static_cast(slots) * sizeof(std::uint32_t); + for (std::uint32_t i = 0; i < n_vp; ++i) + std::memcpy(pc.data() + vp_off + i * sizeof(std::uint32_t), &view_params[i], + sizeof(std::uint32_t)); + if (!consts.empty()) + std::memcpy(pc.data() + vp_off + n_vp * sizeof(std::uint32_t), consts.data(), + sizeof(Scalar) * consts.size()); + if (!rparams.empty()) + std::memcpy(pc.data() + vp_off + n_vp * sizeof(std::uint32_t) + + sizeof(Scalar) * consts.size(), + rparams.data(), sizeof(Scalar) * rparams.size()); + }; + fill_head(pc_ks, 7, k32, matmul_batch, matmul_trans); + // ks_span 在固定头第 7 槽（offset 24） + std::memcpy(pc_ks.data() + 24, &ks_span, sizeof(std::uint32_t)); + // **有效段数** V = ceil(mm_k / span)（span 16 对齐上取 → 尾部可能 + // 存在 k_lo ≥ mm_k 的空段：#ks 早退不写部分和，#kr 必须只加这 V + // 段，否则读到未初始化缓冲的垃圾——ctest 小形状全挂的教训）。 + const std::uint32_t ks_valid = (k32 + ks_span - 1u) / ks_span; + fill_head(pc_kr, 6, ks_valid, matmul_batch, 0u); + + // 3. 描述符集：#ks = N 输入 + 1 部分和；#kr = N 输入 + 部分和 + 输出 const std::size_t n_in = inputs.size(); - const std::size_t n_bindings = - is_disp ? (2 * n_in + 2) : (n_in + 1 + extra_outs.size()); - std::vector bindings(n_bindings); - if (is_disp) + std::vector ks_bindings(n_in + 1); + std::vector kr_bindings(n_in + 2); + for (std::size_t i = 0; i < n_in; ++i) { - // 同一 VkBuffer 绑两次（f32 视图 / f16 视图）；输入对是 readonly、 - // 输出对是 writeonly，shader 每个元素只解引用 PC `prec` 选中的那一个 - // → 另一个从不访问（无别名风险）。 - for (std::size_t i = 0; i < n_in; ++i) - { - bindings[i] = inputs[i]->impl(); - bindings[n_in + i] = inputs[i]->impl(); - } - bindings[2 * n_in] = output.buffer().impl(); - bindings[2 * n_in + 1] = output.buffer().impl(); + ks_bindings[i] = inputs[i]->impl(); + kr_bindings[i] = inputs[i]->impl(); } - else + ks_bindings[n_in] = partial.buffer().impl(); + kr_bindings[n_in] = partial.buffer().impl(); + kr_bindings[n_in + 1] = output.buffer().impl(); + + auto ds1_r = alloc_desc_set_for(ks_pipe.descriptor_layout(), ks_bindings); + NN_TRY_CHECK(ds1_r); + VkDescriptorSet ds_ks = *ds1_r; + auto ds2_r = alloc_desc_set_for(kr_pipe.descriptor_layout(), kr_bindings); + if (!ds2_r) { - for (std::size_t i = 0; i < n_in; ++i) - bindings[i] = inputs[i]->impl(); - bindings[n_in] = output.buffer().impl(); - // P2 多输出：bex 绑定在主输出之后（与生成器 binding 同序） - for (std::size_t k = 0; k < extra_outs.size(); ++k) - bindings[n_in + 1 + k] = extra_outs[k]->buffer().impl(); + release_desc_set_now(ds_ks); + NN_TRY_CHECK(ds2_r); } - auto ds_r = alloc_desc_set_for(pipeline.descriptor_layout(), bindings); - NN_TRY_CHECK(ds_r); - VkDescriptorSet desc_set = *ds_r; + VkDescriptorSet ds_kr = *ds2_r; - // 3. 获取 command buffer + // 4. 录制：inputs 屏障 → #ks → partial 屏障 → #kr → 输出屏障 auto cmd_r = acquire_cmd(); if (!cmd_r) { - release_desc_set_now(desc_set); + release_desc_set_now(ds_kr); + release_desc_set_now(ds_ks); NN_TRY_CHECK(cmd_r); } auto [cmd, owns_cmd] = *cmd_r; - - // 4. 录制：输入屏障 → bind → push constants → dispatch → 输出屏障 std::vector in_bufs; in_bufs.reserve(inputs.size()); for (const auto* b : inputs) in_bufs.push_back(b->impl()); record_input_barriers(cmd, in_bufs); - vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline.handle()); - vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipeline.pipeline_layout(), 0, 1, &desc_set, 0, nullptr); - - // Push constants 布局与 glsl_gen.hpp 一致： - // 逐元素: count, cols, [vp0..], c0.. - // 归约: count, cols, rows, vector_out, [vp0..], c0.. - // matmul: count, cols, rows, mm_k, mm_batch, mm_trans, [vp0..], c0.. - // matmul+归约: count, cols, rows, vector_out, mm_k, mm_batch, mm_trans, - // [vp0..], c0.. - // - // ⚠ 固定头长度**唯一来源 = 注册时的 pc_base_fixed**（本处只读，见上）。 - // 逐形态：逐元素 2 / 逐元素+matmul 6 / 归约 4 / 归约+matmul 7 / - // fold 5 / fold+matmul 7（+ 分派变体 1 个 prec）。 - // 错一个 uint 即常量池整体后移 → shader 从错位处读常量，GPU 上带 - // 常量的归约静默错值而 CPU 正常（实测把 select(cond,1,0) 的常量 - // 读成垃圾）。 - // 分派变体（键 "#x"）在固定头**末尾**多一个 uint prec：生成器同位置 - // 声明（各生成器的 PC 块都在固定槽之后、vp 之前发 `uint prec;`）。 - const std::uint32_t pc_uints = pc_base_fixed + (is_disp ? 1u : 0u); - std::vector pc( - (pc_uints + n_vp) * sizeof(std::uint32_t) + sizeof(Scalar) * consts.size() - + sizeof(Scalar) * rparams.size()); - std::memcpy(pc.data(), &count, sizeof(std::uint32_t)); - // fold：PC cols 槽 = fold_k（生成器视图/块轴用输入列数 K 索引； - // 输出网格列恒 1 不进 PC）。其余形态 = 调用方 cols。 - const std::uint32_t cols32 = is_fold - ? *fold_k : static_cast(cols); - std::memcpy(pc.data() + sizeof(std::uint32_t), &cols32, sizeof(std::uint32_t)); - if (raxis >= 0) - { - const std::uint32_t rows32 = static_cast(rows); - std::memcpy(pc.data() + 2 * sizeof(std::uint32_t), &rows32, - sizeof(std::uint32_t)); - const std::uint32_t vo = vector_out ? 1u : 0u; - std::memcpy(pc.data() + 3 * sizeof(std::uint32_t), &vo, - sizeof(std::uint32_t)); - if (has_mm) - { - // matmul+归约：mm_k + mm_batch + mm_trans（形状参数 + 运行期 - // operand layout，运行时填充） - std::memcpy(pc.data() + 4 * sizeof(std::uint32_t), &*matmul_k, - sizeof(std::uint32_t)); - std::memcpy(pc.data() + 5 * sizeof(std::uint32_t), &matmul_batch, - sizeof(std::uint32_t)); - std::memcpy(pc.data() + 6 * sizeof(std::uint32_t), &matmul_trans, - sizeof(std::uint32_t)); - } - } - else if (has_mm) - { - // matmul 融合：rows + mm_k + mm_batch + mm_trans（形状参数 + 运行期 - // operand layout，运行时按实际 spec 填充 → 同结构不同 K/batch/转置 - // 共享一个融合 shader） - const std::uint32_t rows32 = static_cast(rows); - std::memcpy(pc.data() + 2 * sizeof(std::uint32_t), &rows32, - sizeof(std::uint32_t)); - std::memcpy(pc.data() + 3 * sizeof(std::uint32_t), &*matmul_k, - sizeof(std::uint32_t)); - std::memcpy(pc.data() + 4 * sizeof(std::uint32_t), &matmul_batch, - sizeof(std::uint32_t)); - std::memcpy(pc.data() + 5 * sizeof(std::uint32_t), &matmul_trans, - sizeof(std::uint32_t)); - } - else if (is_fold) - { - // fold：rows + **vector_out = 调用方输出列数** - // （形状无关——veclen 不进 key，PC 运行时填充；v1 生成器 - // 不读此槽、填 1 无害）+ fold_k + [mm_k, mm_batch]（7 槽形态）。 - // 与生成器 PC 声明逐字段一致（4.10：创建侧 range 必须同改） - const std::uint32_t rows32 = static_cast(rows); - const std::uint32_t vo = static_cast(cols); - std::memcpy(pc.data() + 2 * sizeof(std::uint32_t), &rows32, - sizeof(std::uint32_t)); - std::memcpy(pc.data() + 3 * sizeof(std::uint32_t), &vo, - sizeof(std::uint32_t)); - std::memcpy(pc.data() + 4 * sizeof(std::uint32_t), &*fold_k, - sizeof(std::uint32_t)); - if (matmul_k) - { - std::memcpy(pc.data() + 5 * sizeof(std::uint32_t), &*matmul_k, - sizeof(std::uint32_t)); - std::memcpy(pc.data() + 6 * sizeof(std::uint32_t), &matmul_batch, - sizeof(std::uint32_t)); - } - } - // 运行期精度分派：把真实 (输入精度位图, 输出精度) 写进固定头末尾的 - // prec 槽——shader 据此在每个加载/存储点走 uniform 分支。 - if (is_disp) - std::memcpy(pc.data() + pc_base_fixed * sizeof(std::uint32_t), - &prec_sig, sizeof(prec_sig)); - // 运行时视图参数（RowMod 周期 / RotateHalf 块大小），置于固定头之后、常量池之前 - for (std::uint32_t i = 0; i < n_vp; ++i) - std::memcpy(pc.data() + (pc_uints + i) * sizeof(std::uint32_t), - &view_params[i], sizeof(std::uint32_t)); - if (!consts.empty()) - std::memcpy(pc.data() + (pc_uints + n_vp) * sizeof(std::uint32_t), consts.data(), - sizeof(Scalar) * consts.size()); - // 运行时标量参数（优化器 lr/eps/β 等 float rp 槽），置于常量池之后 - if (!rparams.empty()) - std::memcpy(pc.data() + (pc_uints + n_vp) * sizeof(std::uint32_t) + - sizeof(Scalar) * consts.size(), - rparams.data(), sizeof(Scalar) * rparams.size()); - vkCmdPushConstants(cmd, pipeline.pipeline_layout(), - VK_SHADER_STAGE_COMPUTE_BIT, 0, static_cast(pc.size()), pc.data()); - - // dispatch：逐元素 = ceil(count/(256*vec_width))；行归约 = rows 个工作组； - // 列归约 = ceil(cols/32) 个工作组（每工作组 32 列 tile、warp=行块， - // 与生成器 l/wb/col 结构及 OP 级 reduce.comp 契约同源）； - // matmul 分块 = (ceil(cols/BLOCK), ceil(m_per/BLOCK), matmul_batch)， - // BLOCK 与 glsl_gen 生成的输出块一致（EXPR_MATMUL_BLOCK=64：每工作组 - // 64×64 输出块、16×16 线程、每线程 4×4 寄存器分块；网格饥饿时的 - // 细 tile 变体 BLOCK=EXPR_MATMUL_BLOCK_FINE=16，见上方 use_fine） - const std::uint32_t vec_width = kernel.vec_width; - if (has_mm && raxis < 0 && !is_fold) - { - // P0 细 tile：use_fine 时输出块 64 → 16（同一 struct 的两个变体， - // 只是网格粒度；生成器同源，见 EXPR_MATMUL_BLOCK_FINE） - const std::uint32_t mm_block = use_fine - ? nn::EXPR_MATMUL_BLOCK_FINE : nn::EXPR_MATMUL_BLOCK; - const std::uint32_t wg_x = - (static_cast(cols) + mm_block - 1u) / mm_block; - // batch：dispatch z = 批次，A/B 按 batch 垂直切分 → y 只覆盖 - // **批内**行 m_per = rows/mm_batch（生成器 main/load_tiles 均以 m_per - // 为界、写回守卫 rr < m_per）。按总 rows 派 y 会把 batch 在 y/z 数 - // 两遍 → 工作量 ∝ batch²：多余 (BH−1)/BH 的 WG 跑完整条 mm_k 流水 - // 后整块丢弃——结果仍正确、对拍永远测不出，只有计时能发现。rows = - // batch*M 按契约整除。 - const std::uint32_t rows_u = static_cast(rows); - const std::uint32_t m_per = rows_u / matmul_batch; - const std::uint32_t wg_y = (m_per + mm_block - 1u) / mm_block; - // batch：dispatch z = 批次，A/B 按 batch 垂直切分 - ts_dispatch(cmd, "fused_mm", wg_x, wg_y, matmul_batch, - shader_name, - use_fine ? nn::EXPR_MM_FINE_SUFFIX : nullptr); - } - else - { - // fold 的 dispatch 分两种（FoldMeta，与生成器分派同源）： - // v1 标量 fold（每线程一行，row=gl_GlobalInvocationID.x）： - // count=rows×1 → ceil(count/256) 恰 = ceil(rows/256) ✓ - // v2 双域 fold（每 WG NR 行，row=wg*NR+ri，EXPR_FOLD_ROWS_PER_WG - // 与生成器行循环同源）：**wg = ceil(rows/NR)** - // （按 ceil(count/256)=ceil(rows×out/256) 派会丢行！） - const std::uint32_t wg_count = - (is_fold && !fmeta.per_thread_row) - ? (static_cast(rows) - + nn::EXPR_FOLD_ROWS_PER_WG - 1u) - / nn::EXPR_FOLD_ROWS_PER_WG - : (raxis == 0) ? static_cast(rows) - : (raxis == 1) ? (static_cast(cols) + 31u) / 32u - : (count + 256u * vec_width - 1u) / (256u * vec_width); - ts_dispatch(cmd, "fused", wg_count, 1, 1, shader_name); - } + vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, ks_pipe.handle()); + vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, ks_pipe.pipeline_layout(), 0, + 1, &ds_ks, 0, nullptr); + vkCmdPushConstants(cmd, ks_pipe.pipeline_layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, + static_cast(pc_ks.size()), pc_ks.data()); + { + constexpr std::uint32_t Bf = nn::EXPR_MATMUL_BLOCK_FINE; + const std::uint32_t m_per = rows32 / matmul_batch; + const std::uint32_t ks_valid = (k32 + ks_span - 1u) / ks_span; + // 只派发**有效段**（空段的 WG 会早退且不写部分和 → 派了也是 + // 白跑）；x = 列块序号决定段号，与生成器 slice = wg.x / nx 同源 + const std::uint32_t gx = ((cols32 + Bf - 1u) / Bf) * ks_valid; + const std::uint32_t gy = (m_per + Bf - 1u) / Bf; + // 段变体计时进同一 "fused_mm" 桶（#ks/#kr 两行与原单发同维对账） + ts_dispatch(cmd, "fused_mm", gx, gy, matmul_batch, shader_name, + nn::EXPR_MM_KSPLIT_SUFFIX); + } + record_output_barrier(cmd, partial.buffer().impl()); + + vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, kr_pipe.handle()); + vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, kr_pipe.pipeline_layout(), 0, + 1, &ds_kr, 0, nullptr); + vkCmdPushConstants(cmd, kr_pipe.pipeline_layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, + static_cast(pc_kr.size()), pc_kr.data()); + ts_dispatch(cmd, "fused_mm", (count + 256u - 1u) / 256u, 1u, 1u, shader_name, + nn::EXPR_MM_KSPLIT_REDUCE_SUFFIX); record_output_barrier(cmd, output.buffer().impl()); - // P2 多输出：额外输出同样过写后屏障 - for (const GpuTensor* eo : extra_outs) - record_output_barrier(cmd, eo->buffer().impl()); - // 5. 独立模式提交 + // 5. 提交（独立模式等待后两个描述符集均可回收；批量模式由帧池 + // reset 统一回收，与既有路径一致）——本块只在 `owns_cmd` + // （⇒ solo）下可达，故这里的直接释放不违反 RV-15 的登记规则 if (owns_cmd) { - auto r = submit_and_wait(cmd, desc_set); + auto r = submit_and_wait(cmd, ds_ks); + vkFreeDescriptorSets(device_.device(), gpu_tensor_pool_, 1, &ds_kr); NN_TRY_CHECK(r); } if (profile_ops_enabled()) { const long long dt = std::chrono::duration_cast( - std::chrono::steady_clock::now() - prof_t0).count(); - // fold 按形态如实打印（is_fold = 本形态是否 fold）；out= 运行时 - // 列数（veclen 不进 key，注册侧值可能与调用方实际 cols 不同—— - // 打印注册侧列数会误导归因） - std::fprintf(stderr, - "[gpu-profile] fused fold=%d out=%u total=%lldus key=%.16s\n", - static_cast(is_fold), - static_cast(cols), dt, shader_name.c_str()); + std::chrono::steady_clock::now() - prof_t0) + .count(); + std::fprintf(stderr, "[gpu-profile] fused fold=0 out=%u total=%lldus key=%.16s ks\n", + static_cast(cols), dt, shader_name.c_str()); } return output; } - // ══════════════════════════════════════════════════════════════════ - // run_fused_gpu_batch — P3 跨链批量派发（键 = key#b） - // - // 同结构（同 key）、同参数值（consts/rparams/vp）、互不依赖的 N 个实例 - // 合成**一次 dispatch**：实例表（uvec2 数组，见 FusedBatchInst 注释）经 - // vkCmdUpdateBuffer 在录制期内联写入一个小 device-local 缓冲，shader 用 - // buffer device address 按槽位解引用各实例的输入/输出缓冲。 - // - // 与生成器（expr_glsl_gen 的 batched 分支）的契约： - // Bindings: 仅 binding 0（实例表，storage buffer） - // Push Constants = V0 逐元素同布局：count(=max_count), cols, [vp..], - // [c..], [rp..]（批量路径不读 count/cols 两槽——元素数 - // 与列数取每实例尾槽，见生成器 `#define cols`） - // dispatch = (ceil(max_count/256), N, 1)，实例下标 = gl_WorkGroupID.y - // - // 前置（调用方 = execute_fused_graph 保证，本函数只防御性校验）： - // `key#b` 已注册、设备 bda_available、实例全 f32、组内互不依赖。 - // ══════════════════════════════════════════════════════════════════ - [[nodiscard]] Result run_fused_gpu_batch( - const std::string& shader_name, - std::span instances, - std::span consts, - std::span view_params = {}, - std::span rparams = {}) - { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - const auto prof_t0 = std::chrono::steady_clock::now(); - if (!device_.bda_available()) - NN_FAIL("run_fused_gpu_batch: 设备未启用 buffer device address" - "（应先逐 kernel 降级，不该到此）"); - if (!std::string_view(shader_name).ends_with(nn::EXPR_BATCH_SUFFIX)) - NN_FAIL("run_fused_gpu_batch: 非批量变体键（应以 #b 结尾）：" + shader_name); - const auto it = fused_kernels_.find(shader_name); - if (it == fused_kernels_.end()) - NN_FAIL("fused shader not registered: " + shader_name); - const FusedKernel& kernel = it->second; - const VulkanPipeline& pipeline = kernel.pipeline; - if (instances.empty()) - return {}; - // PC 槽位校验（与 run_fused_gpu 同源：vp/rp 计数来自注册表元数据） - const std::uint32_t n_vp = kernel.view_param_count; - const std::uint32_t n_rp = kernel.rparam_count; - if (view_params.size() != n_vp) - NN_FAIL("run_fused_gpu_batch: view_params count mismatch for " + shader_name); - if (rparams.size() != n_rp) - NN_FAIL("run_fused_gpu_batch: rparams count mismatch for " + shader_name); - const std::uint32_t pc_base_fixed = kernel.pc_base_fixed; - - // 实例表布局：stride S = 输入 + 主输出 + extras + (count,cols) 尾槽 - const std::size_t n_in = instances[0].inputs.size(); - const std::size_t n_ex = instances[0].extra_outs.size(); - const std::size_t stride = n_in + 1 + n_ex + 1; - if (kernel.batch_stride != 0 && - kernel.batch_stride != static_cast(stride)) - NN_FAIL("run_fused_gpu_batch: 实例槽数与注册结构不一致：" + shader_name); - for (const FusedBatchInst& inst : instances) - { - if (inst.inputs.size() != n_in || inst.extra_outs.size() != n_ex) - NN_FAIL("run_fused_gpu_batch: 组内实例结构不一致"); - if (inst.out == nullptr) - NN_FAIL("run_fused_gpu_batch: 实例缺主输出缓冲"); - for (const GpuBuffer* b : inst.inputs) - if (b == nullptr) - NN_FAIL("run_fused_gpu_batch: 实例缺输入缓冲"); - for (const GpuBuffer* b : inst.extra_outs) - if (b == nullptr) - NN_FAIL("run_fused_gpu_batch: 实例缺 extras 输出缓冲"); - } - // vkCmdUpdateBuffer 的硬上限：单次 ≤ 64KB、4 字节对齐 - if (instances.size() * stride * sizeof(std::uint64_t) > 65536u) - NN_FAIL("run_fused_gpu_batch: 实例表超过 vkCmdUpdateBuffer 64KB 上限"); - - // ── 实例表内容：uvec2 = (地址低 32 位, 高 32 位)；尾槽 = (count, cols) ── - std::vector table(instances.size() * stride * 2, 0u); - std::uint32_t max_count = 0; - std::uint32_t max_cols = 0; - for (std::size_t i = 0; i < instances.size(); ++i) - { - const FusedBatchInst& inst = instances[i]; - const std::size_t base = i * stride * 2; - bool addr_ok = true; - const auto put_addr = [&](std::size_t slot, const GpuBuffer* b) -> bool - { - VkBufferDeviceAddressInfo ai{}; - ai.sType = VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO; - ai.buffer = b->impl(); - const VkDeviceAddress addr = - vkGetBufferDeviceAddress(device_.device(), &ai); - table[base + slot * 2 + 0] = - static_cast(addr & 0xffffffffull); - table[base + slot * 2 + 1] = - static_cast(addr >> 32); - return addr != 0; - }; - for (std::size_t j = 0; j < n_in; ++j) - addr_ok = put_addr(j, inst.inputs[j]) && addr_ok; - addr_ok = put_addr(n_in, inst.out) && addr_ok; - for (std::size_t k = 0; k < n_ex; ++k) - addr_ok = put_addr(n_in + 1 + k, inst.extra_outs[k]) && addr_ok; - if (!addr_ok) - NN_FAIL("run_fused_gpu_batch: vkGetBufferDeviceAddress 返回 0" - "（buffer 未带 SHADER_DEVICE_ADDRESS？）"); - table[base + (stride - 1) * 2 + 0] = inst.count; - table[base + (stride - 1) * 2 + 1] = inst.cols; - max_count = std::max(max_count, inst.count); - max_cols = std::max(max_cols, inst.cols); - } - - // 实例表缓冲：小 device-local（TRANSFER_DST 走 vkCmdUpdateBuffer）。 - // 数据内联进命令缓冲 → 录制生命周期安全；缓冲本体析构走 pending_destroys_ - // 延迟队列（铁律 6，与其余录制期资源同款）。 - auto inst_r = GpuBuffer::create_device_local( - device_.device(), alloc_pool(), table.size() * sizeof(std::uint32_t), - VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, - keepalive()); - NN_TRY_CHECK(inst_r); - std::shared_ptr inst_buf = - std::make_shared(std::move(*inst_r)); - - // 描述符：key#b pipeline 只绑 binding 0（实例表）——复用 per-pipeline - // 描述符机制 - // 批量变体只有一个 binding：实例表 - const VkBuffer instance_table = inst_buf->impl(); - auto ds_r = alloc_desc_set_for(pipeline.descriptor_layout(), {&instance_table, 1}); - NN_TRY_CHECK(ds_r); - VkDescriptorSet desc_set = *ds_r; - - // command buffer - auto cmd_r = acquire_cmd(); - if (!cmd_r) - { - release_desc_set_now(desc_set); - NN_TRY_CHECK(cmd_r); - } - auto [cmd, owns_cmd] = *cmd_r; - - // 录制：实例表写入 → 输入屏障（含实例表，去重）→ bind → PC → dispatch - vkCmdUpdateBuffer(cmd, inst_buf->impl(), 0, - static_cast( - table.size() * sizeof(std::uint32_t)), - table.data()); - { - // 屏障覆盖组内全部实例的输入 buffer + 实例表（vkCmdUpdateBuffer 是 - // TRANSFER 写，record_input_barriers 的 src 阶段含 TRANSFER）； - // 同一 buffer 去重（组内互不依赖已由引擎保证，仅防重复屏障）。 - std::vector in_bufs; - in_bufs.reserve(instances.size() * (n_in + 1) + 1); - for (const FusedBatchInst& inst : instances) - for (const GpuBuffer* b : inst.inputs) - in_bufs.push_back(b->impl()); - in_bufs.push_back(inst_buf->impl()); - std::sort(in_bufs.begin(), in_bufs.end()); - in_bufs.erase(std::unique(in_bufs.begin(), in_bufs.end()), - in_bufs.end()); - record_input_barriers(cmd, in_bufs); - } - vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline.handle()); - vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipeline.pipeline_layout(), 0, 1, &desc_set, 0, nullptr); - - // Push constants 与 run_fused_gpu 的 V0 逐元素分支同源： - // count(=max_count), cols, [vp0..], [c0..], [rp0..] - // 批量 shader 不读 count/cols 两槽（每实例取尾槽），两值仅保持 V0 布局。 - const std::uint32_t pc_uints = pc_base_fixed + n_vp; - std::vector pc( - (pc_uints) * sizeof(std::uint32_t) + sizeof(Scalar) * consts.size() - + sizeof(Scalar) * rparams.size()); - std::memcpy(pc.data(), &max_count, sizeof(std::uint32_t)); - std::memcpy(pc.data() + sizeof(std::uint32_t), &max_cols, sizeof(std::uint32_t)); - for (std::uint32_t i = 0; i < n_vp; ++i) - std::memcpy(pc.data() + (pc_base_fixed + i) * sizeof(std::uint32_t), - &view_params[i], sizeof(std::uint32_t)); - if (!consts.empty()) - std::memcpy(pc.data() + (pc_uints) * sizeof(std::uint32_t), consts.data(), - sizeof(Scalar) * consts.size()); - if (!rparams.empty()) - std::memcpy(pc.data() + (pc_uints) * sizeof(std::uint32_t) + - sizeof(Scalar) * consts.size(), - rparams.data(), sizeof(Scalar) * rparams.size()); - vkCmdPushConstants(cmd, pipeline.pipeline_layout(), - VK_SHADER_STAGE_COMPUTE_BIT, 0, - static_cast(pc.size()), pc.data()); - - // dispatch = (ceil(max_count/256), n_instances, 1)；批量恒标量 - // （vec_width=1，与生成器同源） - const std::uint32_t wg_x = (max_count + 255u) / 256u; - ts_dispatch(cmd, "fused_b", wg_x, - static_cast(instances.size()), 1, - shader_name); - - // 输出屏障（主输出 + 全部 extras，去重）：组内输出互不重叠（引擎保证） + // 2. 分配描述符集（N 输入 + 1 输出；分派变体 = 2N + 2；extras 追加 k 个输出） + const std::size_t n_in = inputs.size(); + const std::size_t n_bindings = is_disp ? (2 * n_in + 2) : (n_in + 1 + extra_outs.size()); + std::vector bindings(n_bindings); + if (is_disp) + { + // 同一 VkBuffer 绑两次（f32 视图 / f16 视图）；输入对是 readonly、 + // 输出对是 writeonly，shader 每个元素只解引用 PC `prec` 选中的那一个 + // → 另一个从不访问（无别名风险）。 + for (std::size_t i = 0; i < n_in; ++i) { - std::vector out_bufs; - out_bufs.reserve(instances.size() * (1 + n_ex)); - for (const FusedBatchInst& inst : instances) - { - out_bufs.push_back(inst.out->impl()); - for (const GpuBuffer* b : inst.extra_outs) - out_bufs.push_back(b->impl()); - } - std::sort(out_bufs.begin(), out_bufs.end()); - out_bufs.erase(std::unique(out_bufs.begin(), out_bufs.end()), - out_bufs.end()); - for (VkBuffer b : out_bufs) - record_output_barrier(cmd, b); + bindings[i] = inputs[i]->impl(); + bindings[n_in + i] = inputs[i]->impl(); } + bindings[2 * n_in] = output.buffer().impl(); + bindings[2 * n_in + 1] = output.buffer().impl(); + } + else + { + for (std::size_t i = 0; i < n_in; ++i) + bindings[i] = inputs[i]->impl(); + bindings[n_in] = output.buffer().impl(); + // P2 多输出：bex 绑定在主输出之后（与生成器 binding 同序） + for (std::size_t k = 0; k < extra_outs.size(); ++k) + bindings[n_in + 1 + k] = extra_outs[k]->buffer().impl(); + } + auto ds_r = alloc_desc_set_for(pipeline.descriptor_layout(), bindings); + NN_TRY_CHECK(ds_r); + VkDescriptorSet desc_set = *ds_r; - if (owns_cmd) - { - auto r = submit_and_wait(cmd, desc_set); - NN_TRY_CHECK(r); + // 3. 获取 command buffer + auto cmd_r = acquire_cmd(); + if (!cmd_r) + { + release_desc_set_now(desc_set); + NN_TRY_CHECK(cmd_r); + } + auto [cmd, owns_cmd] = *cmd_r; + + // 4. 录制：输入屏障 → bind → push constants → dispatch → 输出屏障 + std::vector in_bufs; + in_bufs.reserve(inputs.size()); + for (const auto* b : inputs) + in_bufs.push_back(b->impl()); + record_input_barriers(cmd, in_bufs); + + vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline.handle()); + vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline.pipeline_layout(), 0, 1, + &desc_set, 0, nullptr); + + // Push constants 布局与 glsl_gen.hpp 一致： + // 逐元素: count, cols, [vp0..], c0.. + // 归约: count, cols, rows, vector_out, [vp0..], c0.. + // matmul: count, cols, rows, mm_k, mm_batch, mm_trans, [vp0..], c0.. + // matmul+归约: count, cols, rows, vector_out, mm_k, mm_batch, mm_trans, + // [vp0..], c0.. + // + // ⚠ 固定头长度**唯一来源 = 注册时的 pc_base_fixed**（本处只读，见上）。 + // 逐形态：逐元素 2 / 逐元素+matmul 6 / 归约 4 / 归约+matmul 7 / + // fold 5 / fold+matmul 7（+ 分派变体 1 个 prec）。 + // 错一个 uint 即常量池整体后移 → shader 从错位处读常量，GPU 上带 + // 常量的归约静默错值而 CPU 正常（实测把 select(cond,1,0) 的常量 + // 读成垃圾）。 + // 分派变体（键 "#x"）在固定头**末尾**多一个 uint prec：生成器同位置 + // 声明（各生成器的 PC 块都在固定槽之后、vp 之前发 `uint prec;`）。 + const std::uint32_t pc_uints = pc_base_fixed + (is_disp ? 1u : 0u); + std::vector pc((pc_uints + n_vp) * sizeof(std::uint32_t) + + sizeof(Scalar) * consts.size() + sizeof(Scalar) * rparams.size()); + std::memcpy(pc.data(), &count, sizeof(std::uint32_t)); + // fold：PC cols 槽 = fold_k（生成器视图/块轴用输入列数 K 索引； + // 输出网格列恒 1 不进 PC）。其余形态 = 调用方 cols。 + const std::uint32_t cols32 = is_fold ? *fold_k : static_cast(cols); + std::memcpy(pc.data() + sizeof(std::uint32_t), &cols32, sizeof(std::uint32_t)); + if (raxis >= 0) + { + const std::uint32_t rows32 = static_cast(rows); + std::memcpy(pc.data() + 2 * sizeof(std::uint32_t), &rows32, sizeof(std::uint32_t)); + const std::uint32_t vo = vector_out ? 1u : 0u; + std::memcpy(pc.data() + 3 * sizeof(std::uint32_t), &vo, sizeof(std::uint32_t)); + if (has_mm) + { + // matmul+归约：mm_k + mm_batch + mm_trans（形状参数 + 运行期 + // operand layout，运行时填充） + std::memcpy(pc.data() + 4 * sizeof(std::uint32_t), &*matmul_k, sizeof(std::uint32_t)); + std::memcpy(pc.data() + 5 * sizeof(std::uint32_t), &matmul_batch, + sizeof(std::uint32_t)); + std::memcpy(pc.data() + 6 * sizeof(std::uint32_t), &matmul_trans, + sizeof(std::uint32_t)); } - ++fused_batch_dispatches_; - if (profile_ops_enabled()) - { - const long long dt = std::chrono::duration_cast( - std::chrono::steady_clock::now() - prof_t0).count(); - std::fprintf(stderr, - "[gpu-profile] fused_batch inst=%zu max_count=%u total=%lldus " - "key=%.16s\n", - instances.size(), max_count, dt, shader_name.c_str()); + } + else if (has_mm) + { + // matmul 融合：rows + mm_k + mm_batch + mm_trans（形状参数 + 运行期 + // operand layout，运行时按实际 spec 填充 → 同结构不同 K/batch/转置 + // 共享一个融合 shader） + const std::uint32_t rows32 = static_cast(rows); + std::memcpy(pc.data() + 2 * sizeof(std::uint32_t), &rows32, sizeof(std::uint32_t)); + std::memcpy(pc.data() + 3 * sizeof(std::uint32_t), &*matmul_k, sizeof(std::uint32_t)); + std::memcpy(pc.data() + 4 * sizeof(std::uint32_t), &matmul_batch, sizeof(std::uint32_t)); + std::memcpy(pc.data() + 5 * sizeof(std::uint32_t), &matmul_trans, sizeof(std::uint32_t)); + } + else if (is_fold) + { + // fold：rows + **vector_out = 调用方输出列数** + // （形状无关——veclen 不进 key，PC 运行时填充；v1 生成器 + // 不读此槽、填 1 无害）+ fold_k + [mm_k, mm_batch]（7 槽形态）。 + // 与生成器 PC 声明逐字段一致（4.10：创建侧 range 必须同改） + const std::uint32_t rows32 = static_cast(rows); + const std::uint32_t vo = static_cast(cols); + std::memcpy(pc.data() + 2 * sizeof(std::uint32_t), &rows32, sizeof(std::uint32_t)); + std::memcpy(pc.data() + 3 * sizeof(std::uint32_t), &vo, sizeof(std::uint32_t)); + std::memcpy(pc.data() + 4 * sizeof(std::uint32_t), &*fold_k, sizeof(std::uint32_t)); + if (matmul_k) + { + std::memcpy(pc.data() + 5 * sizeof(std::uint32_t), &*matmul_k, sizeof(std::uint32_t)); + std::memcpy(pc.data() + 6 * sizeof(std::uint32_t), &matmul_batch, + sizeof(std::uint32_t)); } - return {}; } + // 运行期精度分派：把真实 (输入精度位图, 输出精度) 写进固定头末尾的 + // prec 槽——shader 据此在每个加载/存储点走 uniform 分支。 + if (is_disp) + std::memcpy(pc.data() + pc_base_fixed * sizeof(std::uint32_t), &prec_sig, sizeof(prec_sig)); + // 运行时视图参数（RowMod 周期 / RotateHalf 块大小），置于固定头之后、常量池之前 + for (std::uint32_t i = 0; i < n_vp; ++i) + std::memcpy(pc.data() + (pc_uints + i) * sizeof(std::uint32_t), &view_params[i], + sizeof(std::uint32_t)); + if (!consts.empty()) + std::memcpy(pc.data() + (pc_uints + n_vp) * sizeof(std::uint32_t), consts.data(), + sizeof(Scalar) * consts.size()); + // 运行时标量参数（优化器 lr/eps/β 等 float rp 槽），置于常量池之后 + if (!rparams.empty()) + std::memcpy(pc.data() + (pc_uints + n_vp) * sizeof(std::uint32_t) + + sizeof(Scalar) * consts.size(), + rparams.data(), sizeof(Scalar) * rparams.size()); + vkCmdPushConstants(cmd, pipeline.pipeline_layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, + static_cast(pc.size()), pc.data()); + + // dispatch：逐元素 = ceil(count/(256*vec_width))；行归约 = rows 个工作组； + // 列归约 = ceil(cols/32) 个工作组（每工作组 32 列 tile、warp=行块， + // 与生成器 l/wb/col 结构及 OP 级 reduce.comp 契约同源）； + // matmul 分块 = (ceil(cols/BLOCK), ceil(m_per/BLOCK), matmul_batch)， + // BLOCK 与 glsl_gen 生成的输出块一致（EXPR_MATMUL_BLOCK=64：每工作组 + // 64×64 输出块、16×16 线程、每线程 4×4 寄存器分块；网格饥饿时的 + // 细 tile 变体 BLOCK=EXPR_MATMUL_BLOCK_FINE=16，见上方 use_fine） + const std::uint32_t vec_width = kernel.vec_width; + if (has_mm && raxis < 0 && !is_fold) + { + // P0 细 tile：use_fine 时输出块 64 → 16（同一 struct 的两个变体， + // 只是网格粒度；生成器同源，见 EXPR_MATMUL_BLOCK_FINE） + const std::uint32_t mm_block = + use_fine ? nn::EXPR_MATMUL_BLOCK_FINE : nn::EXPR_MATMUL_BLOCK; + const std::uint32_t wg_x = (static_cast(cols) + mm_block - 1u) / mm_block; + // batch：dispatch z = 批次，A/B 按 batch 垂直切分 → y 只覆盖 + // **批内**行 m_per = rows/mm_batch（生成器 main/load_tiles 均以 m_per + // 为界、写回守卫 rr < m_per）。按总 rows 派 y 会把 batch 在 y/z 数 + // 两遍 → 工作量 ∝ batch²：多余 (BH−1)/BH 的 WG 跑完整条 mm_k 流水 + // 后整块丢弃——结果仍正确、对拍永远测不出，只有计时能发现。rows = + // batch*M 按契约整除。 + const std::uint32_t rows_u = static_cast(rows); + const std::uint32_t m_per = rows_u / matmul_batch; + const std::uint32_t wg_y = (m_per + mm_block - 1u) / mm_block; + // batch：dispatch z = 批次，A/B 按 batch 垂直切分 + ts_dispatch(cmd, "fused_mm", wg_x, wg_y, matmul_batch, shader_name, + use_fine ? nn::EXPR_MM_FINE_SUFFIX : nullptr); + } + else + { + // fold 的 dispatch 分两种（FoldMeta，与生成器分派同源）： + // v1 标量 fold（每线程一行，row=gl_GlobalInvocationID.x）： + // count=rows×1 → ceil(count/256) 恰 = ceil(rows/256) ✓ + // v2 双域 fold（每 WG NR 行，row=wg*NR+ri，EXPR_FOLD_ROWS_PER_WG + // 与生成器行循环同源）：**wg = ceil(rows/NR)** + // （按 ceil(count/256)=ceil(rows×out/256) 派会丢行！） + const std::uint32_t wg_count = + (is_fold && !fmeta.per_thread_row) + ? (static_cast(rows) + nn::EXPR_FOLD_ROWS_PER_WG - 1u) / + nn::EXPR_FOLD_ROWS_PER_WG + : (raxis == 0) ? static_cast(rows) + : (raxis == 1) ? (static_cast(cols) + 31u) / 32u + : (count + 256u * vec_width - 1u) / (256u * vec_width); + ts_dispatch(cmd, "fused", wg_count, 1, 1, shader_name); + } + record_output_barrier(cmd, output.buffer().impl()); + // P2 多输出：额外输出同样过写后屏障 + for (const GpuTensor* eo : extra_outs) + record_output_barrier(cmd, eo->buffer().impl()); - // 查询融合 shader 是否已注册（GpuEngine AOT 匹配前置判断） - [[nodiscard]] bool has_fused_shader(const std::string& shader_name) const noexcept + // 5. 独立模式提交 + if (owns_cmd) + { + auto r = submit_and_wait(cmd, desc_set); + NN_TRY_CHECK(r); + } + if (profile_ops_enabled()) + { + const long long dt = std::chrono::duration_cast( + std::chrono::steady_clock::now() - prof_t0) + .count(); + // fold 按形态如实打印（is_fold = 本形态是否 fold）；out= 运行时 + // 列数（veclen 不进 key，注册侧值可能与调用方实际 cols 不同—— + // 打印注册侧列数会误导归因） + std::fprintf(stderr, "[gpu-profile] fused fold=%d out=%u total=%lldus key=%.16s\n", + static_cast(is_fold), static_cast(cols), dt, + shader_name.c_str()); + } + return output; +} + +// ══════════════════════════════════════════════════════════════════ +// run_fused_gpu_batch — P3 跨链批量派发（键 = key#b） +// +// 同结构（同 key）、同参数值（consts/rparams/vp）、互不依赖的 N 个实例 +// 合成**一次 dispatch**：实例表（uvec2 数组，见 FusedBatchInst 注释）经 +// vkCmdUpdateBuffer 在录制期内联写入一个小 device-local 缓冲，shader 用 +// buffer device address 按槽位解引用各实例的输入/输出缓冲。 +// +// 与生成器（expr_glsl_gen 的 batched 分支）的契约： +// Bindings: 仅 binding 0（实例表，storage buffer） +// Push Constants = V0 逐元素同布局：count(=max_count), cols, [vp..], +// [c..], [rp..]（批量路径不读 count/cols 两槽——元素数 +// 与列数取每实例尾槽，见生成器 `#define cols`） +// dispatch = (ceil(max_count/256), N, 1)，实例下标 = gl_WorkGroupID.y +// +// 前置（调用方 = execute_fused_graph 保证，本函数只防御性校验）： +// `key#b` 已注册、设备 bda_available、实例全 f32、组内互不依赖。 +// ══════════════════════════════════════════════════════════════════ +[[nodiscard]] Result run_fused_gpu_batch(const std::string& shader_name, + std::span instances, + std::span consts, + std::span view_params = {}, + std::span rparams = {}) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + const auto prof_t0 = std::chrono::steady_clock::now(); + if (!device_.bda_available()) + NN_FAIL("run_fused_gpu_batch: 设备未启用 buffer device address" + "（应先逐 kernel 降级，不该到此）"); + if (!std::string_view(shader_name).ends_with(nn::EXPR_BATCH_SUFFIX)) + NN_FAIL("run_fused_gpu_batch: 非批量变体键（应以 #b 结尾）：" + shader_name); + const auto it = fused_kernels_.find(shader_name); + if (it == fused_kernels_.end()) NN_FAIL("fused shader not registered: " + shader_name); + const FusedKernel& kernel = it->second; + const VulkanPipeline& pipeline = kernel.pipeline; + if (instances.empty()) return {}; + // PC 槽位校验（与 run_fused_gpu 同源：vp/rp 计数来自注册表元数据） + const std::uint32_t n_vp = kernel.view_param_count; + const std::uint32_t n_rp = kernel.rparam_count; + if (view_params.size() != n_vp) + NN_FAIL("run_fused_gpu_batch: view_params count mismatch for " + shader_name); + if (rparams.size() != n_rp) + NN_FAIL("run_fused_gpu_batch: rparams count mismatch for " + shader_name); + const std::uint32_t pc_base_fixed = kernel.pc_base_fixed; + + // 实例表布局：stride S = 输入 + 主输出 + extras + (count,cols) 尾槽 + const std::size_t n_in = instances[0].inputs.size(); + const std::size_t n_ex = instances[0].extra_outs.size(); + const std::size_t stride = n_in + 1 + n_ex + 1; + if (kernel.batch_stride != 0 && kernel.batch_stride != static_cast(stride)) + NN_FAIL("run_fused_gpu_batch: 实例槽数与注册结构不一致：" + shader_name); + for (const FusedBatchInst& inst : instances) { - return fused_kernels_.find(shader_name) != fused_kernels_.end(); + if (inst.inputs.size() != n_in || inst.extra_outs.size() != n_ex) + NN_FAIL("run_fused_gpu_batch: 组内实例结构不一致"); + if (inst.out == nullptr) NN_FAIL("run_fused_gpu_batch: 实例缺主输出缓冲"); + for (const GpuBuffer* b : inst.inputs) + if (b == nullptr) NN_FAIL("run_fused_gpu_batch: 实例缺输入缓冲"); + for (const GpuBuffer* b : inst.extra_outs) + if (b == nullptr) NN_FAIL("run_fused_gpu_batch: 实例缺 extras 输出缓冲"); + } + // vkCmdUpdateBuffer 的硬上限：单次 ≤ 64KB、4 字节对齐 + if (instances.size() * stride * sizeof(std::uint64_t) > 65536u) + NN_FAIL("run_fused_gpu_batch: 实例表超过 vkCmdUpdateBuffer 64KB 上限"); + + // ── 实例表内容：uvec2 = (地址低 32 位, 高 32 位)；尾槽 = (count, cols) ── + std::vector table(instances.size() * stride * 2, 0u); + std::uint32_t max_count = 0; + std::uint32_t max_cols = 0; + for (std::size_t i = 0; i < instances.size(); ++i) + { + const FusedBatchInst& inst = instances[i]; + const std::size_t base = i * stride * 2; + bool addr_ok = true; + const auto put_addr = [&](std::size_t slot, const GpuBuffer* b) -> bool + { + VkBufferDeviceAddressInfo ai{}; + ai.sType = VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO; + ai.buffer = b->impl(); + const VkDeviceAddress addr = vkGetBufferDeviceAddress(device_.device(), &ai); + table[base + slot * 2 + 0] = static_cast(addr & 0xffffffffull); + table[base + slot * 2 + 1] = static_cast(addr >> 32); + return addr != 0; + }; + for (std::size_t j = 0; j < n_in; ++j) + addr_ok = put_addr(j, inst.inputs[j]) && addr_ok; + addr_ok = put_addr(n_in, inst.out) && addr_ok; + for (std::size_t k = 0; k < n_ex; ++k) + addr_ok = put_addr(n_in + 1 + k, inst.extra_outs[k]) && addr_ok; + if (!addr_ok) + NN_FAIL("run_fused_gpu_batch: vkGetBufferDeviceAddress 返回 0" + "（buffer 未带 SHADER_DEVICE_ADDRESS？）"); + table[base + (stride - 1) * 2 + 0] = inst.count; + table[base + (stride - 1) * 2 + 1] = inst.cols; + max_count = std::max(max_count, inst.count); + max_cols = std::max(max_cols, inst.cols); } - // P3 批量派发计数（测试/诊断：确认批量路径真的被走到，而非静默降级） - [[nodiscard]] std::uint64_t fused_batch_dispatches() const noexcept - { return fused_batch_dispatches_; } + // 实例表缓冲：小 device-local（TRANSFER_DST 走 vkCmdUpdateBuffer）。 + // 数据内联进命令缓冲 → 录制生命周期安全；缓冲本体析构走 pending_destroys_ + // 延迟队列（铁律 6，与其余录制期资源同款）。 + auto inst_r = GpuBuffer::create_device_local( + device_.device(), alloc_pool(), table.size() * sizeof(std::uint32_t), + VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, keepalive()); + NN_TRY_CHECK(inst_r); + std::shared_ptr inst_buf = std::make_shared(std::move(*inst_r)); + + // 描述符：key#b pipeline 只绑 binding 0（实例表）——复用 per-pipeline + // 描述符机制 + // 批量变体只有一个 binding：实例表 + const VkBuffer instance_table = inst_buf->impl(); + auto ds_r = alloc_desc_set_for(pipeline.descriptor_layout(), {&instance_table, 1}); + NN_TRY_CHECK(ds_r); + VkDescriptorSet desc_set = *ds_r; + + // command buffer + auto cmd_r = acquire_cmd(); + if (!cmd_r) + { + release_desc_set_now(desc_set); + NN_TRY_CHECK(cmd_r); + } + auto [cmd, owns_cmd] = *cmd_r; - // ══════════════════════════════════════════════════════════════════ - // reduce_gpu — 行/列归约原语（sum/max） - // - // Push Constants (20 bytes): rows, cols, mode, reduce_op, chunk_rows - // Bindings: In(0), Out(1) - // mode: 0=row_reduce → out(rows,1), 1=col_reduce → out(1,cols) - // reduce_op: 0=sum, 1=max - // 列归约大行数走两段式 partials（同 cmd 双 dispatch 单次提交，见分支内注释） - // f16_in（Phase C2）：输入缓冲按 float16_t 读（PipelineId::reduce_f16）； - // 输出/ partials 恒 f32 → 两段式 pass2 输入是 f32 partials，恒用 f32 - // pipeline（ds/layout 各随所属 pipeline）。f16 pipeline 未创建时硬报错， - // 由调用方（GpuEngine）提前查 has(PipelineId::reduce_f16) 走引擎内 cast 回退。 - // ══════════════════════════════════════════════════════════════════ - [[nodiscard]] Result reduce_gpu( - const GpuTensor& input, uint32_t mode, uint32_t reduce_op, - bool f16_in = false) - { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - if (!has(PipelineId::reduce)) - NN_FAIL("reduce pipeline not available"); - if (f16_in && !has(PipelineId::reduce_f16)) - NN_FAIL("reduce f16 pipeline not available"); - // 输入侧 pipeline（pass1 / 单段）；输出侧与 pass2 恒 f32 pipeline - const VulkanPipeline& in_pipe = f16_in ? pipelines_.get(PipelineId::reduce_f16) : pipelines_.get(PipelineId::reduce); - - const uint32_t rows = static_cast(input.rows()); - const uint32_t cols = static_cast(input.cols()); - - // 1. 分配输出 Tensor - std::size_t out_rows = (mode == 0) ? rows : 1; - std::size_t out_cols = (mode == 0) ? 1 : cols; - auto output_res = GpuTensor::create_empty(out_rows, out_cols, *this); - NN_TRY_CHECK(output_res); - GpuTensor output = std::move(*output_res); - - // ── 列归约大行数：两段式 partials（ggml rms_norm_partials 模式）── - // 单段 dispatch = ceil(cols/32) 个 WG（5244² 仅 164 vs 136 并发槽 → - // 1.2 波，末波只剩 ~20% 机器喂内存 → 带宽被波尾压约两成）。 - // pass1 把行向切 nchunk 块（dispatch y = nchunk）→ WG 数 ×nchunk， - // 每 WG 归约一行块写 partials(nchunk, cols)；pass2 复用整表路径合并 - // （行数 = nchunk ≤ 8，成本可忽略）。两 pass 录进同一 cmd 单次提交 - // → 固定提交开销不翻倍。小行数保持单段（零开销）。 - uint32_t nchunk = 1u; - if (mode == 1u) - { - const uint32_t tiles = (cols + 31u) / 32u; - nchunk = (rows + 1023u) / 1024u; - if (nchunk > 8u) nchunk = 8u; - const uint32_t want = (512u + tiles - 1u) / tiles; // ≥512 WG ≈ 4 波 - if (want > nchunk) nchunk = want; - if (nchunk > 8u) nchunk = 8u; - // 小形状回落单段：512² 实测两段 −16%（第二遍+屏障对 - // ~10µs kernel 净亏），1024² 起两段持平或转优 → 边界 512K 元素 - if (rows < 256u || - static_cast(rows) * cols < 524'288ull) - nchunk = 1u; - } - // 两段式启用条件（同窗交错 A/B 实测，同码对照噪声地板 ±5%）： - // 5244² 两段 −6%（0.405 vs 0.430；单段 1.2 波尾欠喂被消除） - // 4096² 两段 ~2%（0.277 vs 0.283；单波本已喂满，第二段纯付 ~6µs） - // 512² 两段 −16%（对 ~10µs kernel 第二遍+屏障净亏）→ 尺寸护栏只排除它 - const bool two_pass = (nchunk > 1u); - const uint32_t chunk_rows = two_pass ? (rows + nchunk - 1u) / nchunk : 0u; - - if (two_pass) + // 录制：实例表写入 → 输入屏障（含实例表，去重）→ bind → PC → dispatch + vkCmdUpdateBuffer(cmd, inst_buf->impl(), 0, + static_cast(table.size() * sizeof(std::uint32_t)), + table.data()); + { + // 屏障覆盖组内全部实例的输入 buffer + 实例表（vkCmdUpdateBuffer 是 + // TRANSFER 写，record_input_barriers 的 src 阶段含 TRANSFER）； + // 同一 buffer 去重（组内互不依赖已由引擎保证，仅防重复屏障）。 + std::vector in_bufs; + in_bufs.reserve(instances.size() * (n_in + 1) + 1); + for (const FusedBatchInst& inst : instances) + for (const GpuBuffer* b : inst.inputs) + in_bufs.push_back(b->impl()); + in_bufs.push_back(inst_buf->impl()); + std::sort(in_bufs.begin(), in_bufs.end()); + in_bufs.erase(std::unique(in_bufs.begin(), in_bufs.end()), in_bufs.end()); + record_input_barriers(cmd, in_bufs); + } + vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline.handle()); + vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline.pipeline_layout(), 0, 1, + &desc_set, 0, nullptr); + + // Push constants 与 run_fused_gpu 的 V0 逐元素分支同源： + // count(=max_count), cols, [vp0..], [c0..], [rp0..] + // 批量 shader 不读 count/cols 两槽（每实例取尾槽），两值仅保持 V0 布局。 + const std::uint32_t pc_uints = pc_base_fixed + n_vp; + std::vector pc((pc_uints) * sizeof(std::uint32_t) + + sizeof(Scalar) * consts.size() + sizeof(Scalar) * rparams.size()); + std::memcpy(pc.data(), &max_count, sizeof(std::uint32_t)); + std::memcpy(pc.data() + sizeof(std::uint32_t), &max_cols, sizeof(std::uint32_t)); + for (std::uint32_t i = 0; i < n_vp; ++i) + std::memcpy(pc.data() + (pc_base_fixed + i) * sizeof(std::uint32_t), &view_params[i], + sizeof(std::uint32_t)); + if (!consts.empty()) + std::memcpy(pc.data() + (pc_uints) * sizeof(std::uint32_t), consts.data(), + sizeof(Scalar) * consts.size()); + if (!rparams.empty()) + std::memcpy(pc.data() + (pc_uints) * sizeof(std::uint32_t) + sizeof(Scalar) * consts.size(), + rparams.data(), sizeof(Scalar) * rparams.size()); + vkCmdPushConstants(cmd, pipeline.pipeline_layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, + static_cast(pc.size()), pc.data()); + + // dispatch = (ceil(max_count/256), n_instances, 1)；批量恒标量 + // （vec_width=1，与生成器同源） + const std::uint32_t wg_x = (max_count + 255u) / 256u; + ts_dispatch(cmd, "fused_b", wg_x, static_cast(instances.size()), 1, shader_name); + + // 输出屏障（主输出 + 全部 extras，去重）：组内输出互不重叠（引擎保证） + { + std::vector out_bufs; + out_bufs.reserve(instances.size() * (1 + n_ex)); + for (const FusedBatchInst& inst : instances) { - NN_TRY_CHECK(ensure_reduce_partial_(nchunk, cols)); - - // pass1 读输入（f16_in → f16 pipeline，ds/layout 随之）； - // pass2 读 f32 partials → 恒 f32 pipeline（f16 shader 读 f32 - // partials = 静默错值，必须分开） - const VkBuffer pass1_bindings[2]{ - input.buffer().impl(), reduce_partial_.buffer}; - const VkBuffer pass2_bindings[2]{ - reduce_partial_.buffer, output.buffer().impl()}; - auto ds1_r = alloc_desc_set_for(in_pipe.descriptor_layout(), pass1_bindings); - NN_TRY_CHECK(ds1_r); - VkDescriptorSet ds1 = *ds1_r; - auto ds2_r = alloc_desc_set_for( - pipelines_.get(PipelineId::reduce).descriptor_layout(), pass2_bindings); - NN_TRY_CHECK(ds2_r); - VkDescriptorSet ds2 = *ds2_r; - - auto cmd_r = acquire_cmd(); - if (!cmd_r) - { - release_desc_set_now(ds1); - release_desc_set_now(ds2); - NN_TRY_CHECK(cmd_r); - } - auto [cmd, owns_cmd] = *cmd_r; - - // pass1：行块级部分归约 → partials（读输入 → in_pipe） - record_input_barriers(cmd, {input.buffer().impl()}); - vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - in_pipe.handle()); - vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - in_pipe.pipeline_layout(), 0, 1, &ds1, 0, nullptr); - const uint32_t push1[5] = {rows, cols, 1u, reduce_op, chunk_rows}; - vkCmdPushConstants(cmd, in_pipe.pipeline_layout(), - VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(push1), push1); - ts_dispatch(cmd, "reduce", (cols + 31u) / 32u, nchunk, 1); - - // partials：写 → 读（同 cmd 内跨 dispatch 依赖） - record_input_barriers(cmd, {reduce_partial_.buffer}); - - // pass2：合并 partials（rows = nchunk ≤ 8，复用 mode1 整表路径； - // 输入 = f32 partials → 恒 f32 pipeline） - vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipelines_.get(PipelineId::reduce).handle()); - vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipelines_.get(PipelineId::reduce).pipeline_layout(), 0, 1, &ds2, 0, nullptr); - const uint32_t push2[5] = {nchunk, cols, 1u, reduce_op, 0u}; - vkCmdPushConstants(cmd, pipelines_.get(PipelineId::reduce).pipeline_layout(), - VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(push2), push2); - ts_dispatch(cmd, "reduce", (cols + 31u) / 32u, 1, 1); - record_output_barrier(cmd, output.buffer().impl()); - - if (owns_cmd) - { - auto r = submit_and_wait(cmd, ds1); // 等 fence 后归还 ds1 - NN_TRY_CHECK(r); - vkFreeDescriptorSets(device_.device(), gpu_tensor_pool_, 1, &ds2); - } - return output; + out_bufs.push_back(inst.out->impl()); + for (const GpuBuffer* b : inst.extra_outs) + out_bufs.push_back(b->impl()); } + std::sort(out_bufs.begin(), out_bufs.end()); + out_bufs.erase(std::unique(out_bufs.begin(), out_bufs.end()), out_bufs.end()); + for (VkBuffer b : out_bufs) + record_output_barrier(cmd, b); + } - // 2. 分配描述符集（随 in_pipe：f16_in 时输入按 float16_t 读） - const VkBuffer bindings[2]{ - input.buffer().impl(), output.buffer().impl()}; - auto ds_r = alloc_desc_set_for(in_pipe.descriptor_layout(), bindings); - NN_TRY_CHECK(ds_r); - VkDescriptorSet desc_set = *ds_r; + if (owns_cmd) + { + auto r = submit_and_wait(cmd, desc_set); + NN_TRY_CHECK(r); + } + ++fused_batch_dispatches_; + if (profile_ops_enabled()) + { + const long long dt = std::chrono::duration_cast( + std::chrono::steady_clock::now() - prof_t0) + .count(); + std::fprintf(stderr, + "[gpu-profile] fused_batch inst=%zu max_count=%u total=%lldus " + "key=%.16s\n", + instances.size(), max_count, dt, shader_name.c_str()); + } + return {}; +} + +// 查询融合 shader 是否已注册（GpuEngine AOT 匹配前置判断） +[[nodiscard]] bool has_fused_shader(const std::string& shader_name) const noexcept +{ return fused_kernels_.find(shader_name) != fused_kernels_.end(); } + +// P3 批量派发计数（测试/诊断：确认批量路径真的被走到，而非静默降级） +[[nodiscard]] std::uint64_t fused_batch_dispatches() const noexcept +{ return fused_batch_dispatches_; } + +// ══════════════════════════════════════════════════════════════════ +// reduce_gpu — 行/列归约原语（sum/max） +// +// Push Constants (20 bytes): rows, cols, mode, reduce_op, chunk_rows +// Bindings: In(0), Out(1) +// mode: 0=row_reduce → out(rows,1), 1=col_reduce → out(1,cols) +// reduce_op: 0=sum, 1=max +// 列归约大行数走两段式 partials（同 cmd 双 dispatch 单次提交，见分支内注释） +// f16_in（Phase C2）：输入缓冲按 float16_t 读（PipelineId::reduce_f16）； +// 输出/ partials 恒 f32 → 两段式 pass2 输入是 f32 partials，恒用 f32 +// pipeline（ds/layout 各随所属 pipeline）。f16 pipeline 未创建时硬报错， +// 由调用方（GpuEngine）提前查 has(PipelineId::reduce_f16) 走引擎内 cast 回退。 +// ══════════════════════════════════════════════════════════════════ +[[nodiscard]] Result reduce_gpu(const GpuTensor& input, uint32_t mode, + uint32_t reduce_op, bool f16_in = false) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + if (!has(PipelineId::reduce)) NN_FAIL("reduce pipeline not available"); + if (f16_in && !has(PipelineId::reduce_f16)) NN_FAIL("reduce f16 pipeline not available"); + // 输入侧 pipeline（pass1 / 单段）；输出侧与 pass2 恒 f32 pipeline + const VulkanPipeline& in_pipe = + f16_in ? pipelines_.get(PipelineId::reduce_f16) : pipelines_.get(PipelineId::reduce); + + const uint32_t rows = static_cast(input.rows()); + const uint32_t cols = static_cast(input.cols()); + + // 1. 分配输出 Tensor + std::size_t out_rows = (mode == 0) ? rows : 1; + std::size_t out_cols = (mode == 0) ? 1 : cols; + auto output_res = GpuTensor::create_empty(out_rows, out_cols, *this); + NN_TRY_CHECK(output_res); + GpuTensor output = std::move(*output_res); + + // ── 列归约大行数：两段式 partials（ggml rms_norm_partials 模式）── + // 单段 dispatch = ceil(cols/32) 个 WG（5244² 仅 164 vs 136 并发槽 → + // 1.2 波，末波只剩 ~20% 机器喂内存 → 带宽被波尾压约两成）。 + // pass1 把行向切 nchunk 块（dispatch y = nchunk）→ WG 数 ×nchunk， + // 每 WG 归约一行块写 partials(nchunk, cols)；pass2 复用整表路径合并 + // （行数 = nchunk ≤ 8，成本可忽略）。两 pass 录进同一 cmd 单次提交 + // → 固定提交开销不翻倍。小行数保持单段（零开销）。 + uint32_t nchunk = 1u; + if (mode == 1u) + { + const uint32_t tiles = (cols + 31u) / 32u; + nchunk = (rows + 1023u) / 1024u; + if (nchunk > 8u) nchunk = 8u; + const uint32_t want = (512u + tiles - 1u) / tiles; // ≥512 WG ≈ 4 波 + if (want > nchunk) nchunk = want; + if (nchunk > 8u) nchunk = 8u; + // 小形状回落单段：512² 实测两段 −16%（第二遍+屏障对 + // ~10µs kernel 净亏），1024² 起两段持平或转优 → 边界 512K 元素 + if (rows < 256u || static_cast(rows) * cols < 524'288ull) nchunk = 1u; + } + // 两段式启用条件（同窗交错 A/B 实测，同码对照噪声地板 ±5%）： + // 5244² 两段 −6%（0.405 vs 0.430；单段 1.2 波尾欠喂被消除） + // 4096² 两段 ~2%（0.277 vs 0.283；单波本已喂满，第二段纯付 ~6µs） + // 512² 两段 −16%（对 ~10µs kernel 第二遍+屏障净亏）→ 尺寸护栏只排除它 + const bool two_pass = (nchunk > 1u); + const uint32_t chunk_rows = two_pass ? (rows + nchunk - 1u) / nchunk : 0u; + + if (two_pass) + { + NN_TRY_CHECK(ensure_reduce_partial_(nchunk, cols)); + + // pass1 读输入（f16_in → f16 pipeline，ds/layout 随之）； + // pass2 读 f32 partials → 恒 f32 pipeline（f16 shader 读 f32 + // partials = 静默错值，必须分开） + const VkBuffer pass1_bindings[2]{input.buffer().impl(), reduce_partial_.buffer}; + const VkBuffer pass2_bindings[2]{reduce_partial_.buffer, output.buffer().impl()}; + auto ds1_r = alloc_desc_set_for(in_pipe.descriptor_layout(), pass1_bindings); + NN_TRY_CHECK(ds1_r); + VkDescriptorSet ds1 = *ds1_r; + auto ds2_r = alloc_desc_set_for(pipelines_.get(PipelineId::reduce).descriptor_layout(), + pass2_bindings); + NN_TRY_CHECK(ds2_r); + VkDescriptorSet ds2 = *ds2_r; - // 4. 获取 command buffer auto cmd_r = acquire_cmd(); if (!cmd_r) { - release_desc_set_now(desc_set); + release_desc_set_now(ds1); + release_desc_set_now(ds2); NN_TRY_CHECK(cmd_r); } auto [cmd, owns_cmd] = *cmd_r; - // 5. 录制 + // pass1：行块级部分归约 → partials（读输入 → in_pipe） record_input_barriers(cmd, {input.buffer().impl()}); + vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, in_pipe.handle()); + vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, in_pipe.pipeline_layout(), 0, + 1, &ds1, 0, nullptr); + const uint32_t push1[5] = {rows, cols, 1u, reduce_op, chunk_rows}; + vkCmdPushConstants(cmd, in_pipe.pipeline_layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, + sizeof(push1), push1); + ts_dispatch(cmd, "reduce", (cols + 31u) / 32u, nchunk, 1); + + // partials：写 → 读（同 cmd 内跨 dispatch 依赖） + record_input_barriers(cmd, {reduce_partial_.buffer}); + + // pass2：合并 partials（rows = nchunk ≤ 8，复用 mode1 整表路径； + // 输入 = f32 partials → 恒 f32 pipeline） vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - in_pipe.handle()); + pipelines_.get(PipelineId::reduce).handle()); vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - in_pipe.pipeline_layout(), 0, 1, &desc_set, 0, nullptr); - - const uint32_t push_data[5] = {rows, cols, mode, reduce_op, 0u}; - vkCmdPushConstants(cmd, in_pipe.pipeline_layout(), - VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(push_data), push_data); - - // 行：每 WG 256 线程协作归约一行 → rows 个 WG； - // 列：每 WG 32 列 tile（lane=列、warp=行块，合并访问；与 - // reduce.comp 的 WG 结构契约同源）→ (cols+31)/32 个 WG - const uint32_t workgroups = (mode == 0) ? rows : (cols + 31u) / 32u; - ts_dispatch(cmd, "reduce", workgroups, 1, 1); + pipelines_.get(PipelineId::reduce).pipeline_layout(), 0, 1, &ds2, 0, + nullptr); + const uint32_t push2[5] = {nchunk, cols, 1u, reduce_op, 0u}; + vkCmdPushConstants(cmd, pipelines_.get(PipelineId::reduce).pipeline_layout(), + VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(push2), push2); + ts_dispatch(cmd, "reduce", (cols + 31u) / 32u, 1, 1); record_output_barrier(cmd, output.buffer().impl()); - // 6. 独立模式提交 if (owns_cmd) { - auto r = submit_and_wait(cmd, desc_set); + auto r = submit_and_wait(cmd, ds1); // 等 fence 后归还 ds1 NN_TRY_CHECK(r); + vkFreeDescriptorSets(device_.device(), gpu_tensor_pool_, 1, &ds2); } return output; } - // ══════════════════════════════════════════════════════════════════ - // rearrange_3d_gpu — 3D 维度转置 (M, B, N) ↔ (B, M, N) - // - // Push Constants (20 bytes): M, B, N, inverse, total - // Bindings: In(0), Out(1) - // inverse=0: (M, B*N) → (B*M, N) - // inverse=1: (B*M, N) → (M, B*N) - // ══════════════════════════════════════════════════════════════════ - [[nodiscard]] Result rearrange_3d_gpu( - const GpuTensor& input, - uint32_t M, uint32_t B, uint32_t N, uint32_t inverse, - bool f16_io = false) - { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - if (!has(PipelineId::rearrange_3d)) - NN_FAIL("rearrange_3d pipeline not available"); - if (f16_io && !has(PipelineId::rearrange_3d_f16)) - NN_FAIL("rearrange_3d_gpu: f16 pipeline 不可用（设备无16位存储）"); - const auto& pipeline = f16_io ? pipelines_.get(PipelineId::rearrange_3d_f16) - : pipelines_.get(PipelineId::rearrange_3d); - - const uint32_t total = M * B * N; - if (total != input.rows() * input.cols()) - NN_FAIL("rearrange_3d: element count mismatch"); - // 1. 分配输出 Tensor（f16：2B/元素 + 纯绑定视图，同 matmul_gpu） - const std::size_t out_rows = inverse ? M : (static_cast(B) * M); - const std::size_t out_cols = inverse ? (static_cast(B) * N) : N; - std::optional owned; - if (f16_io) - { - auto f16r = GpuTensorF16::create_empty(out_rows, out_cols, *this); - NN_TRY_CHECK(f16r); - owned.emplace(f16r->shared_buffer(), f16r->rows(), f16r->cols()); - } - else - { - auto output_res = GpuTensor::create_empty(out_rows, out_cols, *this); - NN_TRY_CHECK(output_res); - owned.emplace(std::move(*output_res)); - } - GpuTensor output = *owned; + // 2. 分配描述符集（随 in_pipe：f16_in 时输入按 float16_t 读） + const VkBuffer bindings[2]{input.buffer().impl(), output.buffer().impl()}; + auto ds_r = alloc_desc_set_for(in_pipe.descriptor_layout(), bindings); + NN_TRY_CHECK(ds_r); + VkDescriptorSet desc_set = *ds_r; - // 2. 分配并写入描述符集 - const VkBuffer bindings[2]{ - input.buffer().impl(), output.buffer().impl()}; - auto ds_r = alloc_desc_set_for(pipeline.descriptor_layout(), bindings); - NN_TRY_CHECK(ds_r); - VkDescriptorSet desc_set = *ds_r; + // 4. 获取 command buffer + auto cmd_r = acquire_cmd(); + if (!cmd_r) + { + release_desc_set_now(desc_set); + NN_TRY_CHECK(cmd_r); + } + auto [cmd, owns_cmd] = *cmd_r; + + // 5. 录制 + record_input_barriers(cmd, {input.buffer().impl()}); + vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, in_pipe.handle()); + vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, in_pipe.pipeline_layout(), 0, 1, + &desc_set, 0, nullptr); + + const uint32_t push_data[5] = {rows, cols, mode, reduce_op, 0u}; + vkCmdPushConstants(cmd, in_pipe.pipeline_layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, + sizeof(push_data), push_data); + + // 行：每 WG 256 线程协作归约一行 → rows 个 WG； + // 列：每 WG 32 列 tile（lane=列、warp=行块，合并访问；与 + // reduce.comp 的 WG 结构契约同源）→ (cols+31)/32 个 WG + const uint32_t workgroups = (mode == 0) ? rows : (cols + 31u) / 32u; + ts_dispatch(cmd, "reduce", workgroups, 1, 1); + record_output_barrier(cmd, output.buffer().impl()); + + // 6. 独立模式提交 + if (owns_cmd) + { + auto r = submit_and_wait(cmd, desc_set); + NN_TRY_CHECK(r); + } + return output; +} + +// ══════════════════════════════════════════════════════════════════ +// rearrange_3d_gpu — 3D 维度转置 (M, B, N) ↔ (B, M, N) +// +// Push Constants (20 bytes): M, B, N, inverse, total +// Bindings: In(0), Out(1) +// inverse=0: (M, B*N) → (B*M, N) +// inverse=1: (B*M, N) → (M, B*N) +// ══════════════════════════════════════════════════════════════════ +[[nodiscard]] Result rearrange_3d_gpu(const GpuTensor& input, uint32_t M, uint32_t B, + uint32_t N, uint32_t inverse, bool f16_io = false) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + if (!has(PipelineId::rearrange_3d)) NN_FAIL("rearrange_3d pipeline not available"); + if (f16_io && !has(PipelineId::rearrange_3d_f16)) + NN_FAIL("rearrange_3d_gpu: f16 pipeline 不可用（设备无16位存储）"); + const auto& pipeline = f16_io ? pipelines_.get(PipelineId::rearrange_3d_f16) + : pipelines_.get(PipelineId::rearrange_3d); + + const uint32_t total = M * B * N; + if (total != input.rows() * input.cols()) NN_FAIL("rearrange_3d: element count mismatch"); + // 1. 分配输出 Tensor（f16：2B/元素 + 纯绑定视图，同 matmul_gpu） + const std::size_t out_rows = inverse ? M : (static_cast(B) * M); + const std::size_t out_cols = inverse ? (static_cast(B) * N) : N; + std::optional owned; + if (f16_io) + { + auto f16r = GpuTensorF16::create_empty(out_rows, out_cols, *this); + NN_TRY_CHECK(f16r); + owned.emplace(f16r->shared_buffer(), f16r->rows(), f16r->cols()); + } + else + { + auto output_res = GpuTensor::create_empty(out_rows, out_cols, *this); + NN_TRY_CHECK(output_res); + owned.emplace(std::move(*output_res)); + } + GpuTensor output = *owned; - // 4. 获取 command buffer - auto cmd_r = acquire_cmd(); - if (!cmd_r) - { - release_desc_set_now(desc_set); - NN_TRY_CHECK(cmd_r); - } - auto [cmd, owns_cmd] = *cmd_r; + // 2. 分配并写入描述符集 + const VkBuffer bindings[2]{input.buffer().impl(), output.buffer().impl()}; + auto ds_r = alloc_desc_set_for(pipeline.descriptor_layout(), bindings); + NN_TRY_CHECK(ds_r); + VkDescriptorSet desc_set = *ds_r; - // 5. 录制 - record_input_barriers(cmd, {input.buffer().impl()}); - vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipeline.handle()); - vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipeline.pipeline_layout(), 0, 1, &desc_set, 0, nullptr); + // 4. 获取 command buffer + auto cmd_r = acquire_cmd(); + if (!cmd_r) + { + release_desc_set_now(desc_set); + NN_TRY_CHECK(cmd_r); + } + auto [cmd, owns_cmd] = *cmd_r; - const uint32_t push_data[5] = {M, B, N, inverse, total}; - vkCmdPushConstants(cmd, pipeline.pipeline_layout(), - VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(push_data), push_data); + // 5. 录制 + record_input_barriers(cmd, {input.buffer().impl()}); + vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline.handle()); + vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline.pipeline_layout(), 0, 1, + &desc_set, 0, nullptr); - const uint32_t wg_count = (total + 255) / 256; - ts_dispatch(cmd, "rearrange_3d", wg_count, 1, 1); - record_output_barrier(cmd, output.buffer().impl()); + const uint32_t push_data[5] = {M, B, N, inverse, total}; + vkCmdPushConstants(cmd, pipeline.pipeline_layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, + sizeof(push_data), push_data); - // 6. 独立模式提交 - if (owns_cmd) - { - auto r = submit_and_wait(cmd, desc_set); - NN_TRY_CHECK(r); - } - return output; - } + const uint32_t wg_count = (total + 255) / 256; + ts_dispatch(cmd, "rearrange_3d", wg_count, 1, 1); + record_output_barrier(cmd, output.buffer().impl()); - // ══════════════════════════════════════════════════════════════════ - // fill_zero_gpu — 清零 GPU buffer（使用 vkCmdFillBuffer，字节级操作，§6.3） - // ══════════════════════════════════════════════════════════════════ - template - [[nodiscard]] Result fill_zero_gpu(GpuTensorT

& tensor) + // 6. 独立模式提交 + if (owns_cmd) { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - auto cmd_r = acquire_cmd(); - NN_TRY_CHECK(cmd_r); - auto [cmd, owns_cmd] = *cmd_r; - - vkCmdFillBuffer(cmd, tensor.buffer().impl(), 0, VK_WHOLE_SIZE, 0); - - // 屏障确保 fill 完成后可见 - VkBufferMemoryBarrier b{}; - b.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER; - b.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; - b.dstAccessMask = VK_ACCESS_MEMORY_READ_BIT; - b.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - b.buffer = tensor.buffer().impl(); - b.offset = 0; - b.size = VK_WHOLE_SIZE; - vkCmdPipelineBarrier(cmd, - VK_PIPELINE_STAGE_TRANSFER_BIT, - VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, - 0, 0, nullptr, 1, &b, 0, nullptr); - - if (owns_cmd) - { - // fill_zero 不使用描述符集，传 VK_NULL_HANDLE 跳过释放 - auto r = submit_and_wait(cmd, VK_NULL_HANDLE); - NN_TRY_CHECK(r); - } - return {}; + auto r = submit_and_wait(cmd, desc_set); + NN_TRY_CHECK(r); } - - // ══════════════════════════════════════════════════════════════════ - // copy_buffer_gpu — GPU 内 buffer 拷贝（使用 vkCmdCopyBuffer） - // ══════════════════════════════════════════════════════════════════ - [[nodiscard]] Result copy_buffer_gpu( - VkBuffer src, VkBuffer dst, VkDeviceSize size) + return output; +} + +// ══════════════════════════════════════════════════════════════════ +// fill_zero_gpu — 清零 GPU buffer（使用 vkCmdFillBuffer，字节级操作，§6.3） +// ══════════════════════════════════════════════════════════════════ +template [[nodiscard]] Result fill_zero_gpu(GpuTensorT

& tensor) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + auto cmd_r = acquire_cmd(); + NN_TRY_CHECK(cmd_r); + auto [cmd, owns_cmd] = *cmd_r; + + vkCmdFillBuffer(cmd, tensor.buffer().impl(), 0, VK_WHOLE_SIZE, 0); + + // 屏障确保 fill 完成后可见 + VkBufferMemoryBarrier b{}; + b.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER; + b.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; + b.dstAccessMask = VK_ACCESS_MEMORY_READ_BIT; + b.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; + b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; + b.buffer = tensor.buffer().impl(); + b.offset = 0; + b.size = VK_WHOLE_SIZE; + vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_TRANSFER_BIT, + VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, 0, + 0, nullptr, 1, &b, 0, nullptr); + + if (owns_cmd) { - return copy_buffer_region_gpu(src, 0, dst, 0, size); + // fill_zero 不使用描述符集，传 VK_NULL_HANDLE 跳过释放 + auto r = submit_and_wait(cmd, VK_NULL_HANDLE); + NN_TRY_CHECK(r); } - - // ══════════════════════════════════════════════════════════════════ - // copy_buffer_region_gpu — 带偏移的 buffer 拷贝（offload slab / 通用区域拷贝） - // 从 src[src_offset, src_offset+size) 拷到 dst[dst_offset, dst_offset+size)。 - // batch 模式只录制不提交；非 batch 独立提交等待。 - // ══════════════════════════════════════════════════════════════════ - [[nodiscard]] Result copy_buffer_region_gpu( - VkBuffer src, VkDeviceSize src_offset, - VkBuffer dst, VkDeviceSize dst_offset, - VkDeviceSize size) + return {}; +} + +// ══════════════════════════════════════════════════════════════════ +// copy_buffer_gpu — GPU 内 buffer 拷贝（使用 vkCmdCopyBuffer） +// ══════════════════════════════════════════════════════════════════ +[[nodiscard]] Result copy_buffer_gpu(VkBuffer src, VkBuffer dst, VkDeviceSize size) +{ return copy_buffer_region_gpu(src, 0, dst, 0, size); } + +// ══════════════════════════════════════════════════════════════════ +// copy_buffer_region_gpu — 带偏移的 buffer 拷贝（offload slab / 通用区域拷贝） +// 从 src[src_offset, src_offset+size) 拷到 dst[dst_offset, dst_offset+size)。 +// batch 模式只录制不提交；非 batch 独立提交等待。 +// ══════════════════════════════════════════════════════════════════ +[[nodiscard]] Result copy_buffer_region_gpu(VkBuffer src, VkDeviceSize src_offset, + VkBuffer dst, VkDeviceSize dst_offset, + VkDeviceSize size) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + auto cmd_r = acquire_cmd(); + NN_TRY_CHECK(cmd_r); + auto [cmd, owns_cmd] = *cmd_r; + + VkBufferCopy cp{src_offset, dst_offset, size}; + vkCmdCopyBuffer(cmd, src, dst, 1, &cp); + + VkBufferMemoryBarrier b{}; + b.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER; + b.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; + b.dstAccessMask = VK_ACCESS_MEMORY_READ_BIT; + b.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; + b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; + b.buffer = dst; + b.offset = 0; + b.size = VK_WHOLE_SIZE; + vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_TRANSFER_BIT, + VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, 0, + 0, nullptr, 1, &b, 0, nullptr); + + if (owns_cmd) { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - auto cmd_r = acquire_cmd(); - NN_TRY_CHECK(cmd_r); - auto [cmd, owns_cmd] = *cmd_r; - - VkBufferCopy cp{src_offset, dst_offset, size}; - vkCmdCopyBuffer(cmd, src, dst, 1, &cp); - - VkBufferMemoryBarrier b{}; - b.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER; - b.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; - b.dstAccessMask = VK_ACCESS_MEMORY_READ_BIT; - b.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - b.buffer = dst; - b.offset = 0; - b.size = VK_WHOLE_SIZE; - vkCmdPipelineBarrier(cmd, - VK_PIPELINE_STAGE_TRANSFER_BIT, - VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, - 0, 0, nullptr, 1, &b, 0, nullptr); - - if (owns_cmd) - { - auto r = submit_and_wait(cmd, VK_NULL_HANDLE); - NN_TRY_CHECK(r); - } - return {}; + auto r = submit_and_wait(cmd, VK_NULL_HANDLE); + NN_TRY_CHECK(r); } + return {}; +} + +// ══════════════════════════════════════════════════════════════════ +// 异步标量回读—— 语义见成员声明处 +// ══════════════════════════════════════════════════════════════════ +[[nodiscard]] Result init_scalar_readback_slots() +{ + rb_slots_.resize(SCALAR_READBACK_SLOTS); + for (auto& s : rb_slots_) + { + VkCommandBufferAllocateInfo ca{}; + ca.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO; + ca.commandPool = command_pool_; + ca.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; + ca.commandBufferCount = 1; + auto r = detail::vk_check(vkAllocateCommandBuffers(device_.device(), &ca, &s.cmd), __FILE__, + __LINE__); + NN_TRY_CHECK(r); + VkFenceCreateInfo fi{}; + fi.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO; + r = detail::vk_check(vkCreateFence(device_.device(), &fi, nullptr, &s.fence), __FILE__, + __LINE__); + NN_TRY_CHECK(r); + VkBufferCreateInfo bi{}; + bi.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO; + bi.size = sizeof(float); + bi.usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT; + // P3：同其余缓冲创建点补 SHADER_DEVICE_ADDRESS（池分配侧配对） + if (memory_pool().device_address_alloc()) + bi.usage |= VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT; + bi.sharingMode = VK_SHARING_MODE_EXCLUSIVE; + r = detail::vk_check(vkCreateBuffer(device_.device(), &bi, nullptr, &s.host), __FILE__, + __LINE__); + NN_TRY_CHECK(r); + VkMemoryRequirements mr{}; + vkGetBufferMemoryRequirements(device_.device(), s.host, &mr); + constexpr VkMemoryPropertyFlags kVis = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT; + constexpr VkMemoryPropertyFlags kCoh = VK_MEMORY_PROPERTY_HOST_COHERENT_BIT; + auto ar = memory_pool().allocate(mr, kVis | kCoh, kVis); + NN_TRY_CHECK(ar); + s.alloc = *ar; + r = detail::vk_check( + vkBindBufferMemory(device_.device(), s.host, s.alloc.memory, s.alloc.offset), __FILE__, + __LINE__); + NN_TRY_CHECK(r); + r = detail::vk_check(vkMapMemory(device_.device(), s.alloc.memory, s.alloc.offset, + sizeof(float), 0, &s.mapped), + __FILE__, __LINE__); + NN_TRY_CHECK(r); + } + return {}; +} - // ══════════════════════════════════════════════════════════════════ - // 异步标量回读—— 语义见成员声明处 - // ══════════════════════════════════════════════════════════════════ - [[nodiscard]] Result init_scalar_readback_slots() +void destroy_scalar_readback_slots() +{ + for (auto& s : rb_slots_) { - rb_slots_.resize(SCALAR_READBACK_SLOTS); - for (auto& s : rb_slots_) + if (s.mapped != nullptr && s.alloc.valid()) { - VkCommandBufferAllocateInfo ca{}; - ca.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO; - ca.commandPool = command_pool_; - ca.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; - ca.commandBufferCount = 1; - auto r = detail::vk_check( - vkAllocateCommandBuffers(device_.device(), &ca, &s.cmd), __FILE__, __LINE__); - NN_TRY_CHECK(r); - VkFenceCreateInfo fi{}; - fi.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO; - r = detail::vk_check( - vkCreateFence(device_.device(), &fi, nullptr, &s.fence), __FILE__, __LINE__); - NN_TRY_CHECK(r); - VkBufferCreateInfo bi{}; - bi.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO; - bi.size = sizeof(float); - bi.usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT; - // P3：同其余缓冲创建点补 SHADER_DEVICE_ADDRESS（池分配侧配对） - if (memory_pool().device_address_alloc()) - bi.usage |= VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT; - bi.sharingMode = VK_SHARING_MODE_EXCLUSIVE; - r = detail::vk_check( - vkCreateBuffer(device_.device(), &bi, nullptr, &s.host), __FILE__, __LINE__); - NN_TRY_CHECK(r); - VkMemoryRequirements mr{}; - vkGetBufferMemoryRequirements(device_.device(), s.host, &mr); - constexpr VkMemoryPropertyFlags kVis = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT; - constexpr VkMemoryPropertyFlags kCoh = VK_MEMORY_PROPERTY_HOST_COHERENT_BIT; - auto ar = memory_pool().allocate(mr, kVis | kCoh, kVis); - NN_TRY_CHECK(ar); - s.alloc = *ar; - r = detail::vk_check( - vkBindBufferMemory(device_.device(), s.host, s.alloc.memory, s.alloc.offset), - __FILE__, __LINE__); - NN_TRY_CHECK(r); - r = detail::vk_check( - vkMapMemory(device_.device(), s.alloc.memory, s.alloc.offset, - sizeof(float), 0, &s.mapped), - __FILE__, __LINE__); - NN_TRY_CHECK(r); + vkUnmapMemory(device_.device(), s.alloc.memory); + s.mapped = nullptr; } - return {}; + if (s.host != VK_NULL_HANDLE) vkDestroyBuffer(device_.device(), s.host, nullptr); + if (s.fence != VK_NULL_HANDLE) vkDestroyFence(device_.device(), s.fence, nullptr); + if (s.alloc.valid()) memory_pool().free(s.alloc); + s.host = VK_NULL_HANDLE; + s.fence = VK_NULL_HANDLE; + s.cmd = VK_NULL_HANDLE; // command buffer 随 command pool 释放 + s.pending = false; } - - void destroy_scalar_readback_slots() + rb_slots_.clear(); +} + +// 录制一次 4B D2H 拷贝并提交（**不等待**）。 +// 调用约定：必须在产出该标量的主帧已提交之后调用（同队列 FIFO ⇒ +// 拷贝排在生产者之后；若在主帧提交前提交，会读到上一轮的旧值）。 +[[nodiscard]] Result submit_scalar_readback(std::size_t slot, VkBuffer src) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + if (slot >= rb_slots_.size()) NN_FAIL("scalar readback: slot out of range"); + auto& s = rb_slots_[slot]; + + // 上一轮的值尚未取走：等它就绪（防御路径；正常调用方先 poll） + if (s.pending) { - for (auto& s : rb_slots_) + constexpr uint64_t kTimeoutNs = 60'000'000'000ULL; + const VkResult wr = vkWaitForFences(device_.device(), 1, &s.fence, VK_TRUE, kTimeoutNs); + if (wr == VK_ERROR_DEVICE_LOST) { - if (s.mapped != nullptr && s.alloc.valid()) - { - vkUnmapMemory(device_.device(), s.alloc.memory); - s.mapped = nullptr; - } - if (s.host != VK_NULL_HANDLE) - vkDestroyBuffer(device_.device(), s.host, nullptr); - if (s.fence != VK_NULL_HANDLE) - vkDestroyFence(device_.device(), s.fence, nullptr); - if (s.alloc.valid()) - memory_pool().free(s.alloc); - s.host = VK_NULL_HANDLE; - s.fence = VK_NULL_HANDLE; - s.cmd = VK_NULL_HANDLE; // command buffer 随 command pool 释放 - s.pending = false; + device_lost_ = true; + NN_FAIL("scalar readback: VK_ERROR_DEVICE_LOST（GPU 已被 TDR 重置）"); } - rb_slots_.clear(); + if (wr != VK_SUCCESS) + NN_FAIL("scalar readback: fence wait failed " + std::to_string(static_cast(wr))); + s.pending = false; } - // 录制一次 4B D2H 拷贝并提交（**不等待**）。 - // 调用约定：必须在产出该标量的主帧已提交之后调用（同队列 FIFO ⇒ - // 拷贝排在生产者之后；若在主帧提交前提交，会读到上一轮的旧值）。 - [[nodiscard]] Result submit_scalar_readback(std::size_t slot, VkBuffer src) + auto r = detail::vk_check(vkResetCommandBuffer(s.cmd, 0), __FILE__, __LINE__); + NN_TRY_CHECK(r); + VkCommandBufferBeginInfo cbi{}; + cbi.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO; + cbi.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT; + r = detail::vk_check(vkBeginCommandBuffer(s.cmd, &cbi), __FILE__, __LINE__); + NN_TRY_CHECK(r); + VkBufferCopy cp{}; + cp.srcOffset = 0; + cp.dstOffset = 0; + cp.size = sizeof(float); // (1,1) f32 标量 + vkCmdCopyBuffer(s.cmd, src, s.host, 1, &cp); + + VkBufferMemoryBarrier b{}; + b.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER; + b.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; + b.dstAccessMask = VK_ACCESS_HOST_READ_BIT; + b.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; + b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; + b.buffer = s.host; + b.offset = 0; + b.size = sizeof(float); + vkCmdPipelineBarrier(s.cmd, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_HOST_BIT, 0, 0, + nullptr, 1, &b, 0, nullptr); + + r = detail::vk_check(vkEndCommandBuffer(s.cmd), __FILE__, __LINE__); + NN_TRY_CHECK(r); + r = detail::vk_check(vkResetFences(device_.device(), 1, &s.fence), __FILE__, __LINE__); + NN_TRY_CHECK(r); + { + std::lock_guard lock(queue_mutex_); + VkSubmitInfo si{}; + si.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO; + si.commandBufferCount = 1; + si.pCommandBuffers = &s.cmd; + r = detail::vk_check(vkQueueSubmit(device_.compute_queue(), 1, &si, s.fence), __FILE__, + __LINE__); + } + NN_TRY_CHECK(r); + s.pending = true; + return {}; +} + +// 非阻塞查询：就绪则写出 out 并返回 true；未就绪返回 false（不阻塞）。 +[[nodiscard]] Result poll_scalar_readback(std::size_t slot, float& out) +{ + if (slot >= rb_slots_.size()) NN_FAIL("scalar readback: slot out of range"); + auto& s = rb_slots_[slot]; + if (!s.pending) return false; + const VkResult st = vkGetFenceStatus(device_.device(), s.fence); + if (st == VK_NOT_READY) return false; + if (st == VK_ERROR_DEVICE_LOST) + { + device_lost_ = true; + NN_FAIL("scalar readback: VK_ERROR_DEVICE_LOST（GPU 已被 TDR 重置）"); + } + if (st != VK_SUCCESS) + NN_FAIL("scalar readback: fence status " + std::to_string(static_cast(st))); + // size = VK_WHOLE_SIZE 时 offset 无对齐要求（VU 只约束非 WHOLE_SIZE）。 + VkMappedMemoryRange mr{}; + mr.sType = VK_STRUCTURE_TYPE_MAPPED_MEMORY_RANGE; + mr.memory = s.alloc.memory; + mr.offset = s.alloc.offset; + mr.size = VK_WHOLE_SIZE; + if (vkInvalidateMappedMemoryRanges(device_.device(), 1, &mr) != VK_SUCCESS) + NN_FAIL("scalar readback: invalidate mapped range failed"); + std::memcpy(&out, s.mapped, sizeof(float)); + s.pending = false; + return true; +} + +// ══════════════════════════════════════════════════════════════════ +// clone_gpu — GPU 内深拷贝（分配新 buffer + copy，无 PCIe 传输） +// ══════════════════════════════════════════════════════════════════ +template [[nodiscard]] Result> clone_gpu(const GpuTensorT

& src) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + // 1. 分配同形状的新 buffer + auto dst_res = GpuTensorT

::create_empty(src.rows(), src.cols(), *this); + NN_TRY_CHECK(dst_res); + GpuTensorT

dst = std::move(*dst_res); + + // 2. GPU 内拷贝 + const VkDeviceSize size = static_cast(src.rows() * src.cols() * sizeof(elem

)); + auto r = copy_buffer_gpu(src.buffer().impl(), dst.buffer().impl(), size); + NN_TRY_CHECK(r); + return dst; +} + +// ── 通用精度转换原语（engine.cast 的 GPU 实现，无 PCIe 往返）──────── +// 源/目标均为原始 32-bit word 缓冲；kind 分派精度对： +// 0 = f16→f32（src 2B/pair 打包，dst 4B/word） +// 1 = f32→f16（src 4B/word，dst 2B/pair 打包，round-half-to-even） +// 元素数保持不变（count = rows*cols），仅字节宽度差。后续精度对 +// （BF16/F64/f8 等）扩 kind 枚举即可，绝不改引擎 cast API。 +[[nodiscard]] Result cast_gpu(const GpuBuffer& src, const GpuBuffer& dst, std::size_t count, + std::uint32_t kind) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + if (!has(PipelineId::cast)) NN_FAIL("cast_gpu: cast pipeline not available"); + const VkBuffer bindings[2]{src.impl(), dst.impl()}; + auto ds_r = alloc_desc_set_for(pipelines_.get(PipelineId::cast).descriptor_layout(), bindings); + NN_TRY_CHECK(ds_r); + VkDescriptorSet desc_set = *ds_r; + + auto cmd_r = acquire_cmd(); + if (!cmd_r) { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - if (slot >= rb_slots_.size()) - NN_FAIL("scalar readback: slot out of range"); - auto& s = rb_slots_[slot]; + release_desc_set_now(desc_set); + NN_TRY_CHECK(cmd_r); + } + auto [cmd, owns_cmd] = *cmd_r; - // 上一轮的值尚未取走：等它就绪（防御路径；正常调用方先 poll） - if (s.pending) - { - constexpr uint64_t kTimeoutNs = 60'000'000'000ULL; - const VkResult wr = - vkWaitForFences(device_.device(), 1, &s.fence, VK_TRUE, kTimeoutNs); - if (wr == VK_ERROR_DEVICE_LOST) - { - device_lost_ = true; - NN_FAIL("scalar readback: VK_ERROR_DEVICE_LOST（GPU 已被 TDR 重置）"); - } - if (wr != VK_SUCCESS) - NN_FAIL("scalar readback: fence wait failed " + - std::to_string(static_cast(wr))); - s.pending = false; - } + std::vector in_bufs{src.impl()}; + record_input_barriers(cmd, in_bufs); - auto r = detail::vk_check(vkResetCommandBuffer(s.cmd, 0), __FILE__, __LINE__); - NN_TRY_CHECK(r); - VkCommandBufferBeginInfo cbi{}; - cbi.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO; - cbi.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT; - r = detail::vk_check(vkBeginCommandBuffer(s.cmd, &cbi), __FILE__, __LINE__); + vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, + pipelines_.get(PipelineId::cast).handle()); + vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, + pipelines_.get(PipelineId::cast).pipeline_layout(), 0, 1, &desc_set, 0, + nullptr); + + struct Push + { + std::uint32_t count; + std::uint32_t kind; + }; + Push push{static_cast(count), kind}; + std::vector pc(sizeof(push)); + std::memcpy(pc.data(), &push, sizeof(push)); + vkCmdPushConstants(cmd, pipelines_.get(PipelineId::cast).pipeline_layout(), + VK_SHADER_STAGE_COMPUTE_BIT, 0, static_cast(pc.size()), pc.data()); + + // 每线程处理 2 元素（一对一 half 或单词对） + const auto pairs = (count + 1u) / 2u; + const std::uint32_t wg = static_cast((pairs + 255u) / 256u); + ts_dispatch(cmd, "cast", wg < 1u ? 1u : wg, 1u, 1u); + record_output_barrier(cmd, dst.impl()); + + if (owns_cmd) + { + auto r = submit_and_wait(cmd, desc_set); NN_TRY_CHECK(r); - VkBufferCopy cp{}; - cp.srcOffset = 0; - cp.dstOffset = 0; - cp.size = sizeof(float); // (1,1) f32 标量 - vkCmdCopyBuffer(s.cmd, src, s.host, 1, &cp); + } + return {}; +} + +// ── 三值打包（P2/F3）：W (rows,k) f32 + τ (rows) f32 → 打包字 (rows,wpr) ── +[[nodiscard]] Result trit_pack_gpu(const GpuBuffer& w, const GpuBuffer& tau, + const GpuBuffer& packed, std::uint32_t rows, + std::uint32_t k, std::uint32_t wpr) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + if (!has(PipelineId::trit_pack)) NN_FAIL("trit_pack_gpu: trit_pack pipeline not available"); + const VkBuffer bindings[3]{w.impl(), tau.impl(), packed.impl()}; + auto ds_r = + alloc_desc_set_for(pipelines_.get(PipelineId::trit_pack).descriptor_layout(), bindings); + NN_TRY_CHECK(ds_r); + VkDescriptorSet desc_set = *ds_r; + + auto cmd_r = acquire_cmd(); + if (!cmd_r) + { + release_desc_set_now(desc_set); + NN_TRY_CHECK(cmd_r); + } + auto [cmd, owns_cmd] = *cmd_r; - VkBufferMemoryBarrier b{}; - b.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER; - b.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; - b.dstAccessMask = VK_ACCESS_HOST_READ_BIT; - b.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - b.buffer = s.host; - b.offset = 0; - b.size = sizeof(float); - vkCmdPipelineBarrier(s.cmd, - VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_HOST_BIT, - 0, 0, nullptr, 1, &b, 0, nullptr); + std::vector in_bufs{w.impl(), tau.impl()}; + record_input_barriers(cmd, in_bufs); + + vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, + pipelines_.get(PipelineId::trit_pack).handle()); + vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, + pipelines_.get(PipelineId::trit_pack).pipeline_layout(), 0, 1, + &desc_set, 0, nullptr); - r = detail::vk_check(vkEndCommandBuffer(s.cmd), __FILE__, __LINE__); + struct Push + { + std::uint32_t rows; + std::uint32_t k; + std::uint32_t wpr; + }; + Push push{rows, k, wpr}; + std::vector pc(sizeof(push)); + std::memcpy(pc.data(), &push, sizeof(push)); + vkCmdPushConstants(cmd, pipelines_.get(PipelineId::trit_pack).pipeline_layout(), + VK_SHADER_STAGE_COMPUTE_BIT, 0, static_cast(pc.size()), pc.data()); + + const std::uint32_t total = rows * wpr; + const std::uint32_t wg = (total + 255u) / 256u; + vkCmdDispatch(cmd, wg < 1u ? 1u : wg, 1u, 1u); + record_output_barrier(cmd, packed.impl()); + + if (owns_cmd) + { + auto r = submit_and_wait(cmd, desc_set); NN_TRY_CHECK(r); - r = detail::vk_check(vkResetFences(device_.device(), 1, &s.fence), __FILE__, __LINE__); + } + return {}; +} + +// ── 量化 GEMM（P2/F3）：C = (τ∘A)^(T)·B（A 为打包三值）──────────────── +// f16 = true 时用 f16 存储变体（B/C 均为 half），否则全 f32。 +[[nodiscard]] Result matmul_q_gpu(const GpuBuffer& packed, const GpuBuffer& tau, + const GpuBuffer& b, const GpuBuffer& c, std::uint32_t M, + std::uint32_t N, std::uint32_t K, std::uint32_t transA, + std::uint32_t wpr, bool f16) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + // ── 选址（与稠密 matmul 同口径，实测驱动；docs 21 §4.12）────────── + // 窄 N（≤ GEMV_MAX_N）→ **GEMV 内核**：分块版按 64×64 块算，n=1 时 + // 63/64 列空转（102/40HX 实测打包侧权重带宽只有 2.9 GB/s vs 稠密 + // GEMV 97 GB/s）→ 解码路必须走"每线程沿打包字分片 + 只算真实列"。 + // 宽 N（训练）→ 分块内核（分块解码一次进 shared，计算受限）。 + // GEMV 只在 transA=0（转发形态）实现。**布局契约**：shader 里每个 + // subgroup 负责 RPS=4 行（`m0 = (WGID.x*gl_NumSubgroups + SubgroupID)*4`）， + // 故 grid.x 必须 = ceil(M / (4*(256/subgroup_size)))，且 256 必须被 + // subgroup_size 整除（否则 gl_NumSubgroups 与 host 侧算出不一致）。 + // 设备侧 subgroup≥4 是红线（shader 的 subgroupAdd 归约语义）。 + constexpr std::uint32_t QGEMV_MAX_N = 8u; + constexpr std::uint32_t QGEMV_THREADS = 256u; // 与 matmul_q_gemv.comp 的 local_size_x 一致 + constexpr std::uint32_t QGEMV_RPS = 4u; // 与 matmul_q_gemv.comp 的 RPS 一致 + const std::uint32_t sgs = device_.subgroup_size(); + const bool qgemv_layout_ok = sgs >= 4u && sgs <= QGEMV_THREADS && (QGEMV_THREADS % sgs) == 0u; + const std::uint32_t qgemv_rows = qgemv_layout_ok ? (QGEMV_THREADS / sgs) * QGEMV_RPS : 0u; + const bool use_gemv = (transA == 0u) && N <= QGEMV_MAX_N && qgemv_layout_ok && + has(PipelineId::matmul_q_gemv) && + (!f16 || has(PipelineId::matmul_q_gemv_f16)); + const VulkanPipeline* pipe = nullptr; + if (use_gemv) + pipe = f16 ? &pipelines_.get(PipelineId::matmul_q_gemv_f16) + : &pipelines_.get(PipelineId::matmul_q_gemv); + else if (f16) + { + if (!has(PipelineId::matmul_q_f16)) + NN_FAIL("matmul_q_gpu: f16 pipeline not available（设备缺 16 位存储？）"); + pipe = &pipelines_.get(PipelineId::matmul_q_f16); + } + else + { + if (!has(PipelineId::matmul_q)) NN_FAIL("matmul_q_gpu: matmul_q pipeline not available"); + pipe = &pipelines_.get(PipelineId::matmul_q); + } + const VkBuffer bindings[4]{packed.impl(), tau.impl(), b.impl(), c.impl()}; + auto ds_r = alloc_desc_set_for(pipe->descriptor_layout(), bindings); + NN_TRY_CHECK(ds_r); + VkDescriptorSet desc_set = *ds_r; + + auto cmd_r = acquire_cmd(); + if (!cmd_r) + { + release_desc_set_now(desc_set); + NN_TRY_CHECK(cmd_r); + } + auto [cmd, owns_cmd] = *cmd_r; + + std::vector in_bufs{packed.impl(), tau.impl(), b.impl()}; + record_input_barriers(cmd, in_bufs); + + vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipe->handle()); + vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipe->pipeline_layout(), 0, 1, + &desc_set, 0, nullptr); + + struct Push + { std::uint32_t M, N, K, transA, wpr; }; + Push push{M, N, K, transA, wpr}; + std::vector pc(sizeof(push)); + std::memcpy(pc.data(), &push, sizeof(push)); + vkCmdPushConstants(cmd, pipe->pipeline_layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, + static_cast(pc.size()), pc.data()); + + // 分块版 matmul_q.comp：16×16 线程 / 输出块 64×64（BM=BN=64）。 + // ⚠ 两条不变量（分块内核的"形状契约"，漏掉任一条都是**静默错值**）： + // 1. 网格必须与 shader 的分块常量一致：写少了漏算、写多了重复算 + // （正确但成倍变慢 —— 实测错误网格下宽 batch 训练慢 13×）。 + // 2. **轴向必须对齐**：shader 里 x → block_col（N 方向）、y → + // block_row（M 方向），与 matmul_tiled / batched_matmul 同一约定 + // （两者 dispatch 均为 `(ceil(N/64), ceil(M/64), 1)`）。写成 + // `(ceil(M/64), ceil(N/64))` 会让 M≠N 的形状只算到 + // `ceil(N/64)*64` 行 / `ceil(M/64)*64` 列 —— 方形小形状 + // （M=N≤64，单 workgroup）恰好正确，故此缺陷只在真实形状上暴露 + // （2B4T 的 lm_head M=128256 只算到前 256 行）。 + if (use_gemv) + { + // 单维 grid.x = ceil(M/qgemv_rows)；qgemv_rows 见上面的布局契约注释 + const std::uint32_t gw = (M + qgemv_rows - 1u) / qgemv_rows; + vkCmdDispatch(cmd, gw < 1u ? 1u : gw, 1u, 1u); + } + else + { + const std::uint32_t gx = (N + 63u) / 64u; // x ↔ 输出列（N） + const std::uint32_t gy = (M + 63u) / 64u; // y ↔ 输出行（M） + vkCmdDispatch(cmd, gx < 1u ? 1u : gx, gy < 1u ? 1u : gy, 1u); + } + record_output_barrier(cmd, c.impl()); + + if (owns_cmd) + { + auto r = submit_and_wait(cmd, desc_set); NN_TRY_CHECK(r); - { - std::lock_guard lock(queue_mutex_); - VkSubmitInfo si{}; - si.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO; - si.commandBufferCount = 1; - si.pCommandBuffers = &s.cmd; - r = detail::vk_check( - vkQueueSubmit(device_.compute_queue(), 1, &si, s.fence), __FILE__, __LINE__); - } + } + return {}; +} + +// ── 创建逻辑形状 (rows,cols) 但存储为偶数槽位的 f16 张量 ─────────── +// f32→f16 的 cast 每线程按整 word（两 half）写入；奇数元素 count 时最后一个 +// word 会越界 2 字节。这里把底层缓冲按偶数元素分配（多 1 槽），逻辑形状不变， +// 写 word 永不越界；to_matrix 仍只读前 rows*cols 个元素。 +[[nodiscard]] Result create_f16_tensor(std::size_t rows, std::size_t cols) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + const std::size_t cnt = rows * cols; + const std::size_t cnt_pad = cnt + (cnt & 1u); // 偶数槽位 + auto b = GpuBuffer::create_device_local(device_.device(), *memory_pool_, cnt_pad * 2u, + VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | + VK_BUFFER_USAGE_TRANSFER_SRC_BIT | + VK_BUFFER_USAGE_TRANSFER_DST_BIT, + keepalive()); + NN_TRY_CHECK(b); + return GpuTensorF16(std::make_shared(std::move(*b)), rows, cols); +} +template +[[nodiscard]] Result> slice_rows_gpu(const GpuTensorT

& src, std::size_t start_row, + std::size_t count) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + const std::size_t cols = src.cols(); + if (start_row + count > src.rows()) NN_FAIL("slice_rows_gpu: range out of bounds"); + // 分配目标 buffer + auto dst_res = GpuTensorT

::create_empty(count, cols, *this); + NN_TRY_CHECK(dst_res); + GpuTensorT

dst = std::move(*dst_res); + + // 行区间在行主序下连续：[start_row * cols, (start_row + count) * cols) + const std::size_t elem_size = sizeof(elem

); + const VkDeviceSize src_offset = static_cast(start_row * cols * elem_size); + const VkDeviceSize size = static_cast(count * cols * elem_size); + + auto cmd_r = acquire_cmd(); + NN_TRY_CHECK(cmd_r); + auto [cmd, owns_cmd] = *cmd_r; + + VkBufferCopy cp{src_offset, 0, size}; + vkCmdCopyBuffer(cmd, src.buffer().impl(), dst.buffer().impl(), 1, &cp); + + VkBufferMemoryBarrier b{}; + b.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER; + b.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; + b.dstAccessMask = VK_ACCESS_MEMORY_READ_BIT; + b.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; + b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; + b.buffer = dst.buffer().impl(); + b.offset = 0; + b.size = VK_WHOLE_SIZE; + vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_TRANSFER_BIT, + VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, 0, + 0, nullptr, 1, &b, 0, nullptr); + + if (owns_cmd) + { + auto r = submit_and_wait(cmd, VK_NULL_HANDLE); NN_TRY_CHECK(r); - s.pending = true; - return {}; - } - - // 非阻塞查询：就绪则写出 out 并返回 true；未就绪返回 false（不阻塞）。 - [[nodiscard]] Result poll_scalar_readback(std::size_t slot, float& out) - { - if (slot >= rb_slots_.size()) - NN_FAIL("scalar readback: slot out of range"); - auto& s = rb_slots_[slot]; - if (!s.pending) - return false; - const VkResult st = vkGetFenceStatus(device_.device(), s.fence); - if (st == VK_NOT_READY) - return false; - if (st == VK_ERROR_DEVICE_LOST) - { - device_lost_ = true; - NN_FAIL("scalar readback: VK_ERROR_DEVICE_LOST（GPU 已被 TDR 重置）"); - } - if (st != VK_SUCCESS) - NN_FAIL("scalar readback: fence status " + - std::to_string(static_cast(st))); - // size = VK_WHOLE_SIZE 时 offset 无对齐要求（VU 只约束非 WHOLE_SIZE）。 - VkMappedMemoryRange mr{}; - mr.sType = VK_STRUCTURE_TYPE_MAPPED_MEMORY_RANGE; - mr.memory = s.alloc.memory; - mr.offset = s.alloc.offset; - mr.size = VK_WHOLE_SIZE; - if (vkInvalidateMappedMemoryRanges(device_.device(), 1, &mr) != VK_SUCCESS) - NN_FAIL("scalar readback: invalidate mapped range failed"); - std::memcpy(&out, s.mapped, sizeof(float)); - s.pending = false; - return true; - } - - // ══════════════════════════════════════════════════════════════════ - // clone_gpu — GPU 内深拷贝（分配新 buffer + copy，无 PCIe 传输） - // ══════════════════════════════════════════════════════════════════ - template - [[nodiscard]] Result> clone_gpu(const GpuTensorT

& src) - { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - // 1. 分配同形状的新 buffer - auto dst_res = GpuTensorT

::create_empty(src.rows(), src.cols(), *this); - NN_TRY_CHECK(dst_res); - GpuTensorT

dst = std::move(*dst_res); - - // 2. GPU 内拷贝 - const VkDeviceSize size = static_cast( - src.rows() * src.cols() * sizeof(elem

)); - auto r = copy_buffer_gpu(src.buffer().impl(), dst.buffer().impl(), size); + } + return dst; +} + +// ══════════════════════════════════════════════════════════════════ +// insert_rows_gpu — 行插入（GPU 内就地写入，无 PCIe 传输） +// 将 src 的所有行写入 dst 的行 [dst_start_row, dst_start_row + src.rows()) +// 真·就地修改（vkCmdCopyBuffer with dstOffset）。 +// 注意：dst 必须以 TRANSFER_DST_BIT 创建（create_empty 已包含）。 +// ══════════════════════════════════════════════════════════════════ +template +[[nodiscard]] Result insert_rows_gpu(GpuTensorT

& dst, std::size_t dst_start_row, + const GpuTensorT

& src) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + const std::size_t cols = dst.cols(); + if (src.cols() != cols) NN_FAIL("insert_rows_gpu: column count mismatch"); + if (dst_start_row + src.rows() > dst.rows()) NN_FAIL("insert_rows_gpu: range out of bounds"); + const std::size_t elem_size = sizeof(elem

); + const VkDeviceSize dst_offset = static_cast(dst_start_row * cols * elem_size); + const VkDeviceSize size = static_cast(src.rows() * cols * elem_size); + + auto cmd_r = acquire_cmd(); + NN_TRY_CHECK(cmd_r); + auto [cmd, owns_cmd] = *cmd_r; + + VkBufferCopy cp{0, dst_offset, size}; + vkCmdCopyBuffer(cmd, src.buffer().impl(), dst.buffer().impl(), 1, &cp); + + VkBufferMemoryBarrier b{}; + b.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER; + b.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; + b.dstAccessMask = VK_ACCESS_MEMORY_READ_BIT; + b.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; + b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; + b.buffer = dst.buffer().impl(); + b.offset = 0; + b.size = VK_WHOLE_SIZE; + vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_TRANSFER_BIT, + VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, 0, + 0, nullptr, 1, &b, 0, nullptr); + + if (owns_cmd) + { + auto r = submit_and_wait(cmd, VK_NULL_HANDLE); NN_TRY_CHECK(r); - return dst; - } - - // ── 通用精度转换原语（engine.cast 的 GPU 实现，无 PCIe 往返）──────── - // 源/目标均为原始 32-bit word 缓冲；kind 分派精度对： - // 0 = f16→f32（src 2B/pair 打包，dst 4B/word） - // 1 = f32→f16（src 4B/word，dst 2B/pair 打包，round-half-to-even） - // 元素数保持不变（count = rows*cols），仅字节宽度差。后续精度对 - // （BF16/F64/f8 等）扩 kind 枚举即可，绝不改引擎 cast API。 - [[nodiscard]] Result cast_gpu( - const GpuBuffer& src, const GpuBuffer& dst, - std::size_t count, std::uint32_t kind) - { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - if (!has(PipelineId::cast)) - NN_FAIL("cast_gpu: cast pipeline not available"); - const VkBuffer bindings[2]{src.impl(), dst.impl()}; - auto ds_r = alloc_desc_set_for( - pipelines_.get(PipelineId::cast).descriptor_layout(), bindings); - NN_TRY_CHECK(ds_r); - VkDescriptorSet desc_set = *ds_r; - - auto cmd_r = acquire_cmd(); - if (!cmd_r) - { - release_desc_set_now(desc_set); - NN_TRY_CHECK(cmd_r); - } - auto [cmd, owns_cmd] = *cmd_r; - - std::vector in_bufs{src.impl()}; - record_input_barriers(cmd, in_bufs); + } + return {}; +} + +// ── Staging Path：CPU span → GPU → 计算 → GPU → CPU span ───────── +[[nodiscard]] Result matmul_direct(std::span a, std::span b, + std::span c, std::size_t M, std::size_t N, + std::size_t K, uint32_t transA = 0, uint32_t transB = 0) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + // 1. 创建临时 GPU buffers（byte count = 元素数 × sizeof(float)，§6.3） + auto a_buf = GpuBuffer::create_device_local( + device_.device(), *memory_pool_, M * K * sizeof(float), + VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, keepalive()); + NN_TRY_CHECK(a_buf); + auto b_buf = GpuBuffer::create_device_local( + device_.device(), *memory_pool_, K * N * sizeof(float), + VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, keepalive()); + NN_TRY_CHECK(b_buf); + // 2. 上传数据 + auto a_tensor = GpuTensor(std::make_shared(std::move(*a_buf)), M, K); + auto b_tensor = GpuTensor(std::make_shared(std::move(*b_buf)), K, N); + + auto r = upload_blocking(a_tensor, a); + if (!r) return r; + + r = upload_blocking(b_tensor, b); + if (!r) return r; + + // 3. 执行 matmul + auto c_tensor_res = matmul_gpu(a_tensor, b_tensor, transA, transB); + NN_TRY_CHECK(c_tensor_res); + // 4. 下载结果 + r = download_blocking(*c_tensor_res, c); + if (!r) return r; + + return {}; +} + +// ══════════════════════════════════════════════════════════════════ +// transpose_gpu — 矩阵转置 (R, C) → (C, R) +// +// Push Constants (8 bytes): rows, cols +// Bindings: In(0), Out(1) +// ══════════════════════════════════════════════════════════════════ +[[nodiscard]] Result transpose_gpu(const GpuTensor& A, bool f16_io = false) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + if (!has(PipelineId::transpose)) NN_FAIL("transpose pipeline not available"); + // Phase C1 f16 变体：f16_io 请求但设备无 f16 pipeline（能力门控未过） + // → 显式报错，引擎层回退边界 cast（不静默绑错类型） + if (f16_io && !has(PipelineId::transpose_f16)) + NN_FAIL("transpose_gpu: f16 pipeline 不可用（需16位存储 + shaderFloat16）"); + const auto& pipeline = + f16_io ? pipelines_.get(PipelineId::transpose_f16) : pipelines_.get(PipelineId::transpose); + + const uint32_t R = static_cast(A.rows()); + const uint32_t C = static_cast(A.cols()); + + // 输出 (C, R)：f16 按 2B/元素分配 + 纯绑定视图（同 matmul_gpu 做法） + std::optional owned; + if (f16_io) + { + auto f16r = GpuTensorF16::create_empty(C, R, *this); + NN_TRY_CHECK(f16r); + owned.emplace(f16r->shared_buffer(), f16r->rows(), f16r->cols()); + } + else + { + auto out_res = GpuTensor::create_empty(C, R, *this); + NN_TRY_CHECK(out_res); + owned.emplace(std::move(*out_res)); + } + GpuTensor output = *owned; - vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipelines_.get(PipelineId::cast).handle()); - vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipelines_.get(PipelineId::cast).pipeline_layout(), 0, 1, &desc_set, 0, nullptr); - - struct Push { std::uint32_t count; std::uint32_t kind; }; - Push push{static_cast(count), kind}; - std::vector pc(sizeof(push)); - std::memcpy(pc.data(), &push, sizeof(push)); - vkCmdPushConstants(cmd, pipelines_.get(PipelineId::cast).pipeline_layout(), - VK_SHADER_STAGE_COMPUTE_BIT, 0, - static_cast(pc.size()), pc.data()); - - // 每线程处理 2 元素（一对一 half 或单词对） - const auto pairs = (count + 1u) / 2u; - const std::uint32_t wg = static_cast((pairs + 255u) / 256u); - ts_dispatch(cmd, "cast", wg < 1u ? 1u : wg, 1u, 1u); - record_output_barrier(cmd, dst.impl()); + // 分配描述符集 + const VkBuffer bindings[2]{A.buffer().impl(), output.buffer().impl()}; + auto ds_r = alloc_desc_set_for(pipeline.descriptor_layout(), bindings); + NN_TRY_CHECK(ds_r); + VkDescriptorSet desc_set = *ds_r; - if (owns_cmd) - { - auto r = submit_and_wait(cmd, desc_set); - NN_TRY_CHECK(r); - } - return {}; + // 获取 command buffer + auto cmd_r = acquire_cmd(); + if (!cmd_r) + { + release_desc_set_now(desc_set); + NN_TRY_CHECK(cmd_r); } - - // ── 三值打包（P2/F3）：W (rows,k) f32 + τ (rows) f32 → 打包字 (rows,wpr) ── - [[nodiscard]] Result trit_pack_gpu( - const GpuBuffer& w, const GpuBuffer& tau, const GpuBuffer& packed, - std::uint32_t rows, std::uint32_t k, std::uint32_t wpr) + auto [cmd, owns_cmd] = *cmd_r; + + // 录制 + record_input_barriers(cmd, {A.buffer().impl()}); + vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline.handle()); + vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline.pipeline_layout(), 0, 1, + &desc_set, 0, nullptr); + + const uint32_t push_data[2] = {R, C}; + vkCmdPushConstants(cmd, pipeline.pipeline_layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, + sizeof(push_data), push_data); + + // 砖块化 dispatch：(8, 8, n_bricks)——每 WG 覆盖 64×64 输出块， + // 8×8 个 WG 组成一砖（并发足迹聚集，见 transpose.comp 头注释）。 + // 越界 tile 由 shader 早退（dispatch 固定 8×8，末砖可不满）。 + // ⚠️ x 必须是 8 而非 16：shader 按 8 宽砖解算 tile_c = + // (bz%bx_count)*8 + gl_WorkGroupID.x，假定 gl_WorkGroupID.x ∈ [0,8)； + // 派发 16 宽会让 tile_c 跨两砖、越界 tile 只早退一半，行>512 且列>512 + // 时静默只写前 512 行。 + const uint32_t tx = (C + 63u) / 64u; + const uint32_t ty = (R + 63u) / 64u; + const uint32_t bx = (tx + 7u) / 8u; + const uint32_t by = (ty + 7u) / 8u; + ts_dispatch(cmd, "transpose", 8, 8, bx * by); + record_output_barrier(cmd, output.buffer().impl()); + + if (owns_cmd) + { + auto r = submit_and_wait(cmd, desc_set); + NN_TRY_CHECK(r); + } + return output; +} + +// ══════════════════════════════════════════════════════════════════ +// ══════════════════════════════════════════════════════════════════ +// im2col_gpu / col2im_gpu — 卷积/池化窗口展开（纯数据搬运，无算法） +// Push Constants (36 bytes): C,H,W,k,stride,pad,OH,OW,B +// Bindings: In(0), Out(1) +// ══════════════════════════════════════════════════════════════════ +[[nodiscard]] Result im2col_gpu(const GpuTensor& x, std::size_t C, std::size_t H, + std::size_t W, std::size_t k, std::size_t stride, + std::size_t pad, std::size_t OH, std::size_t OW, + bool f16_io = false) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + if (!has(PipelineId::im2col)) NN_FAIL("im2col pipeline not available"); + if (f16_io && !has(PipelineId::im2col_f16)) NN_FAIL("im2col f16 pipeline not available"); + if (C == 0 || H == 0 || W == 0 || k == 0 || stride == 0 || OH == 0 || OW == 0) + NN_FAIL("im2col_gpu: C/H/W/k/stride/OH/OW must be > 0"); + if (x.rows() != C * H * W) NN_FAIL("im2col_gpu: x must be (C*H*W, B)"); + const VulkanPipeline& pipe = + f16_io ? pipelines_.get(PipelineId::im2col_f16) : pipelines_.get(PipelineId::im2col); + + const std::size_t B = x.cols(); + const std::size_t rows = C * k * k; + const std::size_t cols = B * OH * OW; + + // f16 输出按 2B/元素分配 + 纯绑定视图（同 matmul_gpu/elementwise f16_io） + GpuTensor output; + if (f16_io) { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - if (!has(PipelineId::trit_pack)) - NN_FAIL("trit_pack_gpu: trit_pack pipeline not available"); - const VkBuffer bindings[3]{w.impl(), tau.impl(), packed.impl()}; - auto ds_r = alloc_desc_set_for( - pipelines_.get(PipelineId::trit_pack).descriptor_layout(), bindings); - NN_TRY_CHECK(ds_r); - VkDescriptorSet desc_set = *ds_r; + auto f16r = GpuTensorF16::create_empty(rows, cols, *this); + NN_TRY_CHECK(f16r); + output = GpuTensor(f16r->shared_buffer(), rows, cols); + } + else + { + auto out_res = GpuTensor::create_empty(rows, cols, *this); + NN_TRY_CHECK(out_res); + output = std::move(*out_res); + } - auto cmd_r = acquire_cmd(); - if (!cmd_r) - { - release_desc_set_now(desc_set); - NN_TRY_CHECK(cmd_r); - } - auto [cmd, owns_cmd] = *cmd_r; + const VkBuffer bindings[2]{x.buffer().impl(), output.buffer().impl()}; + auto ds_r = alloc_desc_set_for(pipe.descriptor_layout(), bindings); + NN_TRY_CHECK(ds_r); + VkDescriptorSet desc_set = *ds_r; - std::vector in_bufs{w.impl(), tau.impl()}; - record_input_barriers(cmd, in_bufs); + auto cmd_r = acquire_cmd(); + if (!cmd_r) + { + release_desc_set_now(desc_set); + NN_TRY_CHECK(cmd_r); + } + auto [cmd, owns_cmd] = *cmd_r; + + record_input_barriers(cmd, {x.buffer().impl()}); + vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipe.handle()); + vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipe.pipeline_layout(), 0, 1, + &desc_set, 0, nullptr); + + const uint32_t push_data[9] = { + static_cast(C), static_cast(H), static_cast(W), + static_cast(k), static_cast(stride), static_cast(pad), + static_cast(OH), static_cast(OW), static_cast(B)}; + vkCmdPushConstants(cmd, pipe.pipeline_layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, + sizeof(push_data), push_data); + + const std::size_t total = rows * cols; + const uint32_t wg_count = static_cast((total + 255) / 256); + ts_dispatch(cmd, "im2col", wg_count, 1, 1); + record_output_barrier(cmd, output.buffer().impl()); + + if (owns_cmd) + { + auto r = submit_and_wait(cmd, desc_set); + NN_TRY_CHECK(r); + } + return output; +} + +[[nodiscard]] Result col2im_gpu(const GpuTensor& col, std::size_t C, std::size_t H, + std::size_t W, std::size_t k, std::size_t stride, + std::size_t pad, std::size_t OH, std::size_t OW, + bool f16_io = false) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + if (!has(PipelineId::col2im)) NN_FAIL("col2im pipeline not available"); + if (f16_io && !has(PipelineId::col2im_f16)) NN_FAIL("col2im f16 pipeline not available"); + if (C == 0 || H == 0 || W == 0 || k == 0 || stride == 0 || OH == 0 || OW == 0) + NN_FAIL("col2im_gpu: C/H/W/k/stride/OH/OW must be > 0"); + const std::size_t P = OH * OW; + if (col.rows() != C * k * k || col.cols() == 0 || col.cols() % P != 0) + NN_FAIL("col2im_gpu: col must be (C*k*k, B*OH*OW)"); + const VulkanPipeline& pipe = + f16_io ? pipelines_.get(PipelineId::col2im_f16) : pipelines_.get(PipelineId::col2im); + + const std::size_t B = col.cols() / P; + const std::size_t rows = C * H * W; + const std::size_t cols = B; + + GpuTensor output; + if (f16_io) + { + auto f16r = GpuTensorF16::create_empty(rows, cols, *this); + NN_TRY_CHECK(f16r); + output = GpuTensor(f16r->shared_buffer(), rows, cols); + } + else + { + auto out_res = GpuTensor::create_empty(rows, cols, *this); + NN_TRY_CHECK(out_res); + output = std::move(*out_res); + } - vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipelines_.get(PipelineId::trit_pack).handle()); - vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipelines_.get(PipelineId::trit_pack).pipeline_layout(), 0, 1, &desc_set, 0, nullptr); + const VkBuffer bindings[2]{col.buffer().impl(), output.buffer().impl()}; + auto ds_r = alloc_desc_set_for(pipe.descriptor_layout(), bindings); + NN_TRY_CHECK(ds_r); + VkDescriptorSet desc_set = *ds_r; - struct Push { std::uint32_t rows; std::uint32_t k; std::uint32_t wpr; }; - Push push{rows, k, wpr}; - std::vector pc(sizeof(push)); - std::memcpy(pc.data(), &push, sizeof(push)); - vkCmdPushConstants(cmd, pipelines_.get(PipelineId::trit_pack).pipeline_layout(), - VK_SHADER_STAGE_COMPUTE_BIT, 0, - static_cast(pc.size()), pc.data()); + auto cmd_r = acquire_cmd(); + if (!cmd_r) + { + release_desc_set_now(desc_set); + NN_TRY_CHECK(cmd_r); + } + auto [cmd, owns_cmd] = *cmd_r; + + record_input_barriers(cmd, {col.buffer().impl()}); + vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipe.handle()); + vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipe.pipeline_layout(), 0, 1, + &desc_set, 0, nullptr); + + const uint32_t push_data[9] = { + static_cast(C), static_cast(H), static_cast(W), + static_cast(k), static_cast(stride), static_cast(pad), + static_cast(OH), static_cast(OW), static_cast(B)}; + vkCmdPushConstants(cmd, pipe.pipeline_layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, + sizeof(push_data), push_data); + + const std::size_t total = rows * cols; + const uint32_t wg_count = static_cast((total + 255) / 256); + ts_dispatch(cmd, "col2im", wg_count, 1, 1); + record_output_barrier(cmd, output.buffer().impl()); + + if (owns_cmd) + { + auto r = submit_and_wait(cmd, desc_set); + NN_TRY_CHECK(r); + } + return output; +} + +// ══════════════════════════════════════════════════════════════════ +// grouped_reduce_gpu — 分组归约（沿行方向按固定长度 R 分组求和/求最大） +// x (G*R, N) → out (G, N) +// Push Constants (16 bytes): G, R, N, reduce_op +// Bindings: In(0), Out(1)；Dispatch: ceil(G*N / 256) +// f16_in（Phase C2c）：x 按 half 读、输出恒 f32（CNN 调用 P=F32）； +// pipeline 未创建时硬报错，调用方提前查 has(PipelineId::group_reduce_f16) +// 走引擎内 cast 回退。 +// ══════════════════════════════════════════════════════════════════ +[[nodiscard]] Result grouped_reduce_gpu(const GpuTensor& x, std::size_t G, std::size_t R, + bool is_max, bool f16_in = false) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + if (!has(PipelineId::group_reduce)) NN_FAIL("group_reduce pipeline not available"); + if (f16_in && !has(PipelineId::group_reduce_f16)) + NN_FAIL("group_reduce f16 pipeline not available"); + if (G == 0 || R == 0) NN_FAIL("grouped_reduce_gpu: G/R must be > 0"); + if (x.rows() != G * R) NN_FAIL("grouped_reduce_gpu: x must be (G*R, N)"); + const std::size_t N = x.cols(); + if (N == 0) NN_FAIL("grouped_reduce_gpu: N must be > 0"); + const VulkanPipeline& pipe = f16_in ? pipelines_.get(PipelineId::group_reduce_f16) + : pipelines_.get(PipelineId::group_reduce); + + auto out_res = GpuTensor::create_empty(G, N, *this); + NN_TRY_CHECK(out_res); + GpuTensor output = std::move(*out_res); + + const VkBuffer bindings[2]{x.buffer().impl(), output.buffer().impl()}; + auto ds_r = alloc_desc_set_for(pipe.descriptor_layout(), bindings); + NN_TRY_CHECK(ds_r); + VkDescriptorSet desc_set = *ds_r; + + auto cmd_r = acquire_cmd(); + if (!cmd_r) + { + release_desc_set_now(desc_set); + NN_TRY_CHECK(cmd_r); + } + auto [cmd, owns_cmd] = *cmd_r; - const std::uint32_t total = rows * wpr; - const std::uint32_t wg = (total + 255u) / 256u; - vkCmdDispatch(cmd, wg < 1u ? 1u : wg, 1u, 1u); - record_output_barrier(cmd, packed.impl()); + record_input_barriers(cmd, {x.buffer().impl()}); + vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipe.handle()); + vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipe.pipeline_layout(), 0, 1, + &desc_set, 0, nullptr); - if (owns_cmd) - { - auto r = submit_and_wait(cmd, desc_set); - NN_TRY_CHECK(r); - } - return {}; - } - - // ── 量化 GEMM（P2/F3）：C = (τ∘A)^(T)·B（A 为打包三值）──────────────── - // f16 = true 时用 f16 存储变体（B/C 均为 half），否则全 f32。 - [[nodiscard]] Result matmul_q_gpu( - const GpuBuffer& packed, const GpuBuffer& tau, const GpuBuffer& b, - const GpuBuffer& c, std::uint32_t M, std::uint32_t N, std::uint32_t K, - std::uint32_t transA, std::uint32_t wpr, bool f16) - { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - // ── 选址（与稠密 matmul 同口径，实测驱动；docs 21 §4.12）────────── - // 窄 N（≤ GEMV_MAX_N）→ **GEMV 内核**：分块版按 64×64 块算，n=1 时 - // 63/64 列空转（102/40HX 实测打包侧权重带宽只有 2.9 GB/s vs 稠密 - // GEMV 97 GB/s）→ 解码路必须走"每线程沿打包字分片 + 只算真实列"。 - // 宽 N（训练）→ 分块内核（分块解码一次进 shared，计算受限）。 - // GEMV 只在 transA=0（转发形态）实现。**布局契约**：shader 里每个 - // subgroup 负责 RPS=4 行（`m0 = (WGID.x*gl_NumSubgroups + SubgroupID)*4`）， - // 故 grid.x 必须 = ceil(M / (4*(256/subgroup_size)))，且 256 必须被 - // subgroup_size 整除（否则 gl_NumSubgroups 与 host 侧算出不一致）。 - // 设备侧 subgroup≥4 是红线（shader 的 subgroupAdd 归约语义）。 - constexpr std::uint32_t QGEMV_MAX_N = 8u; - constexpr std::uint32_t QGEMV_THREADS = 256u; // 与 matmul_q_gemv.comp 的 local_size_x 一致 - constexpr std::uint32_t QGEMV_RPS = 4u; // 与 matmul_q_gemv.comp 的 RPS 一致 - const std::uint32_t sgs = device_.subgroup_size(); - const bool qgemv_layout_ok = sgs >= 4u && sgs <= QGEMV_THREADS - && (QGEMV_THREADS % sgs) == 0u; - const std::uint32_t qgemv_rows = - qgemv_layout_ok ? (QGEMV_THREADS / sgs) * QGEMV_RPS : 0u; - const bool use_gemv = (transA == 0u) && N <= QGEMV_MAX_N && qgemv_layout_ok - && has(PipelineId::matmul_q_gemv) - && (!f16 || has(PipelineId::matmul_q_gemv_f16)); - const VulkanPipeline* pipe = nullptr; - if (use_gemv) - pipe = f16 ? &pipelines_.get(PipelineId::matmul_q_gemv_f16) : &pipelines_.get(PipelineId::matmul_q_gemv); - else if (f16) - { - if (!has(PipelineId::matmul_q_f16)) - NN_FAIL("matmul_q_gpu: f16 pipeline not available（设备缺 16 位存储？）"); - pipe = &pipelines_.get(PipelineId::matmul_q_f16); - } - else - { - if (!has(PipelineId::matmul_q)) - NN_FAIL("matmul_q_gpu: matmul_q pipeline not available"); - pipe = &pipelines_.get(PipelineId::matmul_q); - } - const VkBuffer bindings[4]{ - packed.impl(), tau.impl(), b.impl(), c.impl()}; - auto ds_r = alloc_desc_set_for(pipe->descriptor_layout(), bindings); - NN_TRY_CHECK(ds_r); - VkDescriptorSet desc_set = *ds_r; + const uint32_t push_data[4] = {static_cast(G), static_cast(R), + static_cast(N), is_max ? 1u : 0u}; + vkCmdPushConstants(cmd, pipe.pipeline_layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, + sizeof(push_data), push_data); - auto cmd_r = acquire_cmd(); - if (!cmd_r) - { - release_desc_set_now(desc_set); - NN_TRY_CHECK(cmd_r); - } - auto [cmd, owns_cmd] = *cmd_r; + const std::size_t total = G * N; + const uint32_t wg_count = static_cast((total + 255) / 256); + ts_dispatch(cmd, "group_reduce", wg_count, 1, 1); + record_output_barrier(cmd, output.buffer().impl()); - std::vector in_bufs{packed.impl(), tau.impl(), b.impl()}; - record_input_barriers(cmd, in_bufs); + if (owns_cmd) + { + auto r = submit_and_wait(cmd, desc_set); + NN_TRY_CHECK(r); + } + return output; +} + +// gather_gpu — 按行索引查表 (GPU-native) +// +// table: (vocab, D), indices: (num,) → output: (num, D) +// output[i] = table[indices[i]]，越界索引返回零行 +// +// Push Constants (12 bytes): vocab, D, num +// Bindings: Table(0), Indices(1), Output(2) +// ══════════════════════════════════════════════════════════════════ +[[nodiscard]] Result gather_gpu(const GpuTensor& table, const GpuTensor& indices, + bool f16_io = false) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + if (!has(PipelineId::gather)) NN_FAIL("gather pipeline not available"); + // Phase C1 f16 变体：table/out f16；**indices 恒 f32**（行号整数值 + // 精度要求，shader 内 float indices 槽不变） + if (f16_io && !has(PipelineId::gather_f16)) + NN_FAIL("gather_gpu: f16 pipeline 不可用（设备无16位存储）"); + const auto& pipeline = + f16_io ? pipelines_.get(PipelineId::gather_f16) : pipelines_.get(PipelineId::gather); + + const uint32_t vocab = static_cast(table.rows()); + const uint32_t D = static_cast(table.cols()); + const uint32_t num = static_cast(indices.rows() * indices.cols()); + + // 输出 (num, D)（f16：2B/元素 + 纯绑定视图，同 matmul_gpu） + std::optional owned; + if (f16_io) + { + auto f16r = GpuTensorF16::create_empty(num, D, *this); + NN_TRY_CHECK(f16r); + owned.emplace(f16r->shared_buffer(), f16r->rows(), f16r->cols()); + } + else + { + auto out_res = GpuTensor::create_empty(num, D, *this); + NN_TRY_CHECK(out_res); + owned.emplace(std::move(*out_res)); + } + GpuTensor output = *owned; + + // 分配描述符集 + const VkBuffer bindings[3]{table.buffer().impl(), indices.buffer().impl(), + output.buffer().impl()}; + auto ds_r = alloc_desc_set_for(pipeline.descriptor_layout(), bindings); + NN_TRY_CHECK(ds_r); + VkDescriptorSet desc_set = *ds_r; + + // 获取 command buffer + auto cmd_r = acquire_cmd(); + if (!cmd_r) + { + release_desc_set_now(desc_set); + NN_TRY_CHECK(cmd_r); + } + auto [cmd, owns_cmd] = *cmd_r; - vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipe->handle()); - vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipe->pipeline_layout(), 0, 1, &desc_set, 0, nullptr); + // 录制 + record_input_barriers(cmd, {table.buffer().impl(), indices.buffer().impl()}); + vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline.handle()); + vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline.pipeline_layout(), 0, 1, + &desc_set, 0, nullptr); - struct Push { - std::uint32_t M, N, K, transA, wpr; - }; - Push push{M, N, K, transA, wpr}; - std::vector pc(sizeof(push)); - std::memcpy(pc.data(), &push, sizeof(push)); - vkCmdPushConstants(cmd, pipe->pipeline_layout(), - VK_SHADER_STAGE_COMPUTE_BIT, 0, - static_cast(pc.size()), pc.data()); - - // 分块版 matmul_q.comp：16×16 线程 / 输出块 64×64（BM=BN=64）。 - // ⚠ 两条不变量（分块内核的"形状契约"，漏掉任一条都是**静默错值**）： - // 1. 网格必须与 shader 的分块常量一致：写少了漏算、写多了重复算 - // （正确但成倍变慢 —— 实测错误网格下宽 batch 训练慢 13×）。 - // 2. **轴向必须对齐**：shader 里 x → block_col（N 方向）、y → - // block_row（M 方向），与 matmul_tiled / batched_matmul 同一约定 - // （两者 dispatch 均为 `(ceil(N/64), ceil(M/64), 1)`）。写成 - // `(ceil(M/64), ceil(N/64))` 会让 M≠N 的形状只算到 - // `ceil(N/64)*64` 行 / `ceil(M/64)*64` 列 —— 方形小形状 - // （M=N≤64，单 workgroup）恰好正确，故此缺陷只在真实形状上暴露 - // （2B4T 的 lm_head M=128256 只算到前 256 行）。 - if (use_gemv) - { - // 单维 grid.x = ceil(M/qgemv_rows)；qgemv_rows 见上面的布局契约注释 - const std::uint32_t gw = (M + qgemv_rows - 1u) / qgemv_rows; - vkCmdDispatch(cmd, gw < 1u ? 1u : gw, 1u, 1u); - } - else - { - const std::uint32_t gx = (N + 63u) / 64u; // x ↔ 输出列（N） - const std::uint32_t gy = (M + 63u) / 64u; // y ↔ 输出行（M） - vkCmdDispatch(cmd, gx < 1u ? 1u : gx, gy < 1u ? 1u : gy, 1u); - } - record_output_barrier(cmd, c.impl()); + const uint32_t push_data[3] = {vocab, D, num}; + vkCmdPushConstants(cmd, pipeline.pipeline_layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, + sizeof(push_data), push_data); - if (owns_cmd) - { - auto r = submit_and_wait(cmd, desc_set); - NN_TRY_CHECK(r); - } - return {}; - } - - // ── 创建逻辑形状 (rows,cols) 但存储为偶数槽位的 f16 张量 ─────────── - // f32→f16 的 cast 每线程按整 word（两 half）写入；奇数元素 count 时最后一个 - // word 会越界 2 字节。这里把底层缓冲按偶数元素分配（多 1 槽），逻辑形状不变， - // 写 word 永不越界；to_matrix 仍只读前 rows*cols 个元素。 - [[nodiscard]] Result create_f16_tensor( - std::size_t rows, std::size_t cols) - { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - const std::size_t cnt = rows * cols; - const std::size_t cnt_pad = cnt + (cnt & 1u); // 偶数槽位 - auto b = GpuBuffer::create_device_local( - device_.device(), *memory_pool_, cnt_pad * 2u, - VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT - | VK_BUFFER_USAGE_TRANSFER_DST_BIT, - keepalive()); - NN_TRY_CHECK(b); - return GpuTensorF16(std::make_shared(std::move(*b)), rows, cols); - } - template - [[nodiscard]] Result> slice_rows_gpu( - const GpuTensorT

& src, std::size_t start_row, std::size_t count) - { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - const std::size_t cols = src.cols(); - if (start_row + count > src.rows()) - NN_FAIL("slice_rows_gpu: range out of bounds"); - // 分配目标 buffer - auto dst_res = GpuTensorT

::create_empty(count, cols, *this); - NN_TRY_CHECK(dst_res); - GpuTensorT

dst = std::move(*dst_res); - - // 行区间在行主序下连续：[start_row * cols, (start_row + count) * cols) - const std::size_t elem_size = sizeof(elem

); - const VkDeviceSize src_offset = static_cast( - start_row * cols * elem_size); - const VkDeviceSize size = static_cast( - count * cols * elem_size); + const uint32_t total = num * D; + // vec4 kernel（与 elementwise/broadcast/DSL 同口径）：每线程 4 元素 + const uint32_t wg_count = (total + 256u * 4u - 1u) / (256u * 4u); + ts_dispatch(cmd, "gather", wg_count, 1, 1); + record_output_barrier(cmd, output.buffer().impl()); - auto cmd_r = acquire_cmd(); + if (owns_cmd) + { + auto r = submit_and_wait(cmd, desc_set); + NN_TRY_CHECK(r); + } + return output; +} + +// ══════════════════════════════════════════════════════════════════ +// scatter_add_gpu — 按行索引原子累加梯度 (GPU-native) +// +// dst: (vocab, D) 原地修改，indices: (num,), grad: (num, D) +// dst[indices[i]][d] += grad[i][d] +// 使用 CAS 循环实现 float atomicAdd +// +// Push Constants (12 bytes): vocab, D, num +// Bindings: Dst(0), Indices(1), Grad(2) +// +// 注意：dst buffer 内容以 uint 视角访问用于 atomicCompSwap。 +// dst 初始值必须为合法 float 位模式（如 0x00000000 = 0.0f）。 +// f16_io（Phase C2b）：dst/grad 为 f16 存储——f16 shader 里 dst 每 uint +// word 打包 2 个 half（CAS 覆盖整 word）、grad half 直读；indices 恒 +// float（行号 >2048 禁 f16）。pipeline 未创建时硬报错，调用方提前查 +// has(PipelineId::scatter_add_f16) 走引擎内 cast 回退。 +// ══════════════════════════════════════════════════════════════════ +[[nodiscard]] Result scatter_add_gpu(GpuTensor& dst, const GpuTensor& indices, + const GpuTensor& grad, bool f16_io = false) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + if (!has(PipelineId::scatter_add)) NN_FAIL("scatter_add pipeline not available"); + const VulkanPipeline& pipe = f16_io ? pipelines_.get(PipelineId::scatter_add_f16) + : pipelines_.get(PipelineId::scatter_add); + if (f16_io && !has(PipelineId::scatter_add_f16)) + NN_FAIL("scatter_add f16 pipeline not available"); + const uint32_t vocab = static_cast(dst.rows()); + const uint32_t D = static_cast(dst.cols()); + const uint32_t num = static_cast(indices.rows() * indices.cols()); + + // 分配描述符集 + // dst 使用 uint 视图用于原子操作 + const VkBuffer bindings[3]{dst.buffer().impl(), indices.buffer().impl(), grad.buffer().impl()}; + auto ds_r = alloc_desc_set_for(pipe.descriptor_layout(), bindings); + NN_TRY_CHECK(ds_r); + VkDescriptorSet desc_set = *ds_r; + + // 获取 command buffer + auto cmd_r = acquire_cmd(); + if (!cmd_r) + { + release_desc_set_now(desc_set); NN_TRY_CHECK(cmd_r); - auto [cmd, owns_cmd] = *cmd_r; - - VkBufferCopy cp{src_offset, 0, size}; - vkCmdCopyBuffer(cmd, src.buffer().impl(), dst.buffer().impl(), 1, &cp); + } + auto [cmd, owns_cmd] = *cmd_r; + // 录制（dst 需要读+写屏障） + { + std::vector barriers; VkBufferMemoryBarrier b{}; b.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER; - b.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; - b.dstAccessMask = VK_ACCESS_MEMORY_READ_BIT; + b.srcAccessMask = VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT; + b.dstAccessMask = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT; b.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - b.buffer = dst.buffer().impl(); b.offset = 0; b.size = VK_WHOLE_SIZE; + b.buffer = dst.buffer().impl(); + barriers.push_back(b); + b.buffer = indices.buffer().impl(); + barriers.push_back(b); + b.buffer = grad.buffer().impl(); + barriers.push_back(b); vkCmdPipelineBarrier(cmd, - VK_PIPELINE_STAGE_TRANSFER_BIT, - VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, - 0, 0, nullptr, 1, &b, 0, nullptr); + VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, + VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 0, nullptr, + static_cast(barriers.size()), barriers.data(), 0, nullptr); + } - if (owns_cmd) - { - auto r = submit_and_wait(cmd, VK_NULL_HANDLE); - NN_TRY_CHECK(r); - } - return dst; - } - - // ══════════════════════════════════════════════════════════════════ - // insert_rows_gpu — 行插入（GPU 内就地写入，无 PCIe 传输） - // 将 src 的所有行写入 dst 的行 [dst_start_row, dst_start_row + src.rows()) - // 真·就地修改（vkCmdCopyBuffer with dstOffset）。 - // 注意：dst 必须以 TRANSFER_DST_BIT 创建（create_empty 已包含）。 - // ══════════════════════════════════════════════════════════════════ - template - [[nodiscard]] Result insert_rows_gpu( - GpuTensorT

& dst, std::size_t dst_start_row, const GpuTensorT

& src) - { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - const std::size_t cols = dst.cols(); - if (src.cols() != cols) - NN_FAIL("insert_rows_gpu: column count mismatch"); - if (dst_start_row + src.rows() > dst.rows()) - NN_FAIL("insert_rows_gpu: range out of bounds"); - const std::size_t elem_size = sizeof(elem

); - const VkDeviceSize dst_offset = static_cast( - dst_start_row * cols * elem_size); - const VkDeviceSize size = static_cast( - src.rows() * cols * elem_size); + vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipe.handle()); + vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipe.pipeline_layout(), 0, 1, + &desc_set, 0, nullptr); - auto cmd_r = acquire_cmd(); - NN_TRY_CHECK(cmd_r); - auto [cmd, owns_cmd] = *cmd_r; + const uint32_t push_data[3] = {vocab, D, num}; + vkCmdPushConstants(cmd, pipe.pipeline_layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, + sizeof(push_data), push_data); - VkBufferCopy cp{0, dst_offset, size}; - vkCmdCopyBuffer(cmd, src.buffer().impl(), dst.buffer().impl(), 1, &cp); + const uint32_t total = num * D; + const uint32_t wg_count = (total + 255) / 256; + ts_dispatch(cmd, "scatter_add", wg_count, 1, 1); + // 输出屏障（dst 被修改，需要 memory barrier） + { VkBufferMemoryBarrier b{}; b.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER; - b.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; - b.dstAccessMask = VK_ACCESS_MEMORY_READ_BIT; + b.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT; + b.dstAccessMask = VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT; b.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; b.buffer = dst.buffer().impl(); b.offset = 0; b.size = VK_WHOLE_SIZE; - vkCmdPipelineBarrier(cmd, - VK_PIPELINE_STAGE_TRANSFER_BIT, - VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, - 0, 0, nullptr, 1, &b, 0, nullptr); - - if (owns_cmd) - { - auto r = submit_and_wait(cmd, VK_NULL_HANDLE); - NN_TRY_CHECK(r); - } - return {}; - } - - // ── Staging Path：CPU span → GPU → 计算 → GPU → CPU span ───────── - [[nodiscard]] Result matmul_direct( - std::span a, std::span b, - std::span c, std::size_t M, std::size_t N, std::size_t K, - uint32_t transA = 0, uint32_t transB = 0) - { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - // 1. 创建临时 GPU buffers（byte count = 元素数 × sizeof(float)，§6.3） - auto a_buf = GpuBuffer::create_device_local( - device_.device(), *memory_pool_, M * K * sizeof(float), - VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, - keepalive()); - NN_TRY_CHECK(a_buf); - auto b_buf = GpuBuffer::create_device_local( - device_.device(), *memory_pool_, K * N * sizeof(float), - VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, - keepalive()); - NN_TRY_CHECK(b_buf); - // 2. 上传数据 - auto a_tensor = GpuTensor(std::make_shared(std::move(*a_buf)), M, K); - auto b_tensor = GpuTensor(std::make_shared(std::move(*b_buf)), K, N); - - auto r = upload_blocking(a_tensor, a); - if (!r) - return r; - - r = upload_blocking(b_tensor, b); - if (!r) - return r; - - // 3. 执行 matmul - auto c_tensor_res = matmul_gpu(a_tensor, b_tensor, transA, transB); - NN_TRY_CHECK(c_tensor_res); - // 4. 下载结果 - r = download_blocking(*c_tensor_res, c); - if (!r) - return r; - - return {}; - } - - // ══════════════════════════════════════════════════════════════════ - // transpose_gpu — 矩阵转置 (R, C) → (C, R) - // - // Push Constants (8 bytes): rows, cols - // Bindings: In(0), Out(1) - // ══════════════════════════════════════════════════════════════════ - [[nodiscard]] Result transpose_gpu(const GpuTensor& A, - bool f16_io = false) - { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - if (!has(PipelineId::transpose)) - NN_FAIL("transpose pipeline not available"); - // Phase C1 f16 变体：f16_io 请求但设备无 f16 pipeline（能力门控未过） - // → 显式报错，引擎层回退边界 cast（不静默绑错类型） - if (f16_io && !has(PipelineId::transpose_f16)) - NN_FAIL("transpose_gpu: f16 pipeline 不可用（需16位存储 + shaderFloat16）"); - const auto& pipeline = f16_io ? pipelines_.get(PipelineId::transpose_f16) : pipelines_.get(PipelineId::transpose); - - const uint32_t R = static_cast(A.rows()); - const uint32_t C = static_cast(A.cols()); - - // 输出 (C, R)：f16 按 2B/元素分配 + 纯绑定视图（同 matmul_gpu 做法） - std::optional owned; - if (f16_io) - { - auto f16r = GpuTensorF16::create_empty(C, R, *this); - NN_TRY_CHECK(f16r); - owned.emplace(f16r->shared_buffer(), f16r->rows(), f16r->cols()); - } - else - { - auto out_res = GpuTensor::create_empty(C, R, *this); - NN_TRY_CHECK(out_res); - owned.emplace(std::move(*out_res)); - } - GpuTensor output = *owned; - - // 分配描述符集 - const VkBuffer bindings[2]{A.buffer().impl(), output.buffer().impl()}; - auto ds_r = alloc_desc_set_for(pipeline.descriptor_layout(), bindings); - NN_TRY_CHECK(ds_r); - VkDescriptorSet desc_set = *ds_r; - - // 获取 command buffer - auto cmd_r = acquire_cmd(); - if (!cmd_r) - { - release_desc_set_now(desc_set); - NN_TRY_CHECK(cmd_r); - } - auto [cmd, owns_cmd] = *cmd_r; - - // 录制 - record_input_barriers(cmd, {A.buffer().impl()}); - vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipeline.handle()); - vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipeline.pipeline_layout(), 0, 1, &desc_set, 0, nullptr); - - const uint32_t push_data[2] = {R, C}; - vkCmdPushConstants(cmd, pipeline.pipeline_layout(), - VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(push_data), push_data); - - // 砖块化 dispatch：(8, 8, n_bricks)——每 WG 覆盖 64×64 输出块， - // 8×8 个 WG 组成一砖（并发足迹聚集，见 transpose.comp 头注释）。 - // 越界 tile 由 shader 早退（dispatch 固定 8×8，末砖可不满）。 - // ⚠️ x 必须是 8 而非 16：shader 按 8 宽砖解算 tile_c = - // (bz%bx_count)*8 + gl_WorkGroupID.x，假定 gl_WorkGroupID.x ∈ [0,8)； - // 派发 16 宽会让 tile_c 跨两砖、越界 tile 只早退一半，行>512 且列>512 - // 时静默只写前 512 行。 - const uint32_t tx = (C + 63u) / 64u; - const uint32_t ty = (R + 63u) / 64u; - const uint32_t bx = (tx + 7u) / 8u; - const uint32_t by = (ty + 7u) / 8u; - ts_dispatch(cmd, "transpose", 8, 8, bx * by); - record_output_barrier(cmd, output.buffer().impl()); - - if (owns_cmd) - { - auto r = submit_and_wait(cmd, desc_set); - NN_TRY_CHECK(r); - } - return output; - } - - // ══════════════════════════════════════════════════════════════════ - // ══════════════════════════════════════════════════════════════════ - // im2col_gpu / col2im_gpu — 卷积/池化窗口展开（纯数据搬运，无算法） - // Push Constants (36 bytes): C,H,W,k,stride,pad,OH,OW,B - // Bindings: In(0), Out(1) - // ══════════════════════════════════════════════════════════════════ - [[nodiscard]] Result im2col_gpu( - const GpuTensor& x, - std::size_t C, std::size_t H, std::size_t W, - std::size_t k, std::size_t stride, std::size_t pad, - std::size_t OH, std::size_t OW, - bool f16_io = false) - { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - if (!has(PipelineId::im2col)) - NN_FAIL("im2col pipeline not available"); - if (f16_io && !has(PipelineId::im2col_f16)) - NN_FAIL("im2col f16 pipeline not available"); - if (C == 0 || H == 0 || W == 0 || k == 0 || stride == 0 || OH == 0 || OW == 0) - NN_FAIL("im2col_gpu: C/H/W/k/stride/OH/OW must be > 0"); - if (x.rows() != C * H * W) - NN_FAIL("im2col_gpu: x must be (C*H*W, B)"); - const VulkanPipeline& pipe = f16_io ? pipelines_.get(PipelineId::im2col_f16) : pipelines_.get(PipelineId::im2col); - - const std::size_t B = x.cols(); - const std::size_t rows = C * k * k; - const std::size_t cols = B * OH * OW; - - // f16 输出按 2B/元素分配 + 纯绑定视图（同 matmul_gpu/elementwise f16_io） - GpuTensor output; - if (f16_io) - { - auto f16r = GpuTensorF16::create_empty(rows, cols, *this); - NN_TRY_CHECK(f16r); - output = GpuTensor(f16r->shared_buffer(), rows, cols); - } - else - { - auto out_res = GpuTensor::create_empty(rows, cols, *this); - NN_TRY_CHECK(out_res); - output = std::move(*out_res); - } - - const VkBuffer bindings[2]{x.buffer().impl(), output.buffer().impl()}; - auto ds_r = alloc_desc_set_for(pipe.descriptor_layout(), bindings); - NN_TRY_CHECK(ds_r); - VkDescriptorSet desc_set = *ds_r; - - auto cmd_r = acquire_cmd(); - if (!cmd_r) - { - release_desc_set_now(desc_set); - NN_TRY_CHECK(cmd_r); - } - auto [cmd, owns_cmd] = *cmd_r; - - record_input_barriers(cmd, {x.buffer().impl()}); - vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipe.handle()); - vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipe.pipeline_layout(), 0, 1, &desc_set, 0, nullptr); - - const uint32_t push_data[9] = { - static_cast(C), static_cast(H), - static_cast(W), static_cast(k), - static_cast(stride), static_cast(pad), - static_cast(OH), static_cast(OW), - static_cast(B)}; - vkCmdPushConstants(cmd, pipe.pipeline_layout(), - VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(push_data), push_data); - - const std::size_t total = rows * cols; - const uint32_t wg_count = static_cast((total + 255) / 256); - ts_dispatch(cmd, "im2col", wg_count, 1, 1); - record_output_barrier(cmd, output.buffer().impl()); - - if (owns_cmd) - { - auto r = submit_and_wait(cmd, desc_set); - NN_TRY_CHECK(r); - } - return output; - } - - [[nodiscard]] Result col2im_gpu( - const GpuTensor& col, - std::size_t C, std::size_t H, std::size_t W, - std::size_t k, std::size_t stride, std::size_t pad, - std::size_t OH, std::size_t OW, - bool f16_io = false) - { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - if (!has(PipelineId::col2im)) - NN_FAIL("col2im pipeline not available"); - if (f16_io && !has(PipelineId::col2im_f16)) - NN_FAIL("col2im f16 pipeline not available"); - if (C == 0 || H == 0 || W == 0 || k == 0 || stride == 0 || OH == 0 || OW == 0) - NN_FAIL("col2im_gpu: C/H/W/k/stride/OH/OW must be > 0"); - const std::size_t P = OH * OW; - if (col.rows() != C * k * k || col.cols() == 0 || col.cols() % P != 0) - NN_FAIL("col2im_gpu: col must be (C*k*k, B*OH*OW)"); - const VulkanPipeline& pipe = f16_io ? pipelines_.get(PipelineId::col2im_f16) : pipelines_.get(PipelineId::col2im); - - const std::size_t B = col.cols() / P; - const std::size_t rows = C * H * W; - const std::size_t cols = B; - - GpuTensor output; - if (f16_io) - { - auto f16r = GpuTensorF16::create_empty(rows, cols, *this); - NN_TRY_CHECK(f16r); - output = GpuTensor(f16r->shared_buffer(), rows, cols); - } - else - { - auto out_res = GpuTensor::create_empty(rows, cols, *this); - NN_TRY_CHECK(out_res); - output = std::move(*out_res); - } - - const VkBuffer bindings[2]{col.buffer().impl(), output.buffer().impl()}; - auto ds_r = alloc_desc_set_for(pipe.descriptor_layout(), bindings); - NN_TRY_CHECK(ds_r); - VkDescriptorSet desc_set = *ds_r; - - auto cmd_r = acquire_cmd(); - if (!cmd_r) - { - release_desc_set_now(desc_set); - NN_TRY_CHECK(cmd_r); - } - auto [cmd, owns_cmd] = *cmd_r; - - record_input_barriers(cmd, {col.buffer().impl()}); - vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipe.handle()); - vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipe.pipeline_layout(), 0, 1, &desc_set, 0, nullptr); - - const uint32_t push_data[9] = { - static_cast(C), static_cast(H), - static_cast(W), static_cast(k), - static_cast(stride), static_cast(pad), - static_cast(OH), static_cast(OW), - static_cast(B)}; - vkCmdPushConstants(cmd, pipe.pipeline_layout(), - VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(push_data), push_data); - - const std::size_t total = rows * cols; - const uint32_t wg_count = static_cast((total + 255) / 256); - ts_dispatch(cmd, "col2im", wg_count, 1, 1); - record_output_barrier(cmd, output.buffer().impl()); - - if (owns_cmd) - { - auto r = submit_and_wait(cmd, desc_set); - NN_TRY_CHECK(r); - } - return output; + vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, + VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, + 0, 0, nullptr, 1, &b, 0, nullptr); } - // ══════════════════════════════════════════════════════════════════ - // grouped_reduce_gpu — 分组归约（沿行方向按固定长度 R 分组求和/求最大） - // x (G*R, N) → out (G, N) - // Push Constants (16 bytes): G, R, N, reduce_op - // Bindings: In(0), Out(1)；Dispatch: ceil(G*N / 256) - // f16_in（Phase C2c）：x 按 half 读、输出恒 f32（CNN 调用 P=F32）； - // pipeline 未创建时硬报错，调用方提前查 has(PipelineId::group_reduce_f16) - // 走引擎内 cast 回退。 - // ══════════════════════════════════════════════════════════════════ - [[nodiscard]] Result grouped_reduce_gpu( - const GpuTensor& x, std::size_t G, std::size_t R, bool is_max, - bool f16_in = false) - { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - if (!has(PipelineId::group_reduce)) - NN_FAIL("group_reduce pipeline not available"); - if (f16_in && !has(PipelineId::group_reduce_f16)) - NN_FAIL("group_reduce f16 pipeline not available"); - if (G == 0 || R == 0) - NN_FAIL("grouped_reduce_gpu: G/R must be > 0"); - if (x.rows() != G * R) - NN_FAIL("grouped_reduce_gpu: x must be (G*R, N)"); - const std::size_t N = x.cols(); - if (N == 0) - NN_FAIL("grouped_reduce_gpu: N must be > 0"); - const VulkanPipeline& pipe = f16_in ? pipelines_.get(PipelineId::group_reduce_f16) - : pipelines_.get(PipelineId::group_reduce); - - auto out_res = GpuTensor::create_empty(G, N, *this); - NN_TRY_CHECK(out_res); - GpuTensor output = std::move(*out_res); - - const VkBuffer bindings[2]{x.buffer().impl(), output.buffer().impl()}; - auto ds_r = alloc_desc_set_for(pipe.descriptor_layout(), bindings); - NN_TRY_CHECK(ds_r); - VkDescriptorSet desc_set = *ds_r; - - auto cmd_r = acquire_cmd(); - if (!cmd_r) - { - release_desc_set_now(desc_set); - NN_TRY_CHECK(cmd_r); - } - auto [cmd, owns_cmd] = *cmd_r; - - record_input_barriers(cmd, {x.buffer().impl()}); - vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipe.handle()); - vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipe.pipeline_layout(), 0, 1, &desc_set, 0, nullptr); - - const uint32_t push_data[4] = { - static_cast(G), static_cast(R), - static_cast(N), is_max ? 1u : 0u}; - vkCmdPushConstants(cmd, pipe.pipeline_layout(), - VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(push_data), push_data); - - const std::size_t total = G * N; - const uint32_t wg_count = static_cast((total + 255) / 256); - ts_dispatch(cmd, "group_reduce", wg_count, 1, 1); - record_output_barrier(cmd, output.buffer().impl()); - - if (owns_cmd) - { - auto r = submit_and_wait(cmd, desc_set); - NN_TRY_CHECK(r); - } - return output; - } - - // gather_gpu — 按行索引查表 (GPU-native) - // - // table: (vocab, D), indices: (num,) → output: (num, D) - // output[i] = table[indices[i]]，越界索引返回零行 - // - // Push Constants (12 bytes): vocab, D, num - // Bindings: Table(0), Indices(1), Output(2) - // ══════════════════════════════════════════════════════════════════ - [[nodiscard]] Result gather_gpu( - const GpuTensor& table, const GpuTensor& indices, - bool f16_io = false) - { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - if (!has(PipelineId::gather)) - NN_FAIL("gather pipeline not available"); - // Phase C1 f16 变体：table/out f16；**indices 恒 f32**（行号整数值 - // 精度要求，shader 内 float indices 槽不变） - if (f16_io && !has(PipelineId::gather_f16)) - NN_FAIL("gather_gpu: f16 pipeline 不可用（设备无16位存储）"); - const auto& pipeline = f16_io ? pipelines_.get(PipelineId::gather_f16) : pipelines_.get(PipelineId::gather); - - const uint32_t vocab = static_cast(table.rows()); - const uint32_t D = static_cast(table.cols()); - const uint32_t num = static_cast(indices.rows() * indices.cols()); - - // 输出 (num, D)（f16：2B/元素 + 纯绑定视图，同 matmul_gpu） - std::optional owned; - if (f16_io) - { - auto f16r = GpuTensorF16::create_empty(num, D, *this); - NN_TRY_CHECK(f16r); - owned.emplace(f16r->shared_buffer(), f16r->rows(), f16r->cols()); - } - else - { - auto out_res = GpuTensor::create_empty(num, D, *this); - NN_TRY_CHECK(out_res); - owned.emplace(std::move(*out_res)); - } - GpuTensor output = *owned; - - // 分配描述符集 - const VkBuffer bindings[3]{ - table.buffer().impl(), indices.buffer().impl(), output.buffer().impl()}; - auto ds_r = alloc_desc_set_for(pipeline.descriptor_layout(), bindings); - NN_TRY_CHECK(ds_r); - VkDescriptorSet desc_set = *ds_r; - - // 获取 command buffer - auto cmd_r = acquire_cmd(); - if (!cmd_r) - { - release_desc_set_now(desc_set); - NN_TRY_CHECK(cmd_r); - } - auto [cmd, owns_cmd] = *cmd_r; - - // 录制 - record_input_barriers(cmd, {table.buffer().impl(), indices.buffer().impl()}); - vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipeline.handle()); - vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipeline.pipeline_layout(), 0, 1, &desc_set, 0, nullptr); - - const uint32_t push_data[3] = {vocab, D, num}; - vkCmdPushConstants(cmd, pipeline.pipeline_layout(), - VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(push_data), push_data); - - const uint32_t total = num * D; - // vec4 kernel（与 elementwise/broadcast/DSL 同口径）：每线程 4 元素 - const uint32_t wg_count = (total + 256u * 4u - 1u) / (256u * 4u); - ts_dispatch(cmd, "gather", wg_count, 1, 1); - record_output_barrier(cmd, output.buffer().impl()); - - if (owns_cmd) - { - auto r = submit_and_wait(cmd, desc_set); - NN_TRY_CHECK(r); - } - return output; - } - - // ══════════════════════════════════════════════════════════════════ - // scatter_add_gpu — 按行索引原子累加梯度 (GPU-native) - // - // dst: (vocab, D) 原地修改，indices: (num,), grad: (num, D) - // dst[indices[i]][d] += grad[i][d] - // 使用 CAS 循环实现 float atomicAdd - // - // Push Constants (12 bytes): vocab, D, num - // Bindings: Dst(0), Indices(1), Grad(2) - // - // 注意：dst buffer 内容以 uint 视角访问用于 atomicCompSwap。 - // dst 初始值必须为合法 float 位模式（如 0x00000000 = 0.0f）。 - // f16_io（Phase C2b）：dst/grad 为 f16 存储——f16 shader 里 dst 每 uint - // word 打包 2 个 half（CAS 覆盖整 word）、grad half 直读；indices 恒 - // float（行号 >2048 禁 f16）。pipeline 未创建时硬报错，调用方提前查 - // has(PipelineId::scatter_add_f16) 走引擎内 cast 回退。 - // ══════════════════════════════════════════════════════════════════ - [[nodiscard]] Result scatter_add_gpu( - GpuTensor& dst, const GpuTensor& indices, const GpuTensor& grad, - bool f16_io = false) - { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - if (!has(PipelineId::scatter_add)) - NN_FAIL("scatter_add pipeline not available"); - const VulkanPipeline& pipe = f16_io ? pipelines_.get(PipelineId::scatter_add_f16) - : pipelines_.get(PipelineId::scatter_add); - if (f16_io && !has(PipelineId::scatter_add_f16)) - NN_FAIL("scatter_add f16 pipeline not available"); - const uint32_t vocab = static_cast(dst.rows()); - const uint32_t D = static_cast(dst.cols()); - const uint32_t num = static_cast(indices.rows() * indices.cols()); - - // 分配描述符集 - // dst 使用 uint 视图用于原子操作 - const VkBuffer bindings[3]{ - dst.buffer().impl(), indices.buffer().impl(), grad.buffer().impl()}; - auto ds_r = alloc_desc_set_for(pipe.descriptor_layout(), bindings); - NN_TRY_CHECK(ds_r); - VkDescriptorSet desc_set = *ds_r; - - // 获取 command buffer - auto cmd_r = acquire_cmd(); - if (!cmd_r) - { - release_desc_set_now(desc_set); - NN_TRY_CHECK(cmd_r); - } - auto [cmd, owns_cmd] = *cmd_r; - - // 录制（dst 需要读+写屏障） - { - std::vector barriers; - VkBufferMemoryBarrier b{}; - b.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER; - b.srcAccessMask = VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT; - b.dstAccessMask = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT; - b.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - b.offset = 0; - b.size = VK_WHOLE_SIZE; - b.buffer = dst.buffer().impl(); - barriers.push_back(b); - b.buffer = indices.buffer().impl(); - barriers.push_back(b); - b.buffer = grad.buffer().impl(); - barriers.push_back(b); - vkCmdPipelineBarrier(cmd, - VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, - VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, - 0, 0, nullptr, - static_cast(barriers.size()), barriers.data(), - 0, nullptr); - } - - vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipe.handle()); - vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, - pipe.pipeline_layout(), 0, 1, &desc_set, 0, nullptr); - - const uint32_t push_data[3] = {vocab, D, num}; - vkCmdPushConstants(cmd, pipe.pipeline_layout(), - VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(push_data), push_data); - - const uint32_t total = num * D; - const uint32_t wg_count = (total + 255) / 256; - ts_dispatch(cmd, "scatter_add", wg_count, 1, 1); - - // 输出屏障（dst 被修改，需要 memory barrier） - { - VkBufferMemoryBarrier b{}; - b.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER; - b.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT; - b.dstAccessMask = VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT; - b.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - b.buffer = dst.buffer().impl(); - b.offset = 0; - b.size = VK_WHOLE_SIZE; - vkCmdPipelineBarrier(cmd, - VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, - VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, - 0, 0, nullptr, 1, &b, 0, nullptr); - } - - if (owns_cmd) - { - auto r = submit_and_wait(cmd, desc_set); - NN_TRY_CHECK(r); - } - return {}; + if (owns_cmd) + { + auto r = submit_and_wait(cmd, desc_set); + NN_TRY_CHECK(r); } + return {}; +} diff --git a/include/neuralnet.cpp/backend/compute_vk_backend_pipelines.inc b/include/neuralnet.cpp/backend/compute_vk_backend_pipelines.inc index 8256a9cb..ca7cac12 100644 --- a/include/neuralnet.cpp/backend/compute_vk_backend_pipelines.inc +++ b/include/neuralnet.cpp/backend/compute_vk_backend_pipelines.inc @@ -4,718 +4,684 @@ // ⚠ 本文件不是独立头：仅由 backend/compute_vk_backend.hpp 在 // `class GpuBackend { ... }` 体内**按序** #include（顺序即成员序）。 // ═══════════════════════════════════════════════════════════════════════════ - // ══════════════════════════════════════════════════════════════════ - // RLA 扫描原语（手写原语，形状契约见 compute_engine.hpp 扫描级原语注释） - // 1 workgroup/头：dispatch (1, B*H, 1)；标量块在头块内逐行重复存放 - // ══════════════════════════════════════════════════════════════════ - - // ── 前缀扫描 + matvec 读出：输出 (rows*5, seq) ───────────────────── - // f16_io（Phase C2f）：输入输出均 half（f16 pipeline）；scratch St 仍 - // f32（引擎分配）。无 f16 pipeline 时硬报错，调用方走引擎内 cast 回退。 - [[nodiscard]] Result scan_prefix_outer_gpu( - const GpuTensor& K, const GpuTensor& V, const GpuTensor& P, const GpuTensor& R, - const GpuTensor& A0, const GpuTensor& B0, bool has_state, - uint32_t dk, uint32_t heads, bool causal, - const GpuTensor& boundary, bool has_bnd, - bool f16_io = false) +// ══════════════════════════════════════════════════════════════════ +// RLA 扫描原语（手写原语，形状契约见 compute_engine.hpp 扫描级原语注释） +// 1 workgroup/头：dispatch (1, B*H, 1)；标量块在头块内逐行重复存放 +// ══════════════════════════════════════════════════════════════════ + +// ── 前缀扫描 + matvec 读出：输出 (rows*5, seq) ───────────────────── +// f16_io（Phase C2f）：输入输出均 half（f16 pipeline）；scratch St 仍 +// f32（引擎分配）。无 f16 pipeline 时硬报错，调用方走引擎内 cast 回退。 +[[nodiscard]] Result scan_prefix_outer_gpu(const GpuTensor& K, const GpuTensor& V, + const GpuTensor& P, const GpuTensor& R, + const GpuTensor& A0, const GpuTensor& B0, + bool has_state, uint32_t dk, uint32_t heads, + bool causal, const GpuTensor& boundary, + bool has_bnd, bool f16_io = false) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + if (dk == 0u || heads == 0u) NN_FAIL("scan_prefix_outer_gpu: dk/heads must be > 0"); + const bool generic = dk > 64u; // 通用路径（状态驻全局 scratch，无 dk 上限） + if (f16_io && !has_scan_f16_pipelines()) + NN_FAIL("scan_prefix_outer_gpu: f16 pipelines not available"); + // P4-2b：dk ≤ 32 时优先走窄共享变体（占用率高 6× ⇒ 32 个 WG 只排 ~1 个 + // 波次）；变体不可用（老 SPIR-V / 裁剪构建）→ 回落宽变体，数值逐位相同。 + const bool narrow = !generic && dk <= 32u; + const VulkanPipeline& pipe = + generic ? (f16_io ? pipelines_.get(PipelineId::scan_prefix_outer_gen_f16) + : pipelines_.get(PipelineId::scan_prefix_outer_gen)) + : (f16_io ? (narrow && has(PipelineId::scan_prefix_outer_d32_f16) + ? pipelines_.get(PipelineId::scan_prefix_outer_d32_f16) + : pipelines_.get(PipelineId::scan_prefix_outer_f16)) + : (narrow && has(PipelineId::scan_prefix_outer_d32) + ? pipelines_.get(PipelineId::scan_prefix_outer_d32) + : pipelines_.get(PipelineId::scan_prefix_outer))); + if (!generic && !has(PipelineId::scan_prefix_outer)) + NN_FAIL("scan_prefix_outer_gpu: pipeline not available"); + if (generic && !has(PipelineId::scan_prefix_outer_gen)) + NN_FAIL("scan_prefix_outer_gpu: generic pipeline not available"); + const auto rows = static_cast(K.rows()); + if (rows % (dk * heads) != 0) NN_FAIL("scan_prefix_outer_gpu: rows not divisible by H*dk"); + if (V.rows() != rows || V.cols() != K.cols() || P.rows() != rows || P.cols() != K.cols() || + R.rows() != rows || R.cols() != K.cols()) + NN_FAIL("scan_prefix_outer_gpu: K/V/P/R shape mismatch"); + if (has_state && + (A0.rows() != heads * dk || A0.cols() != dk || B0.rows() != heads * dk || B0.cols() != dk)) + NN_FAIL("scan_prefix_outer_gpu: A0/B0 must be (H*dk, dk)"); + const auto seq = static_cast(K.cols()); + if (has_bnd) { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - if (dk == 0u || heads == 0u) - NN_FAIL("scan_prefix_outer_gpu: dk/heads must be > 0"); - const bool generic = dk > 64u; // 通用路径（状态驻全局 scratch，无 dk 上限） - if (f16_io && !has_scan_f16_pipelines()) - NN_FAIL("scan_prefix_outer_gpu: f16 pipelines not available"); - // P4-2b：dk ≤ 32 时优先走窄共享变体（占用率高 6× ⇒ 32 个 WG 只排 ~1 个 - // 波次）；变体不可用（老 SPIR-V / 裁剪构建）→ 回落宽变体，数值逐位相同。 - const bool narrow = !generic && dk <= 32u; - const VulkanPipeline& pipe = generic - ? (f16_io ? pipelines_.get(PipelineId::scan_prefix_outer_gen_f16) : pipelines_.get(PipelineId::scan_prefix_outer_gen)) - : (f16_io - ? (narrow && has(PipelineId::scan_prefix_outer_d32_f16) - ? pipelines_.get(PipelineId::scan_prefix_outer_d32_f16) - : pipelines_.get(PipelineId::scan_prefix_outer_f16)) - : (narrow && has(PipelineId::scan_prefix_outer_d32) - ? pipelines_.get(PipelineId::scan_prefix_outer_d32) - : pipelines_.get(PipelineId::scan_prefix_outer))); - if (!generic && !has(PipelineId::scan_prefix_outer)) - NN_FAIL("scan_prefix_outer_gpu: pipeline not available"); - if (generic && !has(PipelineId::scan_prefix_outer_gen)) - NN_FAIL("scan_prefix_outer_gpu: generic pipeline not available"); - const auto rows = static_cast(K.rows()); - if (rows % (dk * heads) != 0) - NN_FAIL("scan_prefix_outer_gpu: rows not divisible by H*dk"); - if (V.rows() != rows || V.cols() != K.cols() || - P.rows() != rows || P.cols() != K.cols() || - R.rows() != rows || R.cols() != K.cols()) - NN_FAIL("scan_prefix_outer_gpu: K/V/P/R shape mismatch"); - if (has_state && (A0.rows() != heads * dk || A0.cols() != dk || - B0.rows() != heads * dk || B0.cols() != dk)) - NN_FAIL("scan_prefix_outer_gpu: A0/B0 must be (H*dk, dk)"); - const auto seq = static_cast(K.cols()); - if (has_bnd) - { - if (boundary.rows() != 1 || boundary.cols() != (rows / (dk * heads)) * seq) - NN_FAIL("scan_prefix_outer_gpu: boundary must be (1, B*seq)"); - } - const auto BH = rows / dk; - const std::size_t out_rows = static_cast(rows) * 5; - GpuTensor C; - if (f16_io) - { - auto f16r = GpuTensorF16::create_empty(out_rows, seq, *this); - NN_TRY_CHECK(f16r); - C = GpuTensor(f16r->shared_buffer(), out_rows, seq); - } - else - { - auto C_res = GpuTensor::create_empty(out_rows, seq, *this); - NN_TRY_CHECK(C_res); - C = std::move(*C_res); - } - struct PushPrefix { uint32_t dk, heads, seq, causal, has_state, has_bnd, rows; }; - PushPrefix push{dk, heads, seq, causal ? 1u : 0u, has_state ? 1u : 0u, - has_bnd ? 1u : 0u, rows}; - std::vector pc(sizeof(push)); - std::memcpy(pc.data(), &push, sizeof(push)); - if (generic) - { - // 通用 d_k > 64：状态驻全局 scratch（B*H*2*dk²），共享内存 O(1)， - // 对任意 dk 安全。GpuTensor 析构走 pending_destroys 延迟归还， - // batch 录制期被引用（descriptor set）也不会过早释放。 - auto St_res = GpuTensor::create_empty( - static_cast(BH) * 2u * dk * dk, 1u, *this); - NN_TRY_CHECK(St_res); - GpuTensor St = std::move(*St_res); - std::vector inputs{K, V, P, R, A0, B0, boundary, St}; - auto r = dispatch_compute(pipe, inputs, C, pc, 1u, BH, 1u, "scan_prefix"); - NN_TRY_CHECK(r); - return C; - } - std::vector inputs{K, V, P, R, A0, B0, boundary}; - auto r = dispatch_compute(pipe, inputs, C, pc, 1u, BH, 1u, "scan_prefix"); - NN_TRY_CHECK(r); - return C; + if (boundary.rows() != 1 || boundary.cols() != (rows / (dk * heads)) * seq) + NN_FAIL("scan_prefix_outer_gpu: boundary must be (1, B*seq)"); } - - // ── 后缀扫描 + matvec 读出：输出 (rows*3, seq) ───────────────────── - // f16_io 同 scan_prefix_outer_gpu（Phase C2c→C2f）。 - [[nodiscard]] Result scan_suffix_outer_gpu( - const GpuTensor& D, const GpuTensor& X, const GpuTensor& Y, - uint32_t dk, uint32_t heads, bool causal, - const GpuTensor& boundary, bool has_bnd, - bool f16_io = false) + const auto BH = rows / dk; + const std::size_t out_rows = static_cast(rows) * 5; + GpuTensor C; + if (f16_io) { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - if (dk == 0u || heads == 0u) - NN_FAIL("scan_suffix_outer_gpu: dk/heads must be > 0"); - const bool generic = dk > 64u; // 通用路径（状态驻全局 scratch，无 dk 上限） - if (f16_io && !has_scan_f16_pipelines()) - NN_FAIL("scan_suffix_outer_gpu: f16 pipelines not available"); - // 注意：**后缀扫描刻意不走窄共享变体**——实测它对 suffix 有害（D 的 - // dk²×seq 列式散布读使其为流量/争用受限：窄变体把每 SM 的驻留 WG 从 3 - // 提到 15，dk=32/BH=32 下 7.9 → 16.2 ms），故只给 scan_prefix 加变体。 - const VulkanPipeline& pipe = generic - ? (f16_io ? pipelines_.get(PipelineId::scan_suffix_outer_gen_f16) : pipelines_.get(PipelineId::scan_suffix_outer_gen)) - : (f16_io ? pipelines_.get(PipelineId::scan_suffix_outer_f16) : pipelines_.get(PipelineId::scan_suffix_outer)); - if (!generic && !has(PipelineId::scan_suffix_outer)) - NN_FAIL("scan_suffix_outer_gpu: pipeline not available"); - if (generic && !has(PipelineId::scan_suffix_outer_gen)) - NN_FAIL("scan_suffix_outer_gpu: generic pipeline not available"); - const auto rows = static_cast(X.rows()); - if (rows % (dk * heads) != 0) - NN_FAIL("scan_suffix_outer_gpu: X rows not divisible by H*dk"); - if (D.rows() != rows * dk || D.cols() != X.cols() || - Y.rows() != rows || Y.cols() != X.cols()) - NN_FAIL("scan_suffix_outer_gpu: D/X/Y shape mismatch"); - const auto seq = static_cast(X.cols()); - if (has_bnd) - { - if (boundary.rows() != 1 || boundary.cols() != (rows / (dk * heads)) * seq) - NN_FAIL("scan_suffix_outer_gpu: boundary must be (1, B*seq)"); - } - const auto BH = rows / dk; - const std::size_t out_rows = static_cast(rows) * 3; - GpuTensor C; - if (f16_io) - { - auto f16r = GpuTensorF16::create_empty(out_rows, seq, *this); - NN_TRY_CHECK(f16r); - C = GpuTensor(f16r->shared_buffer(), out_rows, seq); - } - else - { - auto C_res = GpuTensor::create_empty(out_rows, seq, *this); - NN_TRY_CHECK(C_res); - C = std::move(*C_res); - } - struct PushSuffix { uint32_t dk, heads, seq, causal, has_bnd, rows; }; - PushSuffix push{dk, heads, seq, causal ? 1u : 0u, has_bnd ? 1u : 0u, rows}; - std::vector pc(sizeof(push)); - std::memcpy(pc.data(), &push, sizeof(push)); - if (generic) - { - auto St_res = GpuTensor::create_empty( - static_cast(BH) * dk * dk, 1u, *this); - NN_TRY_CHECK(St_res); - GpuTensor St = std::move(*St_res); - std::vector inputs{D, X, Y, boundary, St}; - auto r = dispatch_compute(pipe, inputs, C, pc, 1u, BH, 1u, "scan_suffix"); - NN_TRY_CHECK(r); - return C; - } - std::vector inputs{D, X, Y, boundary}; - auto r = dispatch_compute(pipe, inputs, C, pc, 1u, BH, 1u); - NN_TRY_CHECK(r); - return C; + auto f16r = GpuTensorF16::create_empty(out_rows, seq, *this); + NN_TRY_CHECK(f16r); + C = GpuTensor(f16r->shared_buffer(), out_rows, seq); } - - // ── 逐列外积：输出 (rows*dk, seq) ────────────────────────────────── - // f16_io 同 scan_prefix_outer_gpu（Phase C2f）。 - [[nodiscard]] Result outer_col_gpu( - const GpuTensor& P, const GpuTensor& R, const GpuTensor& S, - uint32_t dk, bool has_scale, - bool f16_io = false) + else { - if (!initialized_) - NN_FAIL("GPU backend not initialized"); - if (!has(PipelineId::outer_col)) - NN_FAIL("outer_col_gpu: pipeline not available"); - if (f16_io && pipelines_.get(PipelineId::outer_col_f16).handle() == VK_NULL_HANDLE) - NN_FAIL("outer_col_gpu: f16 pipeline not available"); - const VulkanPipeline& pipe = f16_io ? pipelines_.get(PipelineId::outer_col_f16) : pipelines_.get(PipelineId::outer_col); - if (dk == 0u) - NN_FAIL("outer_col_gpu: dk must be > 0"); - const auto rows = static_cast(P.rows()); - if (rows % dk != 0) - NN_FAIL("outer_col_gpu: rows not divisible by dk"); - if (R.rows() != rows || R.cols() != P.cols()) - NN_FAIL("outer_col_gpu: P/R shape mismatch"); - if (has_scale && (S.rows() != rows || S.cols() != P.cols())) - NN_FAIL("outer_col_gpu: S must be (B*H*dk, seq)"); - const auto seq = static_cast(P.cols()); - const std::size_t out_rows = static_cast(rows) * dk; - GpuTensor C; - if (f16_io) - { - auto f16r = GpuTensorF16::create_empty(out_rows, seq, *this); - NN_TRY_CHECK(f16r); - C = GpuTensor(f16r->shared_buffer(), out_rows, seq); - } - else - { - auto C_res = GpuTensor::create_empty(out_rows, seq, *this); - NN_TRY_CHECK(C_res); - C = std::move(*C_res); - } - struct PushOuter { uint32_t dk, seq, rows, has_scale; }; - PushOuter push{dk, seq, rows, has_scale ? 1u : 0u}; - std::vector pc(sizeof(push)); - std::memcpy(pc.data(), &push, sizeof(push)); - std::vector inputs{P, R, S}; - const auto total = static_cast(rows) * dk * seq; - auto r = dispatch_compute(pipe, inputs, C, pc, - (total + 255u) / 256u, 1u, 1u); - NN_TRY_CHECK(r); - return C; + auto C_res = GpuTensor::create_empty(out_rows, seq, *this); + NN_TRY_CHECK(C_res); + C = std::move(*C_res); } - - // ══════════════════════════════════════════════════════════════════ - // 纯 GPU 原语方法（纯 GPU 架构核心） - // - // 所有方法遵循同一模式： - // 1. 分配输出 Tensor（如有输出） - // 2. 分配描述符集，写入 binding - // 3. 获取 command buffer（batch 共享 / 独立分配） - // 4. 录制：输入屏障 → bind pipeline → push constants → dispatch → 输出屏障 - // 5. 独立模式：end → submit → wait → cleanup - // ══════════════════════════════════════════════════════════════════ - - // ── 辅助：分配描述符集 ────────────────────────────────────────────── - [[nodiscard]] Result alloc_desc_set(VkDescriptorSetLayout layout) + struct PushPrefix + { uint32_t dk, heads, seq, causal, has_state, has_bnd, rows; }; + PushPrefix push{dk, heads, seq, causal ? 1u : 0u, has_state ? 1u : 0u, has_bnd ? 1u : 0u, rows}; + std::vector pc(sizeof(push)); + std::memcpy(pc.data(), &push, sizeof(push)); + if (generic) { - VkDescriptorSet desc_set = VK_NULL_HANDLE; - VkDescriptorSetAllocateInfo desc_alloc{}; - desc_alloc.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO; - desc_alloc.descriptorPool = gpu_tensor_pool_; - desc_alloc.descriptorSetCount = 1; - desc_alloc.pSetLayouts = &layout; - - auto r = detail::vk_check( - vkAllocateDescriptorSets(device_.device(), &desc_alloc, &desc_set), - __FILE__, __LINE__); + // 通用 d_k > 64：状态驻全局 scratch（B*H*2*dk²），共享内存 O(1)， + // 对任意 dk 安全。GpuTensor 析构走 pending_destroys 延迟归还， + // batch 录制期被引用（descriptor set）也不会过早释放。 + auto St_res = + GpuTensor::create_empty(static_cast(BH) * 2u * dk * dk, 1u, *this); + NN_TRY_CHECK(St_res); + GpuTensor St = std::move(*St_res); + std::vector inputs{K, V, P, R, A0, B0, boundary, St}; + auto r = dispatch_compute(pipe, inputs, C, pc, 1u, BH, 1u, "scan_prefix"); NN_TRY_CHECK(r); - // batch 模式：描述符集被当前帧的已录制命令引用，归属该帧， - // 帧的 fence 信号后（reap_frame）才释放 - if (batch_mode_) frames_[batch_frame_].desc_sets.push_back(desc_set); - return desc_set; + return C; } - - // ── 辅助：分配描述符集并把 buffers[i] 绑到 binding i ────────────────── - // 取代每个算子各写一遍的 "VkDescriptorBufferInfo[N] + VkWriteDescriptorSet[N] - // + vkUpdateDescriptorSets" 样板。失败路径自己归还集合（RV-15 语义）， - // 调用方只需 NN_TRY_CHECK。 - // - // 16 个以内走栈上数组（原语算子最多 4 个绑定，融合 extras 也在上限内）； - // 超过时退回堆——保持旧实现"支持任意 N"的行为不变。 - [[nodiscard]] Result alloc_desc_set_for( - VkDescriptorSetLayout layout, std::span buffers) + std::vector inputs{K, V, P, R, A0, B0, boundary}; + auto r = dispatch_compute(pipe, inputs, C, pc, 1u, BH, 1u, "scan_prefix"); + NN_TRY_CHECK(r); + return C; +} + +// ── 后缀扫描 + matvec 读出：输出 (rows*3, seq) ───────────────────── +// f16_io 同 scan_prefix_outer_gpu（Phase C2c→C2f）。 +[[nodiscard]] Result scan_suffix_outer_gpu(const GpuTensor& D, const GpuTensor& X, + const GpuTensor& Y, uint32_t dk, + uint32_t heads, bool causal, + const GpuTensor& boundary, bool has_bnd, + bool f16_io = false) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + if (dk == 0u || heads == 0u) NN_FAIL("scan_suffix_outer_gpu: dk/heads must be > 0"); + const bool generic = dk > 64u; // 通用路径（状态驻全局 scratch，无 dk 上限） + if (f16_io && !has_scan_f16_pipelines()) + NN_FAIL("scan_suffix_outer_gpu: f16 pipelines not available"); + // 注意：**后缀扫描刻意不走窄共享变体**——实测它对 suffix 有害（D 的 + // dk²×seq 列式散布读使其为流量/争用受限：窄变体把每 SM 的驻留 WG 从 3 + // 提到 15，dk=32/BH=32 下 7.9 → 16.2 ms），故只给 scan_prefix 加变体。 + const VulkanPipeline& pipe = + generic ? (f16_io ? pipelines_.get(PipelineId::scan_suffix_outer_gen_f16) + : pipelines_.get(PipelineId::scan_suffix_outer_gen)) + : (f16_io ? pipelines_.get(PipelineId::scan_suffix_outer_f16) + : pipelines_.get(PipelineId::scan_suffix_outer)); + if (!generic && !has(PipelineId::scan_suffix_outer)) + NN_FAIL("scan_suffix_outer_gpu: pipeline not available"); + if (generic && !has(PipelineId::scan_suffix_outer_gen)) + NN_FAIL("scan_suffix_outer_gpu: generic pipeline not available"); + const auto rows = static_cast(X.rows()); + if (rows % (dk * heads) != 0) NN_FAIL("scan_suffix_outer_gpu: X rows not divisible by H*dk"); + if (D.rows() != rows * dk || D.cols() != X.cols() || Y.rows() != rows || Y.cols() != X.cols()) + NN_FAIL("scan_suffix_outer_gpu: D/X/Y shape mismatch"); + const auto seq = static_cast(X.cols()); + if (has_bnd) { - auto ds_r = alloc_desc_set(layout); - if (!ds_r) - return ds_r; - const VkDescriptorSet desc_set = *ds_r; - if (buffers.empty()) - return desc_set; - - constexpr std::size_t kInline = 16; - VkDescriptorBufferInfo inline_infos[kInline]; - VkWriteDescriptorSet inline_writes[kInline]; - std::vector heap_infos; - std::vector heap_writes; - VkDescriptorBufferInfo* infos = inline_infos; - VkWriteDescriptorSet* writes = inline_writes; - if (buffers.size() > kInline) - { - heap_infos.resize(buffers.size()); - heap_writes.resize(buffers.size()); - infos = heap_infos.data(); - writes = heap_writes.data(); - } - - for (std::size_t i = 0; i < buffers.size(); ++i) - { - infos[i] = {buffers[i], 0, VK_WHOLE_SIZE}; - writes[i] = {}; - writes[i].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET; - writes[i].dstSet = desc_set; - writes[i].dstBinding = static_cast(i); - writes[i].descriptorCount = 1; - writes[i].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER; - writes[i].pBufferInfo = &infos[i]; - } - vkUpdateDescriptorSets(device_.device(), - static_cast(buffers.size()), writes, 0, nullptr); - return desc_set; + if (boundary.rows() != 1 || boundary.cols() != (rows / (dk * heads)) * seq) + NN_FAIL("scan_suffix_outer_gpu: boundary must be (1, B*seq)"); } - - // ── 辅助：提交前释放描述符集（**错误分支专用**，RV-15）────────────── - // 语义与 alloc_desc_set 的登记规则严格对齐： - // batch 模式：集合已在 frames_[batch_frame_].desc_sets 名下 → **此处不得 - // 释放**（同一句柄被 vkFreeDescriptorSets 两次 = 校验层 VUID 违规 / - // 驱动 UB），交给 reap_frame 统一回收（该帧命令并未绑定它，留着无害）。 - // solo 模式：集合未登记 → 必须立即释放（否则池配额泄漏）。 - // 所有"拿到 cmd 之前的错误分支"（`!cmd_r` / `!ds2_r` …）一律走本函数。 - void release_desc_set_now(VkDescriptorSet set) noexcept + const auto BH = rows / dk; + const std::size_t out_rows = static_cast(rows) * 3; + GpuTensor C; + if (f16_io) { - if (set == VK_NULL_HANDLE || batch_mode_) return; - vkFreeDescriptorSets(device_.device(), gpu_tensor_pool_, 1, &set); + auto f16r = GpuTensorF16::create_empty(out_rows, seq, *this); + NN_TRY_CHECK(f16r); + C = GpuTensor(f16r->shared_buffer(), out_rows, seq); } - - // ── 辅助：跨 submit 数据依赖（多帧流水线核心）──────────────────────── - // 背景：单队列 FIFO 只是执行顺序保证。实测本驱动（NVIDIA + Windows） - // 下，消费方 submit（download/matmul/batch 帧）紧跟上传 submit - // （<~2ms）时，消费方 GPU 操作会读到上传写入的旧值（零）；host 侧 - // sleep ≥5ms 或 fence 等待可规避——即隐式跨 submit 数据依赖不可靠， - // 必须用队列级原语显式建立： - // 上传 submit: pSignalSemaphores = {region 信号量} + value - // 消费 submit: pWaitSemaphores = {所有 in-flight region 信号量} - // pWaitSemaphoreValues = {各自的 value} - // pWaitDstStageMask = {ALL_COMMANDS}（waitSemaphoreCount>0 - // 时规范要求非空，见 - // VUID-VkSubmitInfo-pWaitDstStageMask） - // **必须用时间线信号量**：同一 value 可被任意多个 submit 等待且等待已达成 - // 的 value 是 no-op。二进制信号量一次 signal 只能被一个 wait 消费，第二 - // 个消费者会永久阻塞（VUID-vkQueueSubmit-pWaitSemaphores-03238），AMD - // 老驱动实测直接死锁。 - // host 永不阻塞（多帧流水线收益保留），GPU 在队列内等待数据就绪。 - // 设备不支持时间线信号量时不创建信号量，改由 drain_in_flight() 在 host - // 侧阻塞兜底（正确性优先）。 - struct StagingWait - { - std::vector sems; - std::vector stages; - std::vector values; - [[nodiscard]] bool any() const noexcept { return !sems.empty(); } - }; - [[nodiscard]] StagingWait collect_staging_waits() + else { - StagingWait w; - if (!staging_ring_) - return w; - if (!staging_ring_->timeline()) - { - // 回退路径：host 等在飞上传完成（阻塞），后续 submit 无跨 submit - // 依赖，无需信号量 - staging_ring_->drain_in_flight(); - return w; - } - for (std::size_t i = 0; i < staging_ring_->num_regions(); ++i) - { - if (!staging_ring_->in_flight(i)) - continue; - w.sems.push_back(staging_ring_->semaphore(i)); - w.stages.push_back(VK_PIPELINE_STAGE_ALL_COMMANDS_BIT); - w.values.push_back(staging_ring_->signal_value(i)); - } - return w; + auto C_res = GpuTensor::create_empty(out_rows, seq, *this); + NN_TRY_CHECK(C_res); + C = std::move(*C_res); } - - // 把跨 submit 等待挂到 submit 上（pNext = VkTimelineSemaphoreSubmitInfo） - static void attach_staging_waits( - VkSubmitInfo& si, const StagingWait& sw, VkTimelineSemaphoreSubmitInfo& tsi) + struct PushSuffix + { uint32_t dk, heads, seq, causal, has_bnd, rows; }; + PushSuffix push{dk, heads, seq, causal ? 1u : 0u, has_bnd ? 1u : 0u, rows}; + std::vector pc(sizeof(push)); + std::memcpy(pc.data(), &push, sizeof(push)); + if (generic) { - if (!sw.any()) - return; - si.waitSemaphoreCount = static_cast(sw.sems.size()); - si.pWaitSemaphores = sw.sems.data(); - si.pWaitDstStageMask = sw.stages.data(); - tsi.sType = VK_STRUCTURE_TYPE_TIMELINE_SEMAPHORE_SUBMIT_INFO; - tsi.waitSemaphoreValueCount = static_cast(sw.values.size()); - tsi.pWaitSemaphoreValues = sw.values.data(); - si.pNext = &tsi; + auto St_res = GpuTensor::create_empty(static_cast(BH) * dk * dk, 1u, *this); + NN_TRY_CHECK(St_res); + GpuTensor St = std::move(*St_res); + std::vector inputs{D, X, Y, boundary, St}; + auto r = dispatch_compute(pipe, inputs, C, pc, 1u, BH, 1u, "scan_suffix"); + NN_TRY_CHECK(r); + return C; } - - // 上传 submit 的信号挂载（时间线信号量模式）。 - // 返回 false = 当前设备非时间线模式：不 signal，消费者走 drain 回退。 - // 调用方提供的 sem_storage/value_storage 必须活到 vkQueueSubmit（ - // pSignalSemaphores/pSignalSemaphoreValues 是借用指针）。 - // 必须在 mark_in_flight() 之前调用：消费者只认"in-flight region 的 - // signal_value"，value 必须先写好。 - static bool attach_upload_signal( - VkSubmitInfo& si, StagingRing& ring, std::size_t region_idx, - VkTimelineSemaphoreSubmitInfo& tsi, - VkSemaphore& sem_storage, std::uint64_t& value_storage) + std::vector inputs{D, X, Y, boundary}; + auto r = dispatch_compute(pipe, inputs, C, pc, 1u, BH, 1u); + NN_TRY_CHECK(r); + return C; +} + +// ── 逐列外积：输出 (rows*dk, seq) ────────────────────────────────── +// f16_io 同 scan_prefix_outer_gpu（Phase C2f）。 +[[nodiscard]] Result outer_col_gpu(const GpuTensor& P, const GpuTensor& R, + const GpuTensor& S, uint32_t dk, bool has_scale, + bool f16_io = false) +{ + if (!initialized_) NN_FAIL("GPU backend not initialized"); + if (!has(PipelineId::outer_col)) NN_FAIL("outer_col_gpu: pipeline not available"); + if (f16_io && pipelines_.get(PipelineId::outer_col_f16).handle() == VK_NULL_HANDLE) + NN_FAIL("outer_col_gpu: f16 pipeline not available"); + const VulkanPipeline& pipe = + f16_io ? pipelines_.get(PipelineId::outer_col_f16) : pipelines_.get(PipelineId::outer_col); + if (dk == 0u) NN_FAIL("outer_col_gpu: dk must be > 0"); + const auto rows = static_cast(P.rows()); + if (rows % dk != 0) NN_FAIL("outer_col_gpu: rows not divisible by dk"); + if (R.rows() != rows || R.cols() != P.cols()) NN_FAIL("outer_col_gpu: P/R shape mismatch"); + if (has_scale && (S.rows() != rows || S.cols() != P.cols())) + NN_FAIL("outer_col_gpu: S must be (B*H*dk, seq)"); + const auto seq = static_cast(P.cols()); + const std::size_t out_rows = static_cast(rows) * dk; + GpuTensor C; + if (f16_io) { - if (!ring.timeline()) - return false; - value_storage = ring.next_signal_value(region_idx); - sem_storage = ring.semaphore(region_idx); - tsi.sType = VK_STRUCTURE_TYPE_TIMELINE_SEMAPHORE_SUBMIT_INFO; - tsi.signalSemaphoreValueCount = 1; - tsi.pSignalSemaphoreValues = &value_storage; - si.signalSemaphoreCount = 1; - si.pSignalSemaphores = &sem_storage; - si.pNext = &tsi; - return true; + auto f16r = GpuTensorF16::create_empty(out_rows, seq, *this); + NN_TRY_CHECK(f16r); + C = GpuTensor(f16r->shared_buffer(), out_rows, seq); } - - // ── 辅助：获取 command buffer（batch 或独立）────────────────────── - // ── 算子 host 开销剖析开关（NN_GPU_PROFILE=1 时 stderr 打印分段耗时； - // 未设置时进程内仅查询一次，之后纯 bool 读取）──────────────────── - [[nodiscard]] static bool profile_ops_enabled() + else { - static const bool on = []() { - // 复用 VulkanDevice::get_env（MSVC 下 _dupenv_s，绕开 getenv 弃用） - const std::string v = VulkanDevice::get_env("NN_GPU_PROFILE"); - return !v.empty() && v[0] != '0'; - }(); - return on; + auto C_res = GpuTensor::create_empty(out_rows, seq, *this); + NN_TRY_CHECK(C_res); + C = std::move(*C_res); } - - // ── P0 细 tile（wgrad 网格饥饿）逃生阀：NN_MM_NO_FINE=1 关闭 ──────────── - // A/B 走双进程同窗交错（bench/run_ab_env.ps1 口径），故进程内静态 - // 缓存即可——避免 run_fused_gpu 热路径每次派发都 getenv。 - [[nodiscard]] static bool mm_fine_disabled() + struct PushOuter + { uint32_t dk, seq, rows, has_scale; }; + PushOuter push{dk, seq, rows, has_scale ? 1u : 0u}; + std::vector pc(sizeof(push)); + std::memcpy(pc.data(), &push, sizeof(push)); + std::vector inputs{P, R, S}; + const auto total = static_cast(rows) * dk * seq; + auto r = dispatch_compute(pipe, inputs, C, pc, (total + 255u) / 256u, 1u, 1u); + NN_TRY_CHECK(r); + return C; +} + +// ══════════════════════════════════════════════════════════════════ +// 纯 GPU 原语方法（纯 GPU 架构核心） +// +// 所有方法遵循同一模式： +// 1. 分配输出 Tensor（如有输出） +// 2. 分配描述符集，写入 binding +// 3. 获取 command buffer（batch 共享 / 独立分配） +// 4. 录制：输入屏障 → bind pipeline → push constants → dispatch → 输出屏障 +// 5. 独立模式：end → submit → wait → cleanup +// ══════════════════════════════════════════════════════════════════ + +// ── 辅助：分配描述符集 ────────────────────────────────────────────── +[[nodiscard]] Result alloc_desc_set(VkDescriptorSetLayout layout) +{ + VkDescriptorSet desc_set = VK_NULL_HANDLE; + VkDescriptorSetAllocateInfo desc_alloc{}; + desc_alloc.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO; + desc_alloc.descriptorPool = gpu_tensor_pool_; + desc_alloc.descriptorSetCount = 1; + desc_alloc.pSetLayouts = &layout; + + auto r = detail::vk_check(vkAllocateDescriptorSets(device_.device(), &desc_alloc, &desc_set), + __FILE__, __LINE__); + NN_TRY_CHECK(r); + // batch 模式：描述符集被当前帧的已录制命令引用，归属该帧， + // 帧的 fence 信号后（reap_frame）才释放 + if (batch_mode_) frames_[batch_frame_].desc_sets.push_back(desc_set); + return desc_set; +} + +// ── 辅助：分配描述符集并把 buffers[i] 绑到 binding i ────────────────── +// 取代每个算子各写一遍的 "VkDescriptorBufferInfo[N] + VkWriteDescriptorSet[N] +// + vkUpdateDescriptorSets" 样板。失败路径自己归还集合（RV-15 语义）， +// 调用方只需 NN_TRY_CHECK。 +// +// 16 个以内走栈上数组（原语算子最多 4 个绑定，融合 extras 也在上限内）； +// 超过时退回堆——保持旧实现"支持任意 N"的行为不变。 +[[nodiscard]] Result alloc_desc_set_for(VkDescriptorSetLayout layout, + std::span buffers) +{ + auto ds_r = alloc_desc_set(layout); + if (!ds_r) return ds_r; + const VkDescriptorSet desc_set = *ds_r; + if (buffers.empty()) return desc_set; + + constexpr std::size_t kInline = 16; + VkDescriptorBufferInfo inline_infos[kInline]; + VkWriteDescriptorSet inline_writes[kInline]; + std::vector heap_infos; + std::vector heap_writes; + VkDescriptorBufferInfo* infos = inline_infos; + VkWriteDescriptorSet* writes = inline_writes; + if (buffers.size() > kInline) { - static const bool off = []() { - const std::string v = VulkanDevice::get_env("NN_MM_NO_FINE"); - return !v.empty() && v[0] != '0'; - }(); - return off; + heap_infos.resize(buffers.size()); + heap_writes.resize(buffers.size()); + infos = heap_infos.data(); + writes = heap_writes.data(); } - // ── P0 细粒度实验阀：NN_MM_FINE_NAIVE=1 → 饥饿形状改派 naive - // （matmul.comp：16×16 线程、每线程 1 输出、无 shared/barrier—— - // 线程数 = M×N（细 tiled 的 16×），靠 warp 切换吃满延迟）。 - // 与 tiled 细变体同窗交错 A/B 定夺最终形态。 - [[nodiscard]] static bool mm_fine_naive() + for (std::size_t i = 0; i < buffers.size(); ++i) { - static const bool on = []() { - const std::string v = VulkanDevice::get_env("NN_MM_FINE_NAIVE"); - return !v.empty() && v[0] != '0'; - }(); - return on; + infos[i] = {buffers[i], 0, VK_WHOLE_SIZE}; + writes[i] = {}; + writes[i].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET; + writes[i].dstSet = desc_set; + writes[i].dstBinding = static_cast(i); + writes[i].descriptorCount = 1; + writes[i].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER; + writes[i].pBufferInfo = &infos[i]; } - - // ── K 切分逃生阀 + 调参（wgrad 二次归约）─────────────────────────────── - // NN_MM_NO_KSPLIT=1 关闭 K 切分（回落细 tile / 64-tile 路径，A/B 基线）； - // NN_MM_KSPLIT_S= 强制段数 S（1=auto，调参用，钳到 [2, MAX_S]）。 - // 双进程同窗交错 A/B 口径下进程内静态缓存即可（与 mm_fine_disabled 同）。 - [[nodiscard]] static bool mm_ksplit_disabled() + vkUpdateDescriptorSets(device_.device(), static_cast(buffers.size()), writes, 0, + nullptr); + return desc_set; +} + +// ── 辅助：提交前释放描述符集（**错误分支专用**，RV-15）────────────── +// 语义与 alloc_desc_set 的登记规则严格对齐： +// batch 模式：集合已在 frames_[batch_frame_].desc_sets 名下 → **此处不得 +// 释放**（同一句柄被 vkFreeDescriptorSets 两次 = 校验层 VUID 违规 / +// 驱动 UB），交给 reap_frame 统一回收（该帧命令并未绑定它，留着无害）。 +// solo 模式：集合未登记 → 必须立即释放（否则池配额泄漏）。 +// 所有"拿到 cmd 之前的错误分支"（`!cmd_r` / `!ds2_r` …）一律走本函数。 +void release_desc_set_now(VkDescriptorSet set) noexcept +{ + if (set == VK_NULL_HANDLE || batch_mode_) return; + vkFreeDescriptorSets(device_.device(), gpu_tensor_pool_, 1, &set); +} + +// ── 辅助：跨 submit 数据依赖（多帧流水线核心）──────────────────────── +// 背景：单队列 FIFO 只是执行顺序保证。实测本驱动（NVIDIA + Windows） +// 下，消费方 submit（download/matmul/batch 帧）紧跟上传 submit +// （<~2ms）时，消费方 GPU 操作会读到上传写入的旧值（零）；host 侧 +// sleep ≥5ms 或 fence 等待可规避——即隐式跨 submit 数据依赖不可靠， +// 必须用队列级原语显式建立： +// 上传 submit: pSignalSemaphores = {region 信号量} + value +// 消费 submit: pWaitSemaphores = {所有 in-flight region 信号量} +// pWaitSemaphoreValues = {各自的 value} +// pWaitDstStageMask = {ALL_COMMANDS}（waitSemaphoreCount>0 +// 时规范要求非空，见 +// VUID-VkSubmitInfo-pWaitDstStageMask） +// **必须用时间线信号量**：同一 value 可被任意多个 submit 等待且等待已达成 +// 的 value 是 no-op。二进制信号量一次 signal 只能被一个 wait 消费，第二 +// 个消费者会永久阻塞（VUID-vkQueueSubmit-pWaitSemaphores-03238），AMD +// 老驱动实测直接死锁。 +// host 永不阻塞（多帧流水线收益保留），GPU 在队列内等待数据就绪。 +// 设备不支持时间线信号量时不创建信号量，改由 drain_in_flight() 在 host +// 侧阻塞兜底（正确性优先）。 +struct StagingWait +{ + std::vector sems; + std::vector stages; + std::vector values; + [[nodiscard]] bool any() const noexcept { return !sems.empty(); } +}; +[[nodiscard]] StagingWait collect_staging_waits() +{ + StagingWait w; + if (!staging_ring_) return w; + if (!staging_ring_->timeline()) { - static const bool off = []() { - const std::string v = VulkanDevice::get_env("NN_MM_NO_KSPLIT"); - return !v.empty() && v[0] != '0'; - }(); - return off; + // 回退路径：host 等在飞上传完成（阻塞），后续 submit 无跨 submit + // 依赖，无需信号量 + staging_ring_->drain_in_flight(); + return w; } - - [[nodiscard]] static std::uint32_t mm_ksplit_forced_s() + for (std::size_t i = 0; i < staging_ring_->num_regions(); ++i) { - static const std::uint32_t s = []() -> std::uint32_t { - const std::string v = VulkanDevice::get_env("NN_MM_KSPLIT_S"); - if (v.empty() || v[0] == '0') - return 0u; // auto - long n = 0; - for (char c : v) - { - if (c < '0' || c > '9') break; - n = n * 10 + (c - '0'); - if (n > 64) break; - } - return n >= 2 ? static_cast(n) : 0u; - }(); - return s; + if (!staging_ring_->in_flight(i)) continue; + w.sems.push_back(staging_ring_->semaphore(i)); + w.stages.push_back(VK_PIPELINE_STAGE_ALL_COMMANDS_BIT); + w.values.push_back(staging_ring_->signal_value(i)); } - - // ── 设备侧计时 helper（口径见成员注释）──────────────────────────── - // 录制块开始：reset 本块查询区间 + 游标复位（帧/solo 通用；上一块内容 - // 必已在其 fence 等待点 resolve，见 reap_frame / submit_and_wait） - void ts_begin_block(VkCommandBuffer cmd, TsCursor& cur) + return w; +} + +// 把跨 submit 等待挂到 submit 上（pNext = VkTimelineSemaphoreSubmitInfo） +static void attach_staging_waits(VkSubmitInfo& si, const StagingWait& sw, + VkTimelineSemaphoreSubmitInfo& tsi) +{ + if (!sw.any()) return; + si.waitSemaphoreCount = static_cast(sw.sems.size()); + si.pWaitSemaphores = sw.sems.data(); + si.pWaitDstStageMask = sw.stages.data(); + tsi.sType = VK_STRUCTURE_TYPE_TIMELINE_SEMAPHORE_SUBMIT_INFO; + tsi.waitSemaphoreValueCount = static_cast(sw.values.size()); + tsi.pWaitSemaphoreValues = sw.values.data(); + si.pNext = &tsi; +} + +// 上传 submit 的信号挂载（时间线信号量模式）。 +// 返回 false = 当前设备非时间线模式：不 signal，消费者走 drain 回退。 +// 调用方提供的 sem_storage/value_storage 必须活到 vkQueueSubmit（ +// pSignalSemaphores/pSignalSemaphoreValues 是借用指针）。 +// 必须在 mark_in_flight() 之前调用：消费者只认"in-flight region 的 +// signal_value"，value 必须先写好。 +static bool attach_upload_signal(VkSubmitInfo& si, StagingRing& ring, std::size_t region_idx, + VkTimelineSemaphoreSubmitInfo& tsi, VkSemaphore& sem_storage, + std::uint64_t& value_storage) +{ + if (!ring.timeline()) return false; + value_storage = ring.next_signal_value(region_idx); + sem_storage = ring.semaphore(region_idx); + tsi.sType = VK_STRUCTURE_TYPE_TIMELINE_SEMAPHORE_SUBMIT_INFO; + tsi.signalSemaphoreValueCount = 1; + tsi.pSignalSemaphoreValues = &value_storage; + si.signalSemaphoreCount = 1; + si.pSignalSemaphores = &sem_storage; + si.pNext = &tsi; + return true; +} + +// ── 辅助：获取 command buffer（batch 或独立）────────────────────── +// ── 算子 host 开销剖析开关（NN_GPU_PROFILE=1 时 stderr 打印分段耗时； +// 未设置时进程内仅查询一次，之后纯 bool 读取）──────────────────── +[[nodiscard]] static bool profile_ops_enabled() +{ + static const bool on = []() { - cur.next = 0; - cur.marks.clear(); - ts_cur_ = &cur; - if (!ts_enabled_) - return; - vkCmdResetQueryPool(cmd, ts_pool_, cur.base, kTsPairsPerBlock * 2); - } - - // dispatch 前后各写一个 timestamp（见成员注释的口径） - // detail：可选的 per-key 归因串（运行期 shader key / 原语变体名；传空跳过） - // —— 录制期驻留进 prof 稳定池，resolve 时与 label@stage 一起回填。 - // detail_suffix：拼在 detail 后的静态后缀（细 tile 变体 "#f"；调用点零分配） - void ts_dispatch(VkCommandBuffer cmd, const char* label, - std::uint32_t gx, std::uint32_t gy, std::uint32_t gz, - std::string_view detail = {}, - const char* detail_suffix = nullptr) + // 复用 VulkanDevice::get_env（MSVC 下 _dupenv_s，绕开 getenv 弃用） + const std::string v = VulkanDevice::get_env("NN_GPU_PROFILE"); + return !v.empty() && v[0] != '0'; + }(); + return on; +} + +// ── P0 细 tile（wgrad 网格饥饿）逃生阀：NN_MM_NO_FINE=1 关闭 ──────────── +// A/B 走双进程同窗交错（bench/run_ab_env.ps1 口径），故进程内静态 +// 缓存即可——避免 run_fused_gpu 热路径每次派发都 getenv。 +[[nodiscard]] static bool mm_fine_disabled() +{ + static const bool off = []() { - const bool mark = ts_enabled_ && ts_cur_ != nullptr && - ts_cur_->next < kTsPairsPerBlock; - std::uint32_t q = 0; - if (mark) + const std::string v = VulkanDevice::get_env("NN_MM_NO_FINE"); + return !v.empty() && v[0] != '0'; + }(); + return off; +} + +// ── P0 细粒度实验阀：NN_MM_FINE_NAIVE=1 → 饥饿形状改派 naive +// （matmul.comp：16×16 线程、每线程 1 输出、无 shared/barrier—— +// 线程数 = M×N（细 tiled 的 16×），靠 warp 切换吃满延迟）。 +// 与 tiled 细变体同窗交错 A/B 定夺最终形态。 +[[nodiscard]] static bool mm_fine_naive() +{ + static const bool on = []() + { + const std::string v = VulkanDevice::get_env("NN_MM_FINE_NAIVE"); + return !v.empty() && v[0] != '0'; + }(); + return on; +} + +// ── K 切分逃生阀 + 调参（wgrad 二次归约）─────────────────────────────── +// NN_MM_NO_KSPLIT=1 关闭 K 切分（回落细 tile / 64-tile 路径，A/B 基线）； +// NN_MM_KSPLIT_S= 强制段数 S（1=auto，调参用，钳到 [2, MAX_S]）。 +// 双进程同窗交错 A/B 口径下进程内静态缓存即可（与 mm_fine_disabled 同）。 +[[nodiscard]] static bool mm_ksplit_disabled() +{ + static const bool off = []() + { + const std::string v = VulkanDevice::get_env("NN_MM_NO_KSPLIT"); + return !v.empty() && v[0] != '0'; + }(); + return off; +} + +[[nodiscard]] static std::uint32_t mm_ksplit_forced_s() +{ + static const std::uint32_t s = []() -> std::uint32_t + { + const std::string v = VulkanDevice::get_env("NN_MM_KSPLIT_S"); + if (v.empty() || v[0] == '0') return 0u; // auto + long n = 0; + for (char c : v) { - q = ts_cur_->base + ts_cur_->next * 2; - vkCmdWriteTimestamp(cmd, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, ts_pool_, q); - // 归因到录制期 stage（resolve 时 scope 已出栈）：key = "label@stage"， - // 使设备 kernel 表可与 host 阶段表同维度对账（仅剖析开启时进入） - const char* stage = nn::prof::current_stage(); - const char* key = (stage != nullptr && stage[0] != '\0') - ? nn::prof::interned_pair(label, stage) - : label; - // per-key 细分（"detail[suffix]@label@stage"；detail 空 = 不入细分表） - const char* d = detail.empty() ? nullptr - : nn::prof::interned_detail(detail, detail_suffix, label, stage); - ts_cur_->marks.push_back(TsMark{key, d, ts_cur_->next}); - ++ts_cur_->next; + if (c < '0' || c > '9') break; + n = n * 10 + (c - '0'); + if (n > 64) break; } - vkCmdDispatch(cmd, gx, gy, gz); - if (mark) - vkCmdWriteTimestamp(cmd, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, ts_pool_, q + 1); + return n >= 2 ? static_cast(n) : 0u; + }(); + return s; +} + +// ── 设备侧计时 helper（口径见成员注释）──────────────────────────── +// 录制块开始：reset 本块查询区间 + 游标复位（帧/solo 通用；上一块内容 +// 必已在其 fence 等待点 resolve，见 reap_frame / submit_and_wait） +void ts_begin_block(VkCommandBuffer cmd, TsCursor& cur) +{ + cur.next = 0; + cur.marks.clear(); + ts_cur_ = &cur; + if (!ts_enabled_) return; + vkCmdResetQueryPool(cmd, ts_pool_, cur.base, kTsPairsPerBlock * 2); +} + +// dispatch 前后各写一个 timestamp（见成员注释的口径） +// detail：可选的 per-key 归因串（运行期 shader key / 原语变体名；传空跳过） +// —— 录制期驻留进 prof 稳定池，resolve 时与 label@stage 一起回填。 +// detail_suffix：拼在 detail 后的静态后缀（细 tile 变体 "#f"；调用点零分配） +void ts_dispatch(VkCommandBuffer cmd, const char* label, std::uint32_t gx, std::uint32_t gy, + std::uint32_t gz, std::string_view detail = {}, + const char* detail_suffix = nullptr) +{ + const bool mark = ts_enabled_ && ts_cur_ != nullptr && ts_cur_->next < kTsPairsPerBlock; + std::uint32_t q = 0; + if (mark) + { + q = ts_cur_->base + ts_cur_->next * 2; + vkCmdWriteTimestamp(cmd, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, ts_pool_, q); + // 归因到录制期 stage（resolve 时 scope 已出栈）：key = "label@stage"， + // 使设备 kernel 表可与 host 阶段表同维度对账（仅剖析开启时进入） + const char* stage = nn::prof::current_stage(); + const char* key = + (stage != nullptr && stage[0] != '\0') ? nn::prof::interned_pair(label, stage) : label; + // per-key 细分（"detail[suffix]@label@stage"；detail 空 = 不入细分表） + const char* d = detail.empty() + ? nullptr + : nn::prof::interned_detail(detail, detail_suffix, label, stage); + ts_cur_->marks.push_back(TsMark{key, d, ts_cur_->next}); + ++ts_cur_->next; } - - // resolve：调用点 fence 已等 → 读回本块全部 pair → 按标号聚合 - void ts_resolve(TsCursor& cur) + vkCmdDispatch(cmd, gx, gy, gz); + if (mark) vkCmdWriteTimestamp(cmd, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, ts_pool_, q + 1); +} + +// resolve：调用点 fence 已等 → 读回本块全部 pair → 按标号聚合 +void ts_resolve(TsCursor& cur) +{ + if (!ts_enabled_ || cur.next == 0 || cur.marks.empty()) { - if (!ts_enabled_ || cur.next == 0 || cur.marks.empty()) - { - cur.next = 0; - cur.marks.clear(); - return; - } - std::vector ts(cur.next * 2); - const VkResult r = vkGetQueryPoolResults(device_.device(), ts_pool_, - cur.base, cur.next * 2, ts.size() * sizeof(std::uint64_t), - ts.data(), sizeof(std::uint64_t), - VK_QUERY_RESULT_64_BIT | VK_QUERY_RESULT_WAIT_BIT); - if (r == VK_SUCCESS) - { - for (const auto& m : cur.marks) - { - const std::uint32_t p = m.pair; - const std::uint64_t t0 = ts[static_cast(p) * 2]; - const std::uint64_t t1 = ts[static_cast(p) * 2 + 1]; - if (t1 >= t0) - nn::prof::device_add(m.key, - static_cast( - static_cast(t1 - t0) * ts_period_ns_), - m.detail); - } - } cur.next = 0; cur.marks.clear(); + return; } - - // 返回 (cmd, owns_cmd)。owns_cmd=true 时调用方需 submit_and_wait。 - [[nodiscard]] Result> acquire_cmd() + std::vector ts(cur.next * 2); + const VkResult r = vkGetQueryPoolResults( + device_.device(), ts_pool_, cur.base, cur.next * 2, ts.size() * sizeof(std::uint64_t), + ts.data(), sizeof(std::uint64_t), VK_QUERY_RESULT_64_BIT | VK_QUERY_RESULT_WAIT_BIT); + if (r == VK_SUCCESS) { - if (batch_mode_) + for (const auto& m : cur.marks) { - batch_has_ops_ = true; // 当前帧已含 op（空帧不提交） - ts_cur_ = &frames_[batch_frame_].ts; // 设备侧计时：归属当前帧块 - return std::make_pair(batch_cmd_, false); + const std::uint32_t p = m.pair; + const std::uint64_t t0 = ts[static_cast(p) * 2]; + const std::uint64_t t1 = ts[static_cast(p) * 2 + 1]; + if (t1 >= t0) + nn::prof::device_add( + m.key, static_cast(static_cast(t1 - t0) * ts_period_ns_), + m.detail); } + } + cur.next = 0; + cur.marks.clear(); +} + +// 返回 (cmd, owns_cmd)。owns_cmd=true 时调用方需 submit_and_wait。 +[[nodiscard]] Result> acquire_cmd() +{ + if (batch_mode_) + { + batch_has_ops_ = true; // 当前帧已含 op（空帧不提交） + ts_cur_ = &frames_[batch_frame_].ts; // 设备侧计时：归属当前帧块 + return std::make_pair(batch_cmd_, false); + } - // 独立模式复用 initialize() 预分配的 solo_cmd_：reset + begin， - // 免去每算子 vkAllocateCommandBuffers/vkFreeCommandBuffers。 - // solo_cmd_ 仅在 submit_and_wait 等完 fence 后才会被再次取用， - // 不存在 pending 期 reset；begin 失败时 buffer 留在 reset 后初态 - // （规范允许 reset 从错误态恢复），下次调用可重试。 - VkCommandBuffer cmd = solo_cmd_; - if (cmd == VK_NULL_HANDLE) - NN_FAIL("acquire_cmd: solo_cmd_ not initialized"); - auto r = detail::vk_check(vkResetCommandBuffer(cmd, 0), __FILE__, __LINE__); - NN_TRY_CHECK(r); - VkCommandBufferBeginInfo begin_info{}; - begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO; - begin_info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT; - r = detail::vk_check(vkBeginCommandBuffer(cmd, &begin_info), __FILE__, __LINE__); + // 独立模式复用 initialize() 预分配的 solo_cmd_：reset + begin， + // 免去每算子 vkAllocateCommandBuffers/vkFreeCommandBuffers。 + // solo_cmd_ 仅在 submit_and_wait 等完 fence 后才会被再次取用， + // 不存在 pending 期 reset；begin 失败时 buffer 留在 reset 后初态 + // （规范允许 reset 从错误态恢复），下次调用可重试。 + VkCommandBuffer cmd = solo_cmd_; + if (cmd == VK_NULL_HANDLE) NN_FAIL("acquire_cmd: solo_cmd_ not initialized"); + auto r = detail::vk_check(vkResetCommandBuffer(cmd, 0), __FILE__, __LINE__); + NN_TRY_CHECK(r); + VkCommandBufferBeginInfo begin_info{}; + begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO; + begin_info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT; + r = detail::vk_check(vkBeginCommandBuffer(cmd, &begin_info), __FILE__, __LINE__); + NN_TRY_CHECK(r); + ts_begin_block(cmd, ts_solo_); // 设备侧计时：solo 块开始 + return std::make_pair(cmd, true); +} + +// ── 辅助：独立模式提交+等待（solo fence/cmd 复用，仅归还描述符集）──── +// cmd：acquire_cmd 返回的 solo_cmd_（非 batch 模式才会走到这里） +// desc_set 可为 VK_NULL_HANDLE（fill_zero/copy_buffer 无描述符集） +// NN_GPU_PROFILE=1：stderr 打印 host 分段耗时（µs）用于开销归因 +[[nodiscard]] Result submit_and_wait(VkCommandBuffer cmd, VkDescriptorSet desc_set) +{ + const bool prof = profile_ops_enabled(); + const auto t_begin = std::chrono::steady_clock::now(); + + auto r = detail::vk_check(vkEndCommandBuffer(cmd), __FILE__, __LINE__); + if (!r) + { + if (desc_set != VK_NULL_HANDLE) + vkFreeDescriptorSets(device_.device(), gpu_tensor_pool_, 1, &desc_set); NN_TRY_CHECK(r); - ts_begin_block(cmd, ts_solo_); // 设备侧计时：solo 块开始 - return std::make_pair(cmd, true); } + const auto t_end = std::chrono::steady_clock::now(); - // ── 辅助：独立模式提交+等待（solo fence/cmd 复用，仅归还描述符集）──── - // cmd：acquire_cmd 返回的 solo_cmd_（非 batch 模式才会走到这里） - // desc_set 可为 VK_NULL_HANDLE（fill_zero/copy_buffer 无描述符集） - // NN_GPU_PROFILE=1：stderr 打印 host 分段耗时（µs）用于开销归因 - [[nodiscard]] Result submit_and_wait( - VkCommandBuffer cmd, VkDescriptorSet desc_set) + if (solo_fence_ == VK_NULL_HANDLE) + { + if (desc_set != VK_NULL_HANDLE) + vkFreeDescriptorSets(device_.device(), gpu_tensor_pool_, 1, &desc_set); + NN_FAIL("submit_and_wait: solo_fence_ not initialized"); + } + // 上一次 submit 已 wait 到 signaled → 这里 reset 回 unsignaled 复用 + //（首次调用时是 init 创建的 fresh unsignaled，reset 为合法 no-op） + r = detail::vk_check(vkResetFences(device_.device(), 1, &solo_fence_), __FILE__, __LINE__); + if (!r) { - const bool prof = profile_ops_enabled(); - const auto t_begin = std::chrono::steady_clock::now(); + if (desc_set != VK_NULL_HANDLE) + vkFreeDescriptorSets(device_.device(), gpu_tensor_pool_, 1, &desc_set); + NN_TRY_CHECK(r); + } - auto r = detail::vk_check(vkEndCommandBuffer(cmd), __FILE__, __LINE__); - if (!r) - { - if (desc_set != VK_NULL_HANDLE) - vkFreeDescriptorSets(device_.device(), gpu_tensor_pool_, 1, &desc_set); - NN_TRY_CHECK(r); - } - const auto t_end = std::chrono::steady_clock::now(); + // 跨 submit 数据依赖：输入 buffer 可能刚由在飞上传的 copy 写入 + // （from_matrix 非阻塞提交后本原语立即执行）——等 in-flight region + // 信号量。host 不阻塞；GPU 在队列内等数据就绪。 + const auto sw = collect_staging_waits(); - if (solo_fence_ == VK_NULL_HANDLE) - { - if (desc_set != VK_NULL_HANDLE) - vkFreeDescriptorSets(device_.device(), gpu_tensor_pool_, 1, &desc_set); - NN_FAIL("submit_and_wait: solo_fence_ not initialized"); - } - // 上一次 submit 已 wait 到 signaled → 这里 reset 回 unsignaled 复用 - //（首次调用时是 init 创建的 fresh unsignaled，reset 为合法 no-op） - r = detail::vk_check( - vkResetFences(device_.device(), 1, &solo_fence_), __FILE__, __LINE__); + { + std::lock_guard lock(queue_mutex_); + VkSubmitInfo submit_info{}; + submit_info.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO; + submit_info.commandBufferCount = 1; + submit_info.pCommandBuffers = &cmd; + VkTimelineSemaphoreSubmitInfo timeline_info{}; + attach_staging_waits(submit_info, sw, timeline_info); + r = detail::vk_check(vkQueueSubmit(device_.compute_queue(), 1, &submit_info, solo_fence_), + __FILE__, __LINE__); if (!r) { if (desc_set != VK_NULL_HANDLE) vkFreeDescriptorSets(device_.device(), gpu_tensor_pool_, 1, &desc_set); NN_TRY_CHECK(r); } + } + const auto t_submit = std::chrono::steady_clock::now(); - // 跨 submit 数据依赖：输入 buffer 可能刚由在飞上传的 copy 写入 - // （from_matrix 非阻塞提交后本原语立即执行）——等 in-flight region - // 信号量。host 不阻塞；GPU 在队列内等数据就绪。 - const auto sw = collect_staging_waits(); - - { - std::lock_guard lock(queue_mutex_); - VkSubmitInfo submit_info{}; - submit_info.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO; - submit_info.commandBufferCount = 1; - submit_info.pCommandBuffers = &cmd; - VkTimelineSemaphoreSubmitInfo timeline_info{}; - attach_staging_waits(submit_info, sw, timeline_info); - r = detail::vk_check( - vkQueueSubmit(device_.compute_queue(), 1, &submit_info, solo_fence_), - __FILE__, __LINE__); - if (!r) - { - if (desc_set != VK_NULL_HANDLE) - vkFreeDescriptorSets(device_.device(), gpu_tensor_pool_, 1, &desc_set); - NN_TRY_CHECK(r); - } - } - const auto t_submit = std::chrono::steady_clock::now(); - - // 30 秒超时（单个原语，比 batch 短） - constexpr uint64_t kSingleOpTimeoutNs = 30'000'000'000ULL; - r = detail::vk_check( - vkWaitForFences(device_.device(), 1, &solo_fence_, VK_TRUE, kSingleOpTimeoutNs), - __FILE__, __LINE__); - const auto t_wait = std::chrono::steady_clock::now(); - - ts_resolve(ts_solo_); // 设备侧计时：solo 块 resolve（fence 已等） + // 30 秒超时（单个原语，比 batch 短） + constexpr uint64_t kSingleOpTimeoutNs = 30'000'000'000ULL; + r = detail::vk_check( + vkWaitForFences(device_.device(), 1, &solo_fence_, VK_TRUE, kSingleOpTimeoutNs), __FILE__, + __LINE__); + const auto t_wait = std::chrono::steady_clock::now(); - // 只归还描述符集；solo fence/cmd 留待下次复用（不销毁/不释放） - if (desc_set != VK_NULL_HANDLE) - vkFreeDescriptorSets(device_.device(), gpu_tensor_pool_, 1, &desc_set); - const auto t_done = std::chrono::steady_clock::now(); - - if (prof) - { - const auto us = [](auto a, auto b) { - return static_cast( - std::chrono::duration_cast(b - a).count()); - }; - std::fprintf(stderr, - "[gpu-profile] solo_op: end=%lldus submit=%lldus wait=%lldus " - "cleanup=%lldus total=%lldus\n", - us(t_begin, t_end), us(t_end, t_submit), us(t_submit, t_wait), - us(t_wait, t_done), us(t_begin, t_done)); - } - return r; - } + ts_resolve(ts_solo_); // 设备侧计时：solo 块 resolve（fence 已等） - // ── 辅助：录制输入屏障（多个输入 buffer）────────────────────────── - void record_input_barriers(VkCommandBuffer cmd, - std::initializer_list buffers) - { - record_input_barriers(cmd, - std::span(buffers.begin(), buffers.size())); - } + // 只归还描述符集；solo fence/cmd 留待下次复用（不销毁/不释放） + if (desc_set != VK_NULL_HANDLE) + vkFreeDescriptorSets(device_.device(), gpu_tensor_pool_, 1, &desc_set); + const auto t_done = std::chrono::steady_clock::now(); - void record_input_barriers(VkCommandBuffer cmd, - std::span buffers) + if (prof) { - std::vector barriers; - barriers.reserve(buffers.size()); - for (auto buf : buffers) + const auto us = [](auto a, auto b) { - VkBufferMemoryBarrier b{}; - b.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER; - b.srcAccessMask = VK_ACCESS_MEMORY_WRITE_BIT; - b.dstAccessMask = VK_ACCESS_SHADER_READ_BIT; - b.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - b.buffer = buf; - b.offset = 0; - b.size = VK_WHOLE_SIZE; - barriers.push_back(b); - } - vkCmdPipelineBarrier(cmd, - VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, - VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, - 0, 0, nullptr, - static_cast(barriers.size()), barriers.data(), - 0, nullptr); + return static_cast( + std::chrono::duration_cast(b - a).count()); + }; + std::fprintf(stderr, + "[gpu-profile] solo_op: end=%lldus submit=%lldus wait=%lldus " + "cleanup=%lldus total=%lldus\n", + us(t_begin, t_end), us(t_end, t_submit), us(t_submit, t_wait), + us(t_wait, t_done), us(t_begin, t_done)); } - - // ── 辅助：录制输出屏障 ───────────────────────────────────────────── - void record_output_barrier(VkCommandBuffer cmd, VkBuffer output_buf) + return r; +} + +// ── 辅助：录制输入屏障（多个输入 buffer）────────────────────────── +void record_input_barriers(VkCommandBuffer cmd, std::initializer_list buffers) +{ record_input_barriers(cmd, std::span(buffers.begin(), buffers.size())); } + +void record_input_barriers(VkCommandBuffer cmd, std::span buffers) +{ + std::vector barriers; + barriers.reserve(buffers.size()); + for (auto buf : buffers) { VkBufferMemoryBarrier b{}; b.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER; - b.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT; - b.dstAccessMask = VK_ACCESS_MEMORY_READ_BIT; + b.srcAccessMask = VK_ACCESS_MEMORY_WRITE_BIT; + b.dstAccessMask = VK_ACCESS_SHADER_READ_BIT; b.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - b.buffer = output_buf; + b.buffer = buf; b.offset = 0; b.size = VK_WHOLE_SIZE; - vkCmdPipelineBarrier(cmd, - VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, - VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, - 0, 0, nullptr, 1, &b, 0, nullptr); + barriers.push_back(b); } - + vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, + VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 0, nullptr, + static_cast(barriers.size()), barriers.data(), 0, nullptr); +} + +// ── 辅助：录制输出屏障 ───────────────────────────────────────────── +void record_output_barrier(VkCommandBuffer cmd, VkBuffer output_buf) +{ + VkBufferMemoryBarrier b{}; + b.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER; + b.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT; + b.dstAccessMask = VK_ACCESS_MEMORY_READ_BIT; + b.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; + b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; + b.buffer = output_buf; + b.offset = 0; + b.size = VK_WHOLE_SIZE; + vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, + VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, 0, + 0, nullptr, 1, &b, 0, nullptr); +} diff --git a/include/neuralnet.cpp/backend/compute_vk_device.hpp b/include/neuralnet.cpp/backend/compute_vk_device.hpp index 4273de08..71daae68 100644 --- a/include/neuralnet.cpp/backend/compute_vk_device.hpp +++ b/include/neuralnet.cpp/backend/compute_vk_device.hpp @@ -33,24 +33,22 @@ namespace detail { if (res != VK_SUCCESS) { - NN_FAIL("Vulkan error " + std::to_string(static_cast(res)) + " at " + std::string(file) + ":" + std::to_string(line)); + NN_FAIL("Vulkan error " + std::to_string(static_cast(res)) + " at " + + std::string(file) + ":" + std::to_string(line)); } return {}; } // ── CPU Scalar ↔ GPU float 转换 ───────────────────────────────────────── -inline void convert_scalar_to_float( - std::span src, float* __restrict dst) noexcept +inline void convert_scalar_to_float(std::span src, float* __restrict dst) noexcept { if constexpr (std::is_same_v) std::memcpy(dst, src.data(), src.size() * sizeof(float)); else - std::transform(src.begin(), src.end(), dst, - [](Scalar v) { return static_cast(v); }); + std::transform(src.begin(), src.end(), dst, [](Scalar v) { return static_cast(v); }); } -inline void convert_float_to_scalar( - const float* __restrict src, std::span dst) noexcept +inline void convert_float_to_scalar(const float* __restrict src, std::span dst) noexcept { if constexpr (std::is_same_v) std::memcpy(dst.data(), src, dst.size() * sizeof(float)); @@ -117,10 +115,8 @@ class VulkanDevice ~VulkanDevice() { - if (device_ != VK_NULL_HANDLE) - vkDestroyDevice(device_, nullptr); - if (instance_ != VK_NULL_HANDLE) - vkDestroyInstance(instance_, nullptr); + if (device_ != VK_NULL_HANDLE) vkDestroyDevice(device_, nullptr); + if (instance_ != VK_NULL_HANDLE) vkDestroyInstance(instance_, nullptr); } // 禁止拷贝和移动 @@ -131,24 +127,18 @@ class VulkanDevice // 读取环境变量（RV-09：实现收敛到 core_env.hpp；本包装保留供后端调用） [[nodiscard]] static std::string get_env(const char* name) - { - return ::nn::detail::env_string_raw(name); - } + { return ::nn::detail::env_string_raw(name); } // ── 手动指定计算设备（必须在 initialize() 之前调用）──────────────────── // selector = 枚举索引（"2"）或设备名子串（"40HX" / "NVIDIA"）； // 空字符串 = 自动选择（设备类型 + apiVersion 打分）。 // 优先级高于环境变量 NN_VULKAN_DEVICE；未命中时 initialize() 返回 // 列出全部候选设备的错误。 - void set_device_selector(std::string selector) - { - device_selector_ = std::move(selector); - } + void set_device_selector(std::string selector) { device_selector_ = std::move(selector); } [[nodiscard]] Result initialize() { - if (initialized_) - return {}; + if (initialized_) return {}; // 1. 创建 VkInstance // 实例版本请求到 1.2：时间线信号量（跨 submit 依赖的正确原语）是 @@ -159,8 +149,7 @@ class VulkanDevice vkGetInstanceProcAddr(nullptr, "vkEnumerateInstanceVersion"))) { uint32_t supported = VK_API_VERSION_1_0; - if (enum_version(&supported) == VK_SUCCESS) - loader_version = supported; + if (enum_version(&supported) == VK_SUCCESS) loader_version = supported; } VkApplicationInfo app_info{}; @@ -176,13 +165,11 @@ class VulkanDevice instance_info.pApplicationInfo = &app_info; VkResult res = vkCreateInstance(&instance_info, nullptr, &instance_); - if (res != VK_SUCCESS) - NN_FAIL("vkCreateInstance failed: " + std::to_string(res)); + if (res != VK_SUCCESS) NN_FAIL("vkCreateInstance failed: " + std::to_string(res)); // 2. 选择物理设备 uint32_t device_count = 0; vkEnumeratePhysicalDevices(instance_, &device_count, nullptr); - if (device_count == 0) - NN_FAIL("No Vulkan devices found"); + if (device_count == 0) NN_FAIL("No Vulkan devices found"); std::vector devices(device_count); vkEnumeratePhysicalDevices(instance_, &device_count, devices.data()); @@ -202,8 +189,7 @@ class VulkanDevice // 显式 API（set_device_selector / GpuBackend::initialize(sel)） // > 环境变量 NN_VULKAN_DEVICE > 自动打分（见下） std::string selector = device_selector_; - if (selector.empty()) - selector = get_env("NN_VULKAN_DEVICE"); + if (selector.empty()) selector = get_env("NN_VULKAN_DEVICE"); // selector = 枚举索引（"2"）或设备名称子串（"40HX"/"NVIDIA"） if (!selector.empty()) @@ -211,16 +197,14 @@ class VulkanDevice const std::string_view want(selector); // 纯数字只按索引匹配：否则 "2" 会命中 "R5 240" 这类名称子串 const bool want_is_index = - !want.empty() && - want.find_first_not_of("0123456789") == std::string_view::npos; + !want.empty() && want.find_first_not_of("0123456789") == std::string_view::npos; for (std::size_t i = 0; i < devices.size(); ++i) { VkPhysicalDeviceProperties props; vkGetPhysicalDeviceProperties(devices[i], &props); const std::string_view name(props.deviceName); - const bool matched = want_is_index - ? (want == std::string_view(std::to_string(i))) - : (name.find(want) != std::string_view::npos); + const bool matched = want_is_index ? (want == std::string_view(std::to_string(i))) + : (name.find(want) != std::string_view::npos); if (matched) { physical_device_ = devices[i]; @@ -252,16 +236,13 @@ class VulkanDevice const std::string_view name(props.deviceName); int rank = 1; - if (props.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU) - rank = 3; - else if (props.deviceType == VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU) - rank = 2; + if (props.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU) rank = 3; + else if (props.deviceType == VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU) rank = 2; if (name.rfind("Microsoft Direct3D12", 0) == 0 || name == "Microsoft Basic Render Driver") rank = 0; - if (rank > best_rank || - (rank == best_rank && props.apiVersion > best_api)) + if (rank > best_rank || (rank == best_rank && props.apiVersion > best_api)) { best_rank = rank; best_api = props.apiVersion; @@ -270,8 +251,7 @@ class VulkanDevice } } // 兜底：无可打分设备时使用第一个设备 - if (physical_device_ == VK_NULL_HANDLE) - physical_device_ = devices[0]; + if (physical_device_ == VK_NULL_HANDLE) physical_device_ = devices[0]; // 记录所选设备名（诊断"到底跑在哪块卡上"用） { @@ -286,7 +266,8 @@ class VulkanDevice vkGetPhysicalDeviceQueueFamilyProperties(physical_device_, &queue_family_count, nullptr); std::vector queue_families(queue_family_count); - vkGetPhysicalDeviceQueueFamilyProperties(physical_device_, &queue_family_count, queue_families.data()); + vkGetPhysicalDeviceQueueFamilyProperties(physical_device_, &queue_family_count, + queue_families.data()); bool found_queue = false; for (uint32_t i = 0; i < queue_family_count; ++i) @@ -299,8 +280,7 @@ class VulkanDevice break; } } - if (!found_queue) - NN_FAIL("No compute queue family found"); + if (!found_queue) NN_FAIL("No compute queue family found"); // 4. 创建逻辑设备 // 先探测时间线信号量能力（Vulkan 1.2 核心特性）：跨 submit 数据依赖 // 需要"同一信号量可被多个 submit 等待"的语义，二进制信号量做不到 @@ -312,8 +292,7 @@ class VulkanDevice { VkPhysicalDeviceProperties props{}; vkGetPhysicalDeviceProperties(physical_device_, &props); - if (app_info.apiVersion >= VK_API_VERSION_1_2 && - props.apiVersion >= VK_API_VERSION_1_2) + if (app_info.apiVersion >= VK_API_VERSION_1_2 && props.apiVersion >= VK_API_VERSION_1_2) { timeline_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_TIMELINE_SEMAPHORE_FEATURES; @@ -326,8 +305,7 @@ class VulkanDevice // 逃生阀：某些驱动的 timeline 实现有问题时，用 // NN_VULKAN_NO_TIMELINE=1 强制走"host 等在飞上传"回退路径 // （正确性优先，牺牲非阻塞流水线）。 - if (!get_env("NN_VULKAN_NO_TIMELINE").empty()) - timeline_semaphores_ = false; + if (!get_env("NN_VULKAN_NO_TIMELINE").empty()) timeline_semaphores_ = false; } // 计算 subgroup 尺寸（VK1.1 核心，查询链式同上方 features2 模式）： @@ -360,8 +338,7 @@ class VulkanDevice has_16bit_storage_ = (storage16.storageBuffer16BitAccess == VK_TRUE); } // 逃生阀：某些驱动/校验层对 16 位存储支持不佳时强制走边界 cast 路径 - if (!get_env("NN_VULKAN_NO_16BIT_STORAGE").empty()) - has_16bit_storage_ = false; + if (!get_env("NN_VULKAN_NO_16BIT_STORAGE").empty()) has_16bit_storage_ = false; // f16 ALU（Phase 3 原生 f16 算术）：查询 + 启用 shaderFloat16 // （Vulkan 1.2 核心特性 / VK_KHR_shader_float16_int8 扩展）。 @@ -378,10 +355,8 @@ class VulkanDevice has_shader_float16_ = (f16alu.shaderFloat16 == VK_TRUE); } // 逃生阀：强制回退 f32 算术变体（对照实验 / 驱动问题） - if (!get_env("NN_VULKAN_NO_16BIT_ALU").empty()) - has_shader_float16_ = false; - if (!has_shader_float16_) - f16alu.shaderFloat16 = VK_FALSE; + if (!get_env("NN_VULKAN_NO_16BIT_ALU").empty()) has_shader_float16_ = false; + if (!has_shader_float16_) f16alu.shaderFloat16 = VK_FALSE; // ── int8 点积（P3，W2A8）：查询 + 启用 shaderIntegerDotProduct ────── // VK_KHR_shader_integer_dot_product（Vulkan 1.3 核心）。Turing+ 的 DP4A @@ -398,19 +373,16 @@ class VulkanDevice has_int8_dot_ = (int8dot.shaderIntegerDotProduct == VK_TRUE); VkPhysicalDeviceShaderIntegerDotProductProperties props{}; - props.sType = - VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_INTEGER_DOT_PRODUCT_PROPERTIES; + props.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_INTEGER_DOT_PRODUCT_PROPERTIES; VkPhysicalDeviceProperties2 p2{}; p2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2; p2.pNext = &props; vkGetPhysicalDeviceProperties2(physical_device_, &p2); - int8_dot_accelerated_ = (props.integerDotProduct4x8BitPackedSignedAccelerated == - VK_TRUE); + int8_dot_accelerated_ = + (props.integerDotProduct4x8BitPackedSignedAccelerated == VK_TRUE); } - if (!get_env("NN_VULKAN_NO_INT8_DOT").empty()) - has_int8_dot_ = false; - if (!has_int8_dot_) - int8dot.shaderIntegerDotProduct = VK_FALSE; + if (!get_env("NN_VULKAN_NO_INT8_DOT").empty()) has_int8_dot_ = false; + if (!has_int8_dot_) int8dot.shaderIntegerDotProduct = VK_FALSE; // 查询阶段把整个结构体清零后可能置 VK_TRUE；非启用分支必须显式关掉， // 避免把未启用的特性声明进 pNext（校验层 VUID-VkDeviceCreateInfo 系列）。 f16alu.shaderInt8 = VK_FALSE; // 本库不用 int8 ALU，恒不启用 @@ -431,8 +403,7 @@ class VulkanDevice has_buffer_device_address_ = (bda.bufferDeviceAddress == VK_TRUE); } } - if (!has_buffer_device_address_) - bda.bufferDeviceAddress = VK_FALSE; + if (!has_buffer_device_address_) bda.bufferDeviceAddress = VK_FALSE; float queue_priority = 1.0f; VkDeviceQueueCreateInfo queue_info{}; @@ -446,28 +417,23 @@ class VulkanDevice // pNext 链：按需串起"已启用的特性"（链上顺序无语义，这里统一用前插， // 免得每加一个特性都要重排多层三元表达式 —— 之前那种写法已经很难读了） VkBaseInStructure* chain = nullptr; - const auto link = [&chain](VkBaseInStructure* node) { + const auto link = [&chain](VkBaseInStructure* node) + { node->pNext = chain; chain = node; }; - if (timeline_semaphores_) - link(reinterpret_cast(&timeline_features)); - if (has_16bit_storage_) - link(reinterpret_cast(&storage16)); - if (has_shader_float16_) - link(reinterpret_cast(&f16alu)); - if (has_int8_dot_) - link(reinterpret_cast(&int8dot)); + if (timeline_semaphores_) link(reinterpret_cast(&timeline_features)); + if (has_16bit_storage_) link(reinterpret_cast(&storage16)); + if (has_shader_float16_) link(reinterpret_cast(&f16alu)); + if (has_int8_dot_) link(reinterpret_cast(&int8dot)); // P3 BDA（`#b` 批量变体）：支持即启用，作为 pNext 链头（其余链不变） - if (has_buffer_device_address_) - link(reinterpret_cast(&bda)); + if (has_buffer_device_address_) link(reinterpret_cast(&bda)); device_info.pNext = chain; device_info.queueCreateInfoCount = 1; device_info.pQueueCreateInfos = &queue_info; res = vkCreateDevice(physical_device_, &device_info, nullptr, &device_); - if (res != VK_SUCCESS) - NN_FAIL("vkCreateDevice failed: " + std::to_string(res)); + if (res != VK_SUCCESS) NN_FAIL("vkCreateDevice failed: " + std::to_string(res)); // 5. 获取计算队列 vkGetDeviceQueue(device_, queue_family_index_, 0, &compute_queue_); @@ -490,8 +456,7 @@ class VulkanDevice [[nodiscard]] bool int8_dot_accelerated() const noexcept { return int8_dot_accelerated_; } // buffer device address 是否已启用：P3 批量派发（`key#b`）前提。未启用 // → 后端不登记 `#b` pipeline，运行期逐 kernel 降级（可观测、非错误）。 - [[nodiscard]] bool bda_available() const noexcept - { return has_buffer_device_address_; } + [[nodiscard]] bool bda_available() const noexcept { return has_buffer_device_address_; } // 所选物理设备名（"NVIDIA CMP 40HX" 等；初始化前为空） [[nodiscard]] const std::string& device_name() const noexcept { return device_name_; } // 计算队列 subgroup 尺寸（初始化前 = 兜底值 4；GEMV 分派门禁用） @@ -499,15 +464,10 @@ class VulkanDevice // timestamp query 能力（NN_PROFILE 设备侧计时）： // validBits=0 → 所选队列不支持时间戳；period = 1 tick 的纳秒数 [[nodiscard]] uint32_t queue_timestamp_valid_bits() const noexcept - { - return queue_timestamp_valid_bits_; - } + { return queue_timestamp_valid_bits_; } [[nodiscard]] float timestamp_period_ns() const noexcept { return timestamp_period_ns_; } // 已启用时间线信号量（决定跨 submit 依赖走信号量还是 host 等 fence） - [[nodiscard]] bool has_timeline_semaphores() const noexcept - { - return timeline_semaphores_; - } + [[nodiscard]] bool has_timeline_semaphores() const noexcept { return timeline_semaphores_; } }; // ══════════════════════════════════════════════════════════════════════════ @@ -528,8 +488,7 @@ class VulkanPipeline ~VulkanPipeline() { - if (pipeline_ != VK_NULL_HANDLE) - vkDestroyPipeline(device_, pipeline_, nullptr); + if (pipeline_ != VK_NULL_HANDLE) vkDestroyPipeline(device_, pipeline_, nullptr); if (pipeline_layout_ != VK_NULL_HANDLE) vkDestroyPipelineLayout(device_, pipeline_layout_, nullptr); if (descriptor_layout_ != VK_NULL_HANDLE) @@ -541,8 +500,8 @@ class VulkanPipeline // 移动语义 VulkanPipeline(VulkanPipeline&& o) noexcept : device_(o.device_), shader_module_(o.shader_module_), - descriptor_layout_(o.descriptor_layout_), - pipeline_layout_(o.pipeline_layout_), pipeline_(o.pipeline_) + descriptor_layout_(o.descriptor_layout_), pipeline_layout_(o.pipeline_layout_), + pipeline_(o.pipeline_) { o.device_ = VK_NULL_HANDLE; o.shader_module_ = VK_NULL_HANDLE; @@ -569,8 +528,8 @@ class VulkanPipeline VulkanPipeline& operator=(const VulkanPipeline&) = delete; // 创建 matmul pipeline - [[nodiscard]] static Result create_matmul( - VkDevice device, std::span spirv_code) + [[nodiscard]] static Result create_matmul(VkDevice device, + std::span spirv_code) { VulkanPipeline pl; pl.device_ = device; @@ -582,8 +541,8 @@ class VulkanPipeline module_info.pCode = spirv_code.data(); auto r = detail::vk_check( - vkCreateShaderModule(device, &module_info, nullptr, &pl.shader_module_), - __FILE__, __LINE__); + vkCreateShaderModule(device, &module_info, nullptr, &pl.shader_module_), __FILE__, + __LINE__); NN_TRY_CHECK(r); // 创建 descriptor set layout（3 个 storage buffer） VkDescriptorSetLayoutBinding bindings[3]{}; @@ -630,18 +589,19 @@ class VulkanPipeline pipeline_info.stage.pName = "main"; pipeline_info.layout = pl.pipeline_layout_; - r = detail::vk_check( - vkCreateComputePipelines(device, VK_NULL_HANDLE, 1, &pipeline_info, nullptr, &pl.pipeline_), - __FILE__, __LINE__); + r = detail::vk_check(vkCreateComputePipelines(device, VK_NULL_HANDLE, 1, &pipeline_info, + nullptr, &pl.pipeline_), + __FILE__, __LINE__); NN_TRY_CHECK(r); return pl; } // 创建通用 pipeline（可变 binding 数和 push constant 大小） // 用于 elementwise_v2 / reduce / broadcast 等纯原语着色器 - [[nodiscard]] static Result create_generic( - VkDevice device, std::span spirv_code, - uint32_t num_bindings, uint32_t push_constant_size) + [[nodiscard]] static Result create_generic(VkDevice device, + std::span spirv_code, + uint32_t num_bindings, + uint32_t push_constant_size) { VulkanPipeline pl; pl.device_ = device; @@ -653,8 +613,8 @@ class VulkanPipeline module_info.pCode = spirv_code.data(); auto r = detail::vk_check( - vkCreateShaderModule(device, &module_info, nullptr, &pl.shader_module_), - __FILE__, __LINE__); + vkCreateShaderModule(device, &module_info, nullptr, &pl.shader_module_), __FILE__, + __LINE__); NN_TRY_CHECK(r); // 创建 descriptor set layout（num_bindings 个 storage buffer） std::vector bindings(num_bindings); @@ -701,16 +661,16 @@ class VulkanPipeline pipeline_info.stage.pName = "main"; pipeline_info.layout = pl.pipeline_layout_; - r = detail::vk_check( - vkCreateComputePipelines(device, VK_NULL_HANDLE, 1, &pipeline_info, nullptr, &pl.pipeline_), - __FILE__, __LINE__); + r = detail::vk_check(vkCreateComputePipelines(device, VK_NULL_HANDLE, 1, &pipeline_info, + nullptr, &pl.pipeline_), + __FILE__, __LINE__); NN_TRY_CHECK(r); return pl; } [[nodiscard]] VkPipeline handle() const noexcept { return pipeline_; } [[nodiscard]] VkPipelineLayout pipeline_layout() const noexcept { return pipeline_layout_; } - [[nodiscard]] VkDescriptorSetLayout descriptor_layout() const noexcept { return descriptor_layout_; } + [[nodiscard]] VkDescriptorSetLayout descriptor_layout() const noexcept + { return descriptor_layout_; } }; } // namespace nn - diff --git a/include/neuralnet.cpp/backend/compute_vk_pipeline_catalog.hpp b/include/neuralnet.cpp/backend/compute_vk_pipeline_catalog.hpp index 85420905..1b1e4a88 100644 --- a/include/neuralnet.cpp/backend/compute_vk_pipeline_catalog.hpp +++ b/include/neuralnet.cpp/backend/compute_vk_pipeline_catalog.hpp @@ -90,8 +90,7 @@ enum class PipelineId : std::uint8_t count }; -inline constexpr std::size_t kPipelineCount = - static_cast(PipelineId::count); +inline constexpr std::size_t kPipelineCount = static_cast(PipelineId::count); // 设备能力要求：不满足时**不创建**该 pipeline（调用方看到空句柄后回退）。 enum class PipelineRequirement : std::uint8_t @@ -119,40 +118,30 @@ class PipelineCatalog // · SPIR-V 未嵌入（纯 CPU 构建）→ 跳过； // · 设备不满足 requirement → 跳过； // · 创建失败 → 跳过（该变体不可用，由上层回退）。 - void create_all(VkDevice device, const VulkanDevice& caps, - std::span specs) + void create_all(VkDevice device, const VulkanDevice& caps, std::span specs) { for (const auto& spec : specs) { - if (!supports(caps, spec.requirement)) - continue; + if (!supports(caps, spec.requirement)) continue; const auto& code = spec.spirv(); - if (code.empty()) - continue; - auto created = VulkanPipeline::create_generic( - device, code, spec.bindings, spec.push_constant_bytes); - if (created) - entries_[index(spec.id)] = std::move(*created); + if (code.empty()) continue; + auto created = VulkanPipeline::create_generic(device, code, spec.bindings, + spec.push_constant_bytes); + if (created) entries_[index(spec.id)] = std::move(*created); } } [[nodiscard]] bool has(PipelineId id) const noexcept - { - return entries_[index(id)].handle() != VK_NULL_HANDLE; - } + { return entries_[index(id)].handle() != VK_NULL_HANDLE; } // 调用方契约：先 has(id) 再 get(id)。缺失时返回空句柄的 pipeline， // Vulkan 校验层会报"绑定空 pipeline"，不会静默用错变体。 [[nodiscard]] const VulkanPipeline& get(PipelineId id) const noexcept - { - return entries_[index(id)]; - } + { return entries_[index(id)]; } private: static constexpr std::size_t index(PipelineId id) noexcept - { - return static_cast(id); - } + { return static_cast(id); } static bool supports(const VulkanDevice& caps, PipelineRequirement req) noexcept { diff --git a/include/neuralnet.cpp/backend/compute_vk_pipeline_specs.hpp b/include/neuralnet.cpp/backend/compute_vk_pipeline_specs.hpp index 6394cce8..eede68b2 100644 --- a/include/neuralnet.cpp/backend/compute_vk_pipeline_specs.hpp +++ b/include/neuralnet.cpp/backend/compute_vk_pipeline_specs.hpp @@ -22,90 +22,85 @@ inline std::span GpuBackend::pipeline_specs() using R = PipelineRequirement; static constexpr PipelineSpec kSpecs[] = { // ── 矩阵乘法族 ───────────────────────────────────────────────── - {PipelineId::matmul, "matmul", - get_matmul_spirv, 3, 5 * sizeof(std::uint32_t), R::always}, - {PipelineId::matmul_tiled, "matmul_tiled", - get_matmul_tiled_spirv, 3, 5 * sizeof(std::uint32_t), R::always}, + {PipelineId::matmul, "matmul", get_matmul_spirv, 3, 5 * sizeof(std::uint32_t), R::always}, + {PipelineId::matmul_tiled, "matmul_tiled", get_matmul_tiled_spirv, 3, + 5 * sizeof(std::uint32_t), R::always}, // 细 tile（16×16 输出块，wgrad 网格饥饿；f32 only） - {PipelineId::matmul_tiled_fine, "matmul_tiled_fine", - get_matmul_tiled_fine_spirv, 3, 5 * sizeof(std::uint32_t), R::always}, - {PipelineId::matmul_gemv, "matmul_gemv", - get_matmul_gemv_spirv, 3, 5 * sizeof(std::uint32_t), R::always}, - {PipelineId::matmul_gemv_f16, "matmul_gemv_f16", - get_matmul_gemv_f16_spirv, 3, 5 * sizeof(std::uint32_t), R::needs_16bit_storage}, + {PipelineId::matmul_tiled_fine, "matmul_tiled_fine", get_matmul_tiled_fine_spirv, 3, + 5 * sizeof(std::uint32_t), R::always}, + {PipelineId::matmul_gemv, "matmul_gemv", get_matmul_gemv_spirv, 3, + 5 * sizeof(std::uint32_t), R::always}, + {PipelineId::matmul_gemv_f16, "matmul_gemv_f16", get_matmul_gemv_f16_spirv, 3, + 5 * sizeof(std::uint32_t), R::needs_16bit_storage}, // 批量 matmul：push 多一个输出缩放 alpha（6 uint） - {PipelineId::batched_matmul, "batched_matmul", - get_batched_matmul_spirv, 3, 6 * sizeof(std::uint32_t), R::always}, - {PipelineId::batched_matmul_f16, "batched_matmul_f16", - get_batched_matmul_f16_spirv, 3, 6 * sizeof(std::uint32_t), R::needs_16bit_storage}, - {PipelineId::matmul_tiled_f16, "matmul_tiled_f16", - get_matmul_tiled_f16_spirv, 3, 5 * sizeof(std::uint32_t), R::needs_16bit_storage}, + {PipelineId::batched_matmul, "batched_matmul", get_batched_matmul_spirv, 3, + 6 * sizeof(std::uint32_t), R::always}, + {PipelineId::batched_matmul_f16, "batched_matmul_f16", get_batched_matmul_f16_spirv, 3, + 6 * sizeof(std::uint32_t), R::needs_16bit_storage}, + {PipelineId::matmul_tiled_f16, "matmul_tiled_f16", get_matmul_tiled_f16_spirv, 3, + 5 * sizeof(std::uint32_t), R::needs_16bit_storage}, // ── 逐元素 / 归约 ────────────────────────────────────────────── // elementwise：A/B/C/OUT 四绑定；8 个 uint 的 push - {PipelineId::elementwise_v2, "elementwise_v2", - get_elementwise_v2_spirv, 4, 8 * sizeof(std::uint32_t), R::always}, - {PipelineId::elementwise_v2_f16, "elementwise_v2_f16", - get_elementwise_v2_f16_spirv, 4, 8 * sizeof(std::uint32_t), R::needs_16bit_storage}, + {PipelineId::elementwise_v2, "elementwise_v2", get_elementwise_v2_spirv, 4, + 8 * sizeof(std::uint32_t), R::always}, + {PipelineId::elementwise_v2_f16, "elementwise_v2_f16", get_elementwise_v2_f16_spirv, 4, + 8 * sizeof(std::uint32_t), R::needs_16bit_storage}, // reduce：rows/cols/mode/reduce_op/chunk_rows - {PipelineId::reduce, "reduce", - get_reduce_spirv, 2, 5 * sizeof(std::uint32_t), R::always}, - {PipelineId::reduce_f16, "reduce_f16", - get_reduce_f16_spirv, 2, 5 * sizeof(std::uint32_t), R::needs_16bit_storage}, + {PipelineId::reduce, "reduce", get_reduce_spirv, 2, 5 * sizeof(std::uint32_t), R::always}, + {PipelineId::reduce_f16, "reduce_f16", get_reduce_f16_spirv, 2, 5 * sizeof(std::uint32_t), + R::needs_16bit_storage}, // 分组归约：rows/cols/R/op - {PipelineId::group_reduce, "group_reduce", - get_group_reduce_spirv, 2, 4 * sizeof(std::uint32_t), R::always}, - {PipelineId::group_reduce_f16, "group_reduce_f16", - get_group_reduce_f16_spirv, 2, 4 * sizeof(std::uint32_t), R::needs_16bit_storage}, + {PipelineId::group_reduce, "group_reduce", get_group_reduce_spirv, 2, + 4 * sizeof(std::uint32_t), R::always}, + {PipelineId::group_reduce_f16, "group_reduce_f16", get_group_reduce_f16_spirv, 2, + 4 * sizeof(std::uint32_t), R::needs_16bit_storage}, // ── 数据搬运 ─────────────────────────────────────────────────── // rearrange_3d：五个形状参数 - {PipelineId::rearrange_3d, "rearrange_3d", - get_rearrange_3d_spirv, 2, 5 * sizeof(std::uint32_t), R::always}, - {PipelineId::rearrange_3d_f16, "rearrange_3d_f16", - get_rearrange_3d_f16_spirv, 2, 5 * sizeof(std::uint32_t), R::needs_16bit_storage}, + {PipelineId::rearrange_3d, "rearrange_3d", get_rearrange_3d_spirv, 2, + 5 * sizeof(std::uint32_t), R::always}, + {PipelineId::rearrange_3d_f16, "rearrange_3d_f16", get_rearrange_3d_f16_spirv, 2, + 5 * sizeof(std::uint32_t), R::needs_16bit_storage}, // transpose：rows/cols - {PipelineId::transpose, "transpose", - get_transpose_spirv, 2, 2 * sizeof(std::uint32_t), R::always}, + {PipelineId::transpose, "transpose", get_transpose_spirv, 2, 2 * sizeof(std::uint32_t), + R::always}, // transpose_f16 的 shared float16_t 需要 SPIR-V Float16 capability - {PipelineId::transpose_f16, "transpose_f16", - get_transpose_f16_spirv, 2, 2 * sizeof(std::uint32_t), R::needs_shader_float16}, + {PipelineId::transpose_f16, "transpose_f16", get_transpose_f16_spirv, 2, + 2 * sizeof(std::uint32_t), R::needs_shader_float16}, // gather / scatter_add：3 个形状参数 - {PipelineId::gather, "gather", - get_gather_spirv, 3, 3 * sizeof(std::uint32_t), R::always}, - {PipelineId::gather_f16, "gather_f16", - get_gather_f16_spirv, 3, 3 * sizeof(std::uint32_t), R::needs_16bit_storage}, - {PipelineId::scatter_add, "scatter_add", - get_scatter_add_spirv, 3, 3 * sizeof(std::uint32_t), R::always}, - {PipelineId::scatter_add_f16, "scatter_add_f16", - get_scatter_add_f16_spirv, 3, 3 * sizeof(std::uint32_t), R::needs_16bit_storage}, + {PipelineId::gather, "gather", get_gather_spirv, 3, 3 * sizeof(std::uint32_t), R::always}, + {PipelineId::gather_f16, "gather_f16", get_gather_f16_spirv, 3, 3 * sizeof(std::uint32_t), + R::needs_16bit_storage}, + {PipelineId::scatter_add, "scatter_add", get_scatter_add_spirv, 3, + 3 * sizeof(std::uint32_t), R::always}, + {PipelineId::scatter_add_f16, "scatter_add_f16", get_scatter_add_f16_spirv, 3, + 3 * sizeof(std::uint32_t), R::needs_16bit_storage}, // im2col / col2im：9 个卷积窗口参数 - {PipelineId::im2col, "im2col", - get_im2col_spirv, 2, 9 * sizeof(std::uint32_t), R::always}, - {PipelineId::im2col_f16, "im2col_f16", - get_im2col_f16_spirv, 2, 9 * sizeof(std::uint32_t), R::needs_16bit_storage}, - {PipelineId::col2im, "col2im", - get_col2im_spirv, 2, 9 * sizeof(std::uint32_t), R::always}, - {PipelineId::col2im_f16, "col2im_f16", - get_col2im_f16_spirv, 2, 9 * sizeof(std::uint32_t), R::needs_16bit_storage}, + {PipelineId::im2col, "im2col", get_im2col_spirv, 2, 9 * sizeof(std::uint32_t), R::always}, + {PipelineId::im2col_f16, "im2col_f16", get_im2col_f16_spirv, 2, 9 * sizeof(std::uint32_t), + R::needs_16bit_storage}, + {PipelineId::col2im, "col2im", get_col2im_spirv, 2, 9 * sizeof(std::uint32_t), R::always}, + {PipelineId::col2im_f16, "col2im_f16", get_col2im_f16_spirv, 2, 9 * sizeof(std::uint32_t), + R::needs_16bit_storage}, // ── RLA 扫描原语（手写原语，不进 AOT 融合注册表）──────────────── // 前缀扫描：8 绑定 / 7 个参数（含 St scratch） - {PipelineId::scan_prefix_outer, "scan_prefix_outer", - get_scan_prefix_outer_spirv, 8, 7 * sizeof(std::uint32_t), R::always}, + {PipelineId::scan_prefix_outer, "scan_prefix_outer", get_scan_prefix_outer_spirv, 8, + 7 * sizeof(std::uint32_t), R::always}, // P4-2b 窄共享变体（dk ≤ 32；同布局） {PipelineId::scan_prefix_outer_d32, "scan_prefix_outer_d32", get_scan_prefix_outer_d32_spirv, 8, 7 * sizeof(std::uint32_t), R::always}, // 后缀扫描：5 绑定 / 6 个参数 - {PipelineId::scan_suffix_outer, "scan_suffix_outer", - get_scan_suffix_outer_spirv, 5, 6 * sizeof(std::uint32_t), R::always}, + {PipelineId::scan_suffix_outer, "scan_suffix_outer", get_scan_suffix_outer_spirv, 5, + 6 * sizeof(std::uint32_t), R::always}, // 通用 d_k（>64）：状态驻全局 scratch，多一档绑定 {PipelineId::scan_prefix_outer_gen, "scan_prefix_outer_gen", get_scan_prefix_outer_gen_spirv, 9, 7 * sizeof(std::uint32_t), R::always}, {PipelineId::scan_suffix_outer_gen, "scan_suffix_outer_gen", get_scan_suffix_outer_gen_spirv, 6, 6 * sizeof(std::uint32_t), R::always}, - {PipelineId::outer_col, "outer_col", - get_outer_col_spirv, 4, 4 * sizeof(std::uint32_t), R::always}, + {PipelineId::outer_col, "outer_col", get_outer_col_spirv, 4, 4 * sizeof(std::uint32_t), + R::always}, // f16 扫描变体（输入输出均 half；St scratch 仍 f32） {PipelineId::scan_prefix_outer_f16, "scan_prefix_outer_f16", get_scan_prefix_outer_f16_spirv, 8, 7 * sizeof(std::uint32_t), R::needs_16bit_storage}, @@ -117,26 +112,25 @@ inline std::span GpuBackend::pipeline_specs() get_scan_prefix_outer_gen_f16_spirv, 9, 7 * sizeof(std::uint32_t), R::needs_16bit_storage}, {PipelineId::scan_suffix_outer_gen_f16, "scan_suffix_outer_gen_f16", get_scan_suffix_outer_gen_f16_spirv, 6, 6 * sizeof(std::uint32_t), R::needs_16bit_storage}, - {PipelineId::outer_col_f16, "outer_col_f16", - get_outer_col_f16_spirv, 4, 4 * sizeof(std::uint32_t), R::needs_16bit_storage}, + {PipelineId::outer_col_f16, "outer_col_f16", get_outer_col_f16_spirv, 4, + 4 * sizeof(std::uint32_t), R::needs_16bit_storage}, // ── 精度转换 / 三值打包 ──────────────────────────────────────── // cast：{count, kind} - {PipelineId::cast, "cast", - get_cast_spirv, 2, 2 * sizeof(std::uint32_t), R::always}, + {PipelineId::cast, "cast", get_cast_spirv, 2, 2 * sizeof(std::uint32_t), R::always}, // trit_pack：{rows, k, wpr} - {PipelineId::trit_pack, "trit_pack", - get_trit_pack_spirv, 3, 3 * sizeof(std::uint32_t), R::always}, + {PipelineId::trit_pack, "trit_pack", get_trit_pack_spirv, 3, 3 * sizeof(std::uint32_t), + R::always}, // matmul_q：{M, N, K, transA, wpr} - {PipelineId::matmul_q, "matmul_q", - get_matmul_q_spirv, 4, 5 * sizeof(std::uint32_t), R::always}, - {PipelineId::matmul_q_f16, "matmul_q_f16", - get_matmul_q_f16_spirv, 4, 5 * sizeof(std::uint32_t), R::needs_16bit_storage}, + {PipelineId::matmul_q, "matmul_q", get_matmul_q_spirv, 4, 5 * sizeof(std::uint32_t), + R::always}, + {PipelineId::matmul_q_f16, "matmul_q_f16", get_matmul_q_f16_spirv, 4, + 5 * sizeof(std::uint32_t), R::needs_16bit_storage}, // 窄 N（N ≤ 8）打包 GEMV：解码主路径 - {PipelineId::matmul_q_gemv, "matmul_q_gemv", - get_matmul_q_gemv_spirv, 4, 5 * sizeof(std::uint32_t), R::always}, - {PipelineId::matmul_q_gemv_f16, "matmul_q_gemv_f16", - get_matmul_q_gemv_f16_spirv, 4, 5 * sizeof(std::uint32_t), R::needs_16bit_storage}, + {PipelineId::matmul_q_gemv, "matmul_q_gemv", get_matmul_q_gemv_spirv, 4, + 5 * sizeof(std::uint32_t), R::always}, + {PipelineId::matmul_q_gemv_f16, "matmul_q_gemv_f16", get_matmul_q_gemv_f16_spirv, 4, + 5 * sizeof(std::uint32_t), R::needs_16bit_storage}, }; return kSpecs; } diff --git a/include/neuralnet.cpp/batch_source.hpp b/include/neuralnet.cpp/batch_source.hpp index a7ec4785..e51bd155 100644 --- a/include/neuralnet.cpp/batch_source.hpp +++ b/include/neuralnet.cpp/batch_source.hpp @@ -57,17 +57,12 @@ class BatchSource // 每 epoch 步数；nullopt = 无固定 epoch（骨架退化为步进驱动）。 [[nodiscard]] virtual std::optional steps_per_epoch() const - { - return std::nullopt; - } + { return std::nullopt; } // 断点定位：把内部游标移到 (epoch, start_step)。无固定 epoch 的域用空实现。 // 有固定 epoch 的域在此重排样本顺序（每 epoch shuffle）。 - [[nodiscard]] virtual Result rewind(std::size_t /*epoch*/, - std::size_t /*start_step*/) - { - return {}; - } + [[nodiscard]] virtual Result rewind(std::size_t /*epoch*/, std::size_t /*start_step*/) + { return {}; } }; // ── 预置适配器：样本表 + 每 epoch shuffle + 按下标装配（文本/视觉复用）──── @@ -88,15 +83,12 @@ class IndexedBatchSource : public BatchSource // 工厂（参数校验；构造函数私有 —— 与 TrainingSession/TextTrainer 同口径） [[nodiscard]] static Result create(std::size_t sample_count, std::size_t batch_size, - std::uint64_t seed, - AssembleFn assemble) + std::uint64_t seed, AssembleFn assemble) { - if (batch_size == 0) - NN_FAIL("IndexedBatchSource::create: batch_size 必须 > 0"); + if (batch_size == 0) NN_FAIL("IndexedBatchSource::create: batch_size 必须 > 0"); if (sample_count == 0) NN_FAIL("IndexedBatchSource::create: sample_count 必须 > 0（训练集为空）"); - if (!assemble) - NN_FAIL("IndexedBatchSource::create: 装配回调不能为空"); + if (!assemble) NN_FAIL("IndexedBatchSource::create: 装配回调不能为空"); return IndexedBatchSource(sample_count, batch_size, seed, std::move(assemble)); } @@ -122,9 +114,7 @@ class IndexedBatchSource : public BatchSource } [[nodiscard]] std::optional steps_per_epoch() const override - { - return (sample_count_ + batch_size_ - 1) / batch_size_; - } + { return (sample_count_ + batch_size_ - 1) / batch_size_; } [[nodiscard]] Result rewind(std::size_t epoch, std::size_t start_step) override { @@ -135,7 +125,8 @@ class IndexedBatchSource : public BatchSource // 单 RNG 流语义：从恒等序起，用同一 RNG 依次重放 0..epoch 共 epoch+1 次 // shuffle —— 与"每 epoch 顺序推进同一 RNG 流"的取样语义一致，同时 // (seed, epoch) 唯一决定顺序 ⇒ 幂等（重复 rewind 同一 epoch 结果不变）。 - for (std::size_t i = 0; i < sample_count_; ++i) order_[i] = i; + for (std::size_t i = 0; i < sample_count_; ++i) + order_[i] = i; std::mt19937_64 rng(seed_); for (std::size_t e = 0; e <= epoch; ++e) std::shuffle(order_.begin(), order_.end(), rng); @@ -152,13 +143,14 @@ class IndexedBatchSource : public BatchSource [[nodiscard]] std::span order() const noexcept { return order_; } private: - IndexedBatchSource(std::size_t sample_count, std::size_t batch_size, - std::uint64_t seed, AssembleFn assemble) + IndexedBatchSource(std::size_t sample_count, std::size_t batch_size, std::uint64_t seed, + AssembleFn assemble) : sample_count_(sample_count), batch_size_(batch_size), seed_(seed), assemble_(std::move(assemble)) { order_.resize(sample_count_); - for (std::size_t i = 0; i < sample_count_; ++i) order_[i] = i; + for (std::size_t i = 0; i < sample_count_; ++i) + order_[i] = i; } std::size_t sample_count_ = 0; diff --git a/include/neuralnet.cpp/cli/cli_engine_factory.hpp b/include/neuralnet.cpp/cli/cli_engine_factory.hpp index 63d39e43..7ffa90f8 100644 --- a/include/neuralnet.cpp/cli/cli_engine_factory.hpp +++ b/include/neuralnet.cpp/cli/cli_engine_factory.hpp @@ -31,41 +31,42 @@ namespace nn::cli { - // ── 引擎选择配置 ─────────────────────────────────────────────────────── - struct EngineConfig - { - bool use_gpu = false; // 启用 Vulkan GpuEngine +// ── 引擎选择配置 ─────────────────────────────────────────────────────── +struct EngineConfig +{ + bool use_gpu = false; // 启用 Vulkan GpuEngine // 手动指定 Vulkan 计算设备（索引 "2" 或名称子串 "40HX"/"NVIDIA"）； // 空 = 自动选择（设备类型 + apiVersion 打分）。 // 注意保留默认成员初始化器：聚合初始化 `EngineConfig{a, b}` 在 // -Wextra -Werror 下会对无初始化器的成员报 -Wmissing-field-initializers。 - std::string gpu_device = {}; - }; + std::string gpu_device = {}; +}; - // ── 创建计算引擎 ─────────────────────────────────────────────────────── - // 根据 use_gpu 选择 GpuEngine/CpuEngine。 - // 返回 Result：显式请求的后端不可用（初始化失败/未编译）时返回错误， - // **不回退 CPU**；成功返回引擎并打印选择信息到 log。 - [[nodiscard]] inline nn::Result> create_engine( - const EngineConfig &cfg, [[maybe_unused]] std::ostream &log = std::cerr) +// ── 创建计算引擎 ─────────────────────────────────────────────────────── +// 根据 use_gpu 选择 GpuEngine/CpuEngine。 +// 返回 Result：显式请求的后端不可用（初始化失败/未编译）时返回错误， +// **不回退 CPU**；成功返回引擎并打印选择信息到 log。 +[[nodiscard]] inline nn::Result> +create_engine(const EngineConfig& cfg, [[maybe_unused]] std::ostream& log = std::cerr) +{ + if (cfg.use_gpu) { - if (cfg.use_gpu) - { #ifdef NN_HAS_VULKAN - auto backend = nn::GpuBackend::shared(); - auto init_r = backend->initialize(cfg.gpu_device); - if (init_r) - { - log << "GPU 加速已启用 (Vulkan GpuEngine, GPU: " - << backend->device().device_name() << ")\n\n"; - return nn::GpuEngine::create(std::move(backend)); - } - NN_FAIL("请求 --gpu 但 GPU 初始化失败: " + init_r.error().message + "（不回退 CPU；请移除 --gpu 或改用 CPU 模式）"); + auto backend = nn::GpuBackend::shared(); + auto init_r = backend->initialize(cfg.gpu_device); + if (init_r) + { + log << "GPU 加速已启用 (Vulkan GpuEngine, GPU: " << backend->device().device_name() + << ")\n\n"; + return nn::GpuEngine::create(std::move(backend)); + } + NN_FAIL("请求 --gpu 但 GPU 初始化失败: " + init_r.error().message + + "（不回退 CPU；请移除 --gpu 或改用 CPU 模式）"); #else - NN_FAIL("请求 --gpu 但未编译 Vulkan 支持（NN_HAS_VULKAN）；不回退 CPU"); + NN_FAIL("请求 --gpu 但未编译 Vulkan 支持（NN_HAS_VULKAN）；不回退 CPU"); #endif - } - - return nn::CpuEngine::create(); } + + return nn::CpuEngine::create(); +} } // namespace nn::cli diff --git a/include/neuralnet.cpp/cli/cli_gpu_option.hpp b/include/neuralnet.cpp/cli/cli_gpu_option.hpp index a91364d0..75651136 100644 --- a/include/neuralnet.cpp/cli/cli_gpu_option.hpp +++ b/include/neuralnet.cpp/cli/cli_gpu_option.hpp @@ -24,44 +24,41 @@ namespace nn::cli { - // 纯数字 token（枚举索引）判定 - [[nodiscard]] inline bool is_device_index_token(const char *token) noexcept +// 纯数字 token（枚举索引）判定 +[[nodiscard]] inline bool is_device_index_token(const char* token) noexcept +{ + if (token == nullptr || *token == '\0') return false; + for (const char* p = token; *p != '\0'; ++p) { - if (token == nullptr || *token == '\0') - return false; - for (const char *p = token; *p != '\0'; ++p) - { - if (*p < '0' || *p > '9') - return false; - } - return true; + if (*p < '0' || *p > '9') return false; } + return true; +} - // 在 argv[i] 处尝试解析 `--gpu` 选项。 - // 返回 std::nullopt → 不是 --gpu（i 不变，调用方继续判断其他 flag） - // 返回 空 string → 是 --gpu，未指定设备（i 不变） - // 返回 非空 string → 是 --gpu 且指定了设备（i 递增到本次消费的索引） - // - // allow_name_value：空格形式是否连名称也收（仅无位置参数的程序可开）。 - [[nodiscard]] inline std::optional parse_gpu_option( - int argc, char *argv[], int &i, bool allow_name_value = false) +// 在 argv[i] 处尝试解析 `--gpu` 选项。 +// 返回 std::nullopt → 不是 --gpu（i 不变，调用方继续判断其他 flag） +// 返回 空 string → 是 --gpu，未指定设备（i 不变） +// 返回 非空 string → 是 --gpu 且指定了设备（i 递增到本次消费的索引） +// +// allow_name_value：空格形式是否连名称也收（仅无位置参数的程序可开）。 +[[nodiscard]] inline std::optional parse_gpu_option(int argc, char* argv[], int& i, + bool allow_name_value = false) +{ + const std::string_view arg = argv[i]; + if (arg == "--gpu") { - const std::string_view arg = argv[i]; - if (arg == "--gpu") + if (i + 1 < argc && argv[i + 1] != nullptr) { - if (i + 1 < argc && argv[i + 1] != nullptr) - { - const bool is_index = is_device_index_token(argv[i + 1]); - if (is_index || (allow_name_value && argv[i + 1][0] != '-')) - return std::string(argv[++i]); - } - return std::string{}; + const bool is_index = is_device_index_token(argv[i + 1]); + if (is_index || (allow_name_value && argv[i + 1][0] != '-')) + return std::string(argv[++i]); } + return std::string{}; + } - constexpr std::string_view kPrefix = "--gpu="; - if (arg.rfind(kPrefix, 0) == 0) - return std::string(arg.substr(kPrefix.size())); + constexpr std::string_view kPrefix = "--gpu="; + if (arg.rfind(kPrefix, 0) == 0) return std::string(arg.substr(kPrefix.size())); - return std::nullopt; - } + return std::nullopt; +} } // namespace nn::cli diff --git a/include/neuralnet.cpp/cli/cli_help.hpp b/include/neuralnet.cpp/cli/cli_help.hpp index c732275d..87b1e8e9 100644 --- a/include/neuralnet.cpp/cli/cli_help.hpp +++ b/include/neuralnet.cpp/cli/cli_help.hpp @@ -45,264 +45,268 @@ namespace nn::cli { - // ── 终端显示宽度：ASCII 记 1 列，CJK/全角记 2 列 ───────────────────── - // 覆盖常见东亚宽字符区段（含全角标点、CJK 扩展）；非法 UTF-8 字节按 1 列。 - [[nodiscard]] inline bool is_wide_codepoint(std::uint32_t cp) noexcept - { - if (cp >= 0x1100u && cp <= 0x115Fu) return true; // 谚文字母 - if (cp >= 0x2E80u && cp <= 0x303Eu) return true; // CJK 部首/符号/假名 - if (cp >= 0x3041u && cp <= 0x33FFu) return true; // 假名/注音/CJK 兼容 - if (cp >= 0x3400u && cp <= 0x4DBFu) return true; // CJK 扩展 A - if (cp >= 0x4E00u && cp <= 0x9FFFu) return true; // CJK 统一表意 - if (cp >= 0xA000u && cp <= 0xA4CFu) return true; // 彝文 - if (cp >= 0xAC00u && cp <= 0xD7A3u) return true; // 谚文音节 - if (cp >= 0xF900u && cp <= 0xFAFFu) return true; // CJK 兼容表意 - if (cp >= 0xFE30u && cp <= 0xFE4Fu) return true; // CJK 兼容形式 - if (cp >= 0xFF00u && cp <= 0xFF60u) return true; // 全角形式 - if (cp >= 0xFFE0u && cp <= 0xFFE6u) return true; // 全角符号 - if (cp >= 0x20000u && cp <= 0x3FFFD) return true; // CJK 扩展 B+ - return false; - } +// ── 终端显示宽度：ASCII 记 1 列，CJK/全角记 2 列 ───────────────────── +// 覆盖常见东亚宽字符区段（含全角标点、CJK 扩展）；非法 UTF-8 字节按 1 列。 +[[nodiscard]] inline bool is_wide_codepoint(std::uint32_t cp) noexcept +{ + if (cp >= 0x1100u && cp <= 0x115Fu) return true; // 谚文字母 + if (cp >= 0x2E80u && cp <= 0x303Eu) return true; // CJK 部首/符号/假名 + if (cp >= 0x3041u && cp <= 0x33FFu) return true; // 假名/注音/CJK 兼容 + if (cp >= 0x3400u && cp <= 0x4DBFu) return true; // CJK 扩展 A + if (cp >= 0x4E00u && cp <= 0x9FFFu) return true; // CJK 统一表意 + if (cp >= 0xA000u && cp <= 0xA4CFu) return true; // 彝文 + if (cp >= 0xAC00u && cp <= 0xD7A3u) return true; // 谚文音节 + if (cp >= 0xF900u && cp <= 0xFAFFu) return true; // CJK 兼容表意 + if (cp >= 0xFE30u && cp <= 0xFE4Fu) return true; // CJK 兼容形式 + if (cp >= 0xFF00u && cp <= 0xFF60u) return true; // 全角形式 + if (cp >= 0xFFE0u && cp <= 0xFFE6u) return true; // 全角符号 + if (cp >= 0x20000u && cp <= 0x3FFFD) return true; // CJK 扩展 B+ + return false; +} - [[nodiscard]] inline std::size_t display_width(std::string_view s) noexcept +[[nodiscard]] inline std::size_t display_width(std::string_view s) noexcept +{ + std::size_t width = 0; + for (std::size_t i = 0; i < s.size();) { - std::size_t width = 0; - for (std::size_t i = 0; i < s.size();) - { - const auto lead = static_cast(s[i]); - std::size_t len = 1; - if (lead >= 0xF0u) len = 4; - else if (lead >= 0xE0u) len = 3; - else if (lead >= 0xC0u) len = 2; + const auto lead = static_cast(s[i]); + std::size_t len = 1; + if (lead >= 0xF0u) len = 4; + else if (lead >= 0xE0u) len = 3; + else if (lead >= 0xC0u) len = 2; - std::uint32_t cp = lead; - bool valid = (i + len <= s.size()); - if (valid && len > 1) + std::uint32_t cp = lead; + bool valid = (i + len <= s.size()); + if (valid && len > 1) + { + cp = lead & (0xFFu >> (len + 1)); + for (std::size_t k = 1; k < len; ++k) { - cp = lead & (0xFFu >> (len + 1)); - for (std::size_t k = 1; k < len; ++k) + const auto cont = static_cast(s[i + k]); + if ((cont & 0xC0u) != 0x80u) { - const auto cont = static_cast(s[i + k]); - if ((cont & 0xC0u) != 0x80u) { valid = false; break; } - cp = (cp << 6) | (cont & 0x3Fu); + valid = false; + break; } + cp = (cp << 6) | (cont & 0x3Fu); } - width += (valid && is_wide_codepoint(cp)) ? 2 : 1; - i += valid ? len : 1; } - return width; + width += (valid && is_wide_codepoint(cp)) ? 2 : 1; + i += valid ? len : 1; } + return width; +} - // ── 帮助排版器 ─────────────────────────────────────────────────────── - // 用法：在 print_usage 里构造局部对象，链式填表；析构时自动渲染一次 - // （也可显式 flush()）。每条调用写在同一物理行（见文件头审计约定）。 - class Help +// ── 帮助排版器 ─────────────────────────────────────────────────────── +// 用法：在 print_usage 里构造局部对象，链式填表；析构时自动渲染一次 +// （也可显式 flush()）。每条调用写在同一物理行（见文件头审计约定）。 +class Help +{ +public: + Help(std::ostream& out, std::string_view prog, std::string_view title) + : out_(out), prog_(prog), title_(title) { - public: - Help(std::ostream &out, std::string_view prog, std::string_view title) - : out_(out), prog_(prog), title_(title) - { - } + } - ~Help() { flush(); } + ~Help() { flush(); } - Help(const Help &) = delete; - Help &operator=(const Help &) = delete; + Help(const Help&) = delete; + Help& operator=(const Help&) = delete; - // 用法行：渲染为 ` `；desc 非空时对齐到用法列。 - Help &usage(std::string_view args, std::string_view desc = {}) - { - usage_.push_back(Row{Kind::Entry, prefix_prog(std::string(args)), std::string(desc)}); - return *this; - } + // 用法行：渲染为 ` `；desc 非空时对齐到用法列。 + Help& usage(std::string_view args, std::string_view desc = {}) + { + usage_.push_back(Row{Kind::Entry, prefix_prog(std::string(args)), std::string(desc)}); + return *this; + } - // 开一个新节（标题后自动空行） - Help §ion(std::string_view title) - { - sections_.push_back(Section{std::string(title), {}}); - return *this; - } + // 开一个新节（标题后自动空行） + Help& section(std::string_view title) + { + sections_.push_back(Section{std::string(title), {}}); + return *this; + } - // 节内说明段落，4 空格缩进 - Help ¬e(std::string_view line) - { - current().push_back(Row{Kind::Note, std::string(line), {}}); - return *this; - } + // 节内说明段落，4 空格缩进 + Help& note(std::string_view line) + { + current().push_back(Row{Kind::Note, std::string(line), {}}); + return *this; + } - // 自由文本行，2 空格缩进 - Help &text(std::string_view line) - { - current().push_back(Row{Kind::Text, std::string(line), {}}); - return *this; - } + // 自由文本行，2 空格缩进 + Help& text(std::string_view line) + { + current().push_back(Row{Kind::Text, std::string(line), {}}); + return *this; + } - // 位置参数 / 词条（无 `--` 前缀），与选项同列对齐 - Help &item(std::string_view name, std::string_view desc) - { - current().push_back(Row{Kind::Entry, std::string(name), std::string(desc)}); - return *this; - } + // 位置参数 / 词条（无 `--` 前缀），与选项同列对齐 + Help& item(std::string_view name, std::string_view desc) + { + current().push_back(Row{Kind::Entry, std::string(name), std::string(desc)}); + return *this; + } - // 选项：flags 例如 "--gpu [索引|名称]"；desc 内的 '\n' 为续行 - Help &opt(std::string_view flags, std::string_view desc) - { - current().push_back(Row{Kind::Entry, std::string(flags), std::string(desc)}); - return *this; - } + // 选项：flags 例如 "--gpu [索引|名称]"；desc 内的 '\n' 为续行 + Help& opt(std::string_view flags, std::string_view desc) + { + current().push_back(Row{Kind::Entry, std::string(flags), std::string(desc)}); + return *this; + } - void flush() - { - if (rendered_) return; - rendered_ = true; + void flush() + { + if (rendered_) return; + rendered_ = true; - out_ << title_ << "\n"; + out_ << title_ << "\n"; - if (!usage_.empty()) - { - out_ << "\n用法:\n"; - const std::size_t usage_col = column_for(usage_, kMaxUsageWidth); - for (const Row &r : usage_) - write_entry(r.name, r.desc, usage_col, kMaxUsageWidth, kUsageOverflow); - } + if (!usage_.empty()) + { + out_ << "\n用法:\n"; + const std::size_t usage_col = column_for(usage_, kMaxUsageWidth); + for (const Row& r : usage_) + write_entry(r.name, r.desc, usage_col, kMaxUsageWidth, kUsageOverflow); + } - // 选项列宽全局统一（跨节一致），不被某一节的短条目带偏 - std::vector entries; - for (const Section &s : sections_) - for (const Row &r : s.rows) - if (r.kind == Kind::Entry) - entries.push_back(r); - const std::size_t col = column_for(entries, kMaxOptionWidth); + // 选项列宽全局统一（跨节一致），不被某一节的短条目带偏 + std::vector entries; + for (const Section& s : sections_) + for (const Row& r : s.rows) + if (r.kind == Kind::Entry) entries.push_back(r); + const std::size_t col = column_for(entries, kMaxOptionWidth); - for (const Section &s : sections_) + for (const Section& s : sections_) + { + if (!s.title.empty()) out_ << "\n" << s.title << ":\n"; + for (const Row& r : s.rows) { - if (!s.title.empty()) out_ << "\n" << s.title << ":\n"; - for (const Row &r : s.rows) + switch (r.kind) { - switch (r.kind) - { - case Kind::Entry: - write_entry(r.name, r.desc, col, kMaxOptionWidth, col); - break; - case Kind::Note: - write_paragraph(r.name, kNoteIndent); - break; - case Kind::Text: - write_paragraph(r.name, kTextIndent); - break; - } + case Kind::Entry: + write_entry(r.name, r.desc, col, kMaxOptionWidth, col); + break; + case Kind::Note: + write_paragraph(r.name, kNoteIndent); + break; + case Kind::Text: + write_paragraph(r.name, kTextIndent); + break; } } } + } - private: - enum Kind { Entry, Note, Text }; +private: + enum Kind + { + Entry, + Note, + Text + }; - struct Row - { - Kind kind; - std::string name; - std::string desc; - }; + struct Row + { + Kind kind; + std::string name; + std::string desc; + }; - struct Section - { - std::string title; - std::vector rows; - }; + struct Section + { + std::string title; + std::vector rows; + }; - static constexpr std::size_t kMinColumn = 16; // 条目名最小列宽 - static constexpr std::size_t kMaxOptionWidth = 34; // 选项名最大列宽 - static constexpr std::size_t kMaxUsageWidth = 56; // 用法行最大列宽 - static constexpr std::size_t kUsageOverflow = 4; // 用法行过长时说明的缩进 - static constexpr std::size_t kNoteIndent = 4; - static constexpr std::size_t kTextIndent = 2; + static constexpr std::size_t kMinColumn = 16; // 条目名最小列宽 + static constexpr std::size_t kMaxOptionWidth = 34; // 选项名最大列宽 + static constexpr std::size_t kMaxUsageWidth = 56; // 用法行最大列宽 + static constexpr std::size_t kUsageOverflow = 4; // 用法行过长时说明的缩进 + static constexpr std::size_t kNoteIndent = 4; + static constexpr std::size_t kTextIndent = 2; - std::vector ¤t() - { - if (sections_.empty()) section(""); - return sections_.back().rows; - } + std::vector& current() + { + if (sections_.empty()) section(""); + return sections_.back().rows; + } - [[nodiscard]] std::string prefix_prog(const std::string &args) const + [[nodiscard]] std::string prefix_prog(const std::string& args) const + { + std::string left = prog_; + if (!args.empty()) { - std::string left = prog_; - if (!args.empty()) - { - left += ' '; - left += args; - } - return left; + left += ' '; + left += args; } + return left; + } - // 描述列 = 2（缩进）+ left + 2（间隔），left 夹在 [kMinColumn, max] - [[nodiscard]] static std::size_t column_for( - const std::vector &rows, std::size_t max_width) noexcept - { - std::size_t left = kMinColumn; - for (const Row &r : rows) - if (r.kind == Entry) - left = std::max(left, display_width(r.name)); - left = std::min(left, max_width); - return 2 + left + 2; - } + // 描述列 = 2（缩进）+ left + 2（间隔），left 夹在 [kMinColumn, max] + [[nodiscard]] static std::size_t column_for(const std::vector& rows, + std::size_t max_width) noexcept + { + std::size_t left = kMinColumn; + for (const Row& r : rows) + if (r.kind == Entry) left = std::max(left, display_width(r.name)); + left = std::min(left, max_width); + return 2 + left + 2; + } - // 段落：每行统一缩进 indent（note/text 用） - void write_paragraph(const std::string &body, std::size_t indent) + // 段落：每行统一缩进 indent（note/text 用） + void write_paragraph(const std::string& body, std::size_t indent) + { + const std::string pad(indent, ' '); + std::size_t pos = 0; + while (true) { - const std::string pad(indent, ' '); - std::size_t pos = 0; - while (true) - { - const std::size_t nl = body.find('\n', pos); - const std::string_view part = - (nl == std::string::npos) ? std::string_view(body).substr(pos) + const std::size_t nl = body.find('\n', pos); + const std::string_view part = (nl == std::string::npos) + ? std::string_view(body).substr(pos) : std::string_view(body).substr(pos, nl - pos); - out_ << pad << part << "\n"; - if (nl == std::string::npos) break; - pos = nl + 1; - } + out_ << pad << part << "\n"; + if (nl == std::string::npos) break; + pos = nl + 1; } + } - void write_entry(const std::string &name, const std::string &desc, - std::size_t desc_col, std::size_t max_width, - std::size_t over_indent) + void write_entry(const std::string& name, const std::string& desc, std::size_t desc_col, + std::size_t max_width, std::size_t over_indent) + { + out_ << " " << name; + if (desc.empty()) { - out_ << " " << name; - if (desc.empty()) - { - out_ << "\n"; - return; - } + out_ << "\n"; + return; + } - const std::size_t w = display_width(name); - const bool same_line = (w <= max_width); - if (same_line) - out_ << std::string(desc_col - 2 - w, ' '); - else - out_ << "\n" << std::string(over_indent, ' '); + const std::size_t w = display_width(name); + const bool same_line = (w <= max_width); + if (same_line) out_ << std::string(desc_col - 2 - w, ' '); + else out_ << "\n" << std::string(over_indent, ' '); - std::size_t pos = 0; - bool first = true; - while (true) - { - const std::size_t nl = desc.find('\n', pos); - const std::string_view part = - (nl == std::string::npos) ? std::string_view(desc).substr(pos) + std::size_t pos = 0; + bool first = true; + while (true) + { + const std::size_t nl = desc.find('\n', pos); + const std::string_view part = (nl == std::string::npos) + ? std::string_view(desc).substr(pos) : std::string_view(desc).substr(pos, nl - pos); - if (!first) out_ << std::string(over_indent, ' '); - out_ << part << "\n"; - first = false; - if (nl == std::string::npos) break; - pos = nl + 1; - } + if (!first) out_ << std::string(over_indent, ' '); + out_ << part << "\n"; + first = false; + if (nl == std::string::npos) break; + pos = nl + 1; } + } - std::ostream &out_; - std::string prog_; - std::string title_; - std::vector usage_; - std::vector

sections_; - bool rendered_ = false; - }; + std::ostream& out_; + std::string prog_; + std::string title_; + std::vector usage_; + std::vector
sections_; + bool rendered_ = false; +}; } // namespace nn::cli diff --git a/include/neuralnet.cpp/cli/cli_lr_scheduler.hpp b/include/neuralnet.cpp/cli/cli_lr_scheduler.hpp index 95140c3a..da11484a 100644 --- a/include/neuralnet.cpp/cli/cli_lr_scheduler.hpp +++ b/include/neuralnet.cpp/cli/cli_lr_scheduler.hpp @@ -11,8 +11,8 @@ namespace nn::cli { -using nn::LrScheduleConfig; -using nn::StepLrScheduleConfig; using nn::compute_epoch_lr; using nn::compute_step_lr; +using nn::LrScheduleConfig; +using nn::StepLrScheduleConfig; } // namespace nn::cli diff --git a/include/neuralnet.cpp/cli/cli_mnist_io.hpp b/include/neuralnet.cpp/cli/cli_mnist_io.hpp index 37d7b8ca..609830df 100644 --- a/include/neuralnet.cpp/cli/cli_mnist_io.hpp +++ b/include/neuralnet.cpp/cli/cli_mnist_io.hpp @@ -36,203 +36,194 @@ namespace nn::cli { - // ── 从 CSV 加载 MNIST 数据 ───────────────────────────────────────────── - // 返回 (feat_mat[feat_dim, N], label_mat[10, N])，label_mat 为 one-hot。 - // max_samples = 0 表示加载全部样本。 - [[nodiscard]] inline nn::Result> - load_mnist_csv(const std::string &path, std::size_t max_samples = 0) +// ── 从 CSV 加载 MNIST 数据 ───────────────────────────────────────────── +// 返回 (feat_mat[feat_dim, N], label_mat[10, N])，label_mat 为 one-hot。 +// max_samples = 0 表示加载全部样本。 +[[nodiscard]] inline nn::Result> +load_mnist_csv(const std::string& path, std::size_t max_samples = 0) +{ + std::ifstream file(path, std::ios::binary | std::ios::ate); + if (!file.is_open()) NN_FAIL("Cannot open file: " + path); + const auto file_size = file.tellg(); + file.seekg(0); + + std::string buffer(static_cast(file_size), '\0'); + file.read(buffer.data(), file_size); + file.close(); + + std::size_t row_count = 0; + for (char c : buffer) + if (c == '\n') ++row_count; + // 末行无换行符时也要计入（否则最后一行样本被静默丢弃） + if (!buffer.empty() && buffer.back() != '\n') ++row_count; + if (row_count == 0) NN_FAIL("CSV file is empty or malformed: " + path); + if (max_samples > 0 && max_samples < row_count) row_count = max_samples; + + const char* ptr = buffer.data(); + // std::string::iterator 是 __wrap_iter 类（libc++），不能当 const char* 用； + // from_chars 需要裸指针，data()+size() 是唯一边界 + const char* end = buffer.data() + buffer.size(); + + // 探测第一行以确定特征维度 + int first_label = 0; + std::size_t feat_dim = 0; { - std::ifstream file(path, std::ios::binary | std::ios::ate); - if (!file.is_open()) - NN_FAIL("Cannot open file: " + path); - const auto file_size = file.tellg(); - file.seekg(0); - - std::string buffer(static_cast(file_size), '\0'); - file.read(buffer.data(), file_size); - file.close(); - - std::size_t row_count = 0; - for (char c : buffer) - if (c == '\n') ++row_count; - // 末行无换行符时也要计入（否则最后一行样本被静默丢弃） - if (!buffer.empty() && buffer.back() != '\n') - ++row_count; - if (row_count == 0) - NN_FAIL("CSV file is empty or malformed: " + path); - if (max_samples > 0 && max_samples < row_count) - row_count = max_samples; - - const char *ptr = buffer.data(); - // std::string::iterator 是 __wrap_iter 类（libc++），不能当 const char* 用； - // from_chars 需要裸指针，data()+size() 是唯一边界 - const char *end = buffer.data() + buffer.size(); - - // 探测第一行以确定特征维度 - int first_label = 0; - std::size_t feat_dim = 0; + const char* p = ptr; + auto [p1, ec1] = std::from_chars(p, end, first_label); + p = p1; + std::size_t cnt = 0; + while (p < end && *p != '\n' && *p != '\r') { - const char *p = ptr; - auto [p1, ec1] = std::from_chars(p, end, first_label); - p = p1; - std::size_t cnt = 0; - while (p < end && *p != '\n' && *p != '\r') + if (*p == ',') { - if (*p == ',') - { - ++cnt; - nn::Scalar tmp; - auto [p2, ec2] = std::from_chars(p + 1, end, tmp); - p = p2; - } - else - ++p; + ++cnt; + nn::Scalar tmp; + auto [p2, ec2] = std::from_chars(p + 1, end, tmp); + p = p2; } - feat_dim = cnt; + else ++p; } + feat_dim = cnt; + } - std::vector features(row_count * feat_dim); - std::vector labels(row_count); - - std::size_t row = 0; - ptr = buffer.data(); + std::vector features(row_count * feat_dim); + std::vector labels(row_count); - while (ptr < end && row < row_count) - { - auto [p_label, ec_label] = std::from_chars(ptr, end, labels[row]); - if (ec_label != std::errc{}) - NN_FAIL("Failed to parse label at row " + std::to_string(row)); - ptr = p_label; + std::size_t row = 0; + ptr = buffer.data(); - for (std::size_t j = 0; j < feat_dim; ++j) - { - if (ptr < end && *ptr == ',') ++ptr; - nn::Scalar val; - auto [p_feat, ec_feat] = std::from_chars(ptr, end, val); - if (ec_feat != std::errc{}) - NN_FAIL("Failed to parse feature at row " + std::to_string(row)); - features[row * feat_dim + j] = val; - ptr = p_feat; - } + while (ptr < end && row < row_count) + { + auto [p_label, ec_label] = std::from_chars(ptr, end, labels[row]); + if (ec_label != std::errc{}) NN_FAIL("Failed to parse label at row " + std::to_string(row)); + ptr = p_label; - while (ptr < end && (*ptr == '\r' || *ptr == '\n')) ++ptr; - ++row; + for (std::size_t j = 0; j < feat_dim; ++j) + { + if (ptr < end && *ptr == ',') ++ptr; + nn::Scalar val; + auto [p_feat, ec_feat] = std::from_chars(ptr, end, val); + if (ec_feat != std::errc{}) + NN_FAIL("Failed to parse feature at row " + std::to_string(row)); + features[row * feat_dim + j] = val; + ptr = p_feat; } - const std::size_t N = row; + while (ptr < end && (*ptr == '\r' || *ptr == '\n')) + ++ptr; + ++row; + } - nn::Matrix feat_mat(feat_dim, N); - for (std::size_t i = 0; i < N; ++i) - for (std::size_t j = 0; j < feat_dim; ++j) - feat_mat.set_value_unchecked(j, i, features[i * feat_dim + j]); + const std::size_t N = row; - nn::Matrix label_mat(10, N); - for (std::size_t i = 0; i < N; ++i) - { - int lbl = labels[i]; - if (lbl < 0 || lbl >= 10) - NN_FAIL("Label out of range: " + std::to_string(lbl)); - label_mat.set_value_unchecked(lbl, i, 1.0); - } + nn::Matrix feat_mat(feat_dim, N); + for (std::size_t i = 0; i < N; ++i) + for (std::size_t j = 0; j < feat_dim; ++j) + feat_mat.set_value_unchecked(j, i, features[i * feat_dim + j]); - return std::pair{std::move(feat_mat), std::move(label_mat)}; + nn::Matrix label_mat(10, N); + for (std::size_t i = 0; i < N; ++i) + { + int lbl = labels[i]; + if (lbl < 0 || lbl >= 10) NN_FAIL("Label out of range: " + std::to_string(lbl)); + label_mat.set_value_unchecked(lbl, i, 1.0); } - // ── 评估模型准确率 ───────────────────────────────────────────────────── - // 分块前向 → 每块下载到 CPU 做 argmax → 汇总 top-1 准确率。 - // eval_samples > 0 时只评估前 N 个样本（用于 Transformer 等评估成本高的场景）。 - // eval_batch > 0 时每块 N 个样本前向（**CNN 必须分块**）： - // 卷积层的 im2col 展开是 k²×C_in 倍，全量评估 60000 样本单次 forward 需 - // ~6.4 GB 显存（实测 `vkAllocateMemory failed: -2`）；分块后峰值 ≈ 1/块数， - // 每块结束归还空闲池块（release_idle_pool_blocks），显存池不留高水位。 - // 0 = 不分块（整批单次前向，仅适合 MLP 等小激活模型）。 - // - // 注：签名用 nn::Model& —— nn::Model 不派生自 nn::Layer，且 Model::forward - // 自带 engine 绑定，签名不同于 Layer::forward(engine, ...)； - // 与调用点 model.forward(...) 一致。 - [[nodiscard]] inline nn::Result - evaluate_mnist(nn::Model &model, nn::ComputeEngine &engine, - const nn::Matrix &x, const nn::Matrix &y_onehot, - std::size_t eval_samples = 0, std::size_t eval_batch = 1000) + return std::pair{std::move(feat_mat), std::move(label_mat)}; +} + +// ── 评估模型准确率 ───────────────────────────────────────────────────── +// 分块前向 → 每块下载到 CPU 做 argmax → 汇总 top-1 准确率。 +// eval_samples > 0 时只评估前 N 个样本（用于 Transformer 等评估成本高的场景）。 +// eval_batch > 0 时每块 N 个样本前向（**CNN 必须分块**）： +// 卷积层的 im2col 展开是 k²×C_in 倍，全量评估 60000 样本单次 forward 需 +// ~6.4 GB 显存（实测 `vkAllocateMemory failed: -2`）；分块后峰值 ≈ 1/块数， +// 每块结束归还空闲池块（release_idle_pool_blocks），显存池不留高水位。 +// 0 = 不分块（整批单次前向，仅适合 MLP 等小激活模型）。 +// +// 注：签名用 nn::Model& —— nn::Model 不派生自 nn::Layer，且 Model::forward +// 自带 engine 绑定，签名不同于 Layer::forward(engine, ...)； +// 与调用点 model.forward(...) 一致。 +[[nodiscard]] inline nn::Result +evaluate_mnist(nn::Model& model, nn::ComputeEngine& engine, const nn::Matrix& x, + const nn::Matrix& y_onehot, std::size_t eval_samples = 0, + std::size_t eval_batch = 1000) +{ + const std::size_t N = (eval_samples > 0) ? std::min(x.cols(), eval_samples) : x.cols(); + if (N == 0) NN_FAIL("evaluate_mnist: empty dataset"); + const std::size_t chunk = (eval_batch > 0) ? std::min(eval_batch, N) : N; + + // 评估一律用推理模式：BatchNorm 使用 running 统计量而非 batch 统计量。 + // 评估结束后恢复训练模式（调用方默认为训练场景）。 + model.set_training(false); + + // RAII guard：无论成功与否，函数退出时都恢复训练模式。 + // 避免 begin_batch/forward/end_batch 失败提前返回时，模型被静默留在 + // 推理模式（BatchNorm 会错误使用 running 统计量，破坏后续训练）。 + struct TrainingGuard { - const std::size_t N = - (eval_samples > 0) ? std::min(x.cols(), eval_samples) : x.cols(); - if (N == 0) - NN_FAIL("evaluate_mnist: empty dataset"); - const std::size_t chunk = (eval_batch > 0) ? std::min(eval_batch, N) : N; - - // 评估一律用推理模式：BatchNorm 使用 running 统计量而非 batch 统计量。 - // 评估结束后恢复训练模式（调用方默认为训练场景）。 - model.set_training(false); - - // RAII guard：无论成功与否，函数退出时都恢复训练模式。 - // 避免 begin_batch/forward/end_batch 失败提前返回时，模型被静默留在 - // 推理模式（BatchNorm 会错误使用 running 统计量，破坏后续训练）。 - struct TrainingGuard - { - nn::Model &m; - ~TrainingGuard() { m.set_training(true); } - } training_guard{model}; + nn::Model& m; + ~TrainingGuard() { m.set_training(true); } + } training_guard{model}; - std::size_t correct = 0; - for (std::size_t start = 0; start < N; start += chunk) + std::size_t correct = 0; + for (std::size_t start = 0; start < N; start += chunk) + { + const std::size_t cnt = std::min(chunk, N - start); + + // 取第 [start, start+cnt) 列的子矩阵（CPU 侧拷贝，MNIST 量级很小） + nn::Matrix xb(x.rows(), cnt); + for (std::size_t i = 0; i < cnt; ++i) + for (std::size_t r = 0; r < x.rows(); ++r) + xb.set_value_unchecked(r, i, x.at_unchecked(r, start + i)); + + auto x_tensor_r = engine.from_matrix(xb); + NN_TRY_CHECK(x_tensor_r); + // batch 模式加速 forward（GPU 下消除 per-primitive 提交开销） + auto bb = engine.begin_batch(); + NN_TRY_CHECK(bb); + auto out_tensor_r = model.forward(*x_tensor_r); + if (!out_tensor_r) { - const std::size_t cnt = std::min(chunk, N - start); - - // 取第 [start, start+cnt) 列的子矩阵（CPU 侧拷贝，MNIST 量级很小） - nn::Matrix xb(x.rows(), cnt); - for (std::size_t i = 0; i < cnt; ++i) - for (std::size_t r = 0; r < x.rows(); ++r) - xb.set_value_unchecked(r, i, x.at_unchecked(r, start + i)); - - auto x_tensor_r = engine.from_matrix(xb); - NN_TRY_CHECK(x_tensor_r); - // batch 模式加速 forward（GPU 下消除 per-primitive 提交开销） - auto bb = engine.begin_batch(); - NN_TRY_CHECK(bb); - auto out_tensor_r = model.forward(*x_tensor_r); - if (!out_tensor_r) - { - (void)engine.end_batch(); // 出错也收尾，避免录制状态泄漏 - NN_TRY_CHECK(out_tensor_r); - } + (void)engine.end_batch(); // 出错也收尾，避免录制状态泄漏 + NN_TRY_CHECK(out_tensor_r); + } - auto eb = engine.end_batch(); - NN_TRY_CHECK(eb); - auto out_r = engine.to_matrix(*out_tensor_r); - NN_TRY_CHECK(out_r); - const auto &out = *out_r; + auto eb = engine.end_batch(); + NN_TRY_CHECK(eb); + auto out_r = engine.to_matrix(*out_tensor_r); + NN_TRY_CHECK(out_r); + const auto& out = *out_r; - for (std::size_t i = 0; i < cnt; ++i) + for (std::size_t i = 0; i < cnt; ++i) + { + nn::Scalar max_val = out.at_unchecked(0, i); + int pred = 0; + for (int j = 1; j < static_cast(nn::MNIST_NUM_CLASSES); ++j) { - nn::Scalar max_val = out.at_unchecked(0, i); - int pred = 0; - for (int j = 1; j < static_cast(nn::MNIST_NUM_CLASSES); ++j) + nn::Scalar val = out.at_unchecked(j, i); + if (val > max_val) { - nn::Scalar val = out.at_unchecked(j, i); - if (val > max_val) - { - max_val = val; - pred = j; - } + max_val = val; + pred = j; } - int true_label = -1; - for (int j = 0; j < static_cast(nn::MNIST_NUM_CLASSES); ++j) + } + int true_label = -1; + for (int j = 0; j < static_cast(nn::MNIST_NUM_CLASSES); ++j) + { + if (y_onehot.at_unchecked(j, start + i) == 1.0) { - if (y_onehot.at_unchecked(j, start + i) == 1.0) - { - true_label = j; - break; - } + true_label = j; + break; } - if (pred == true_label) - ++correct; } - - // 归还本块峰值产生的空闲池块：否则显存池高水位会留在 device - // （CPU 引擎 no-op）。 - (void)engine.release_idle_pool_blocks(); + if (pred == true_label) ++correct; } - return static_cast(correct) / static_cast(N); + + // 归还本块峰值产生的空闲池块：否则显存池高水位会留在 device + // （CPU 引擎 no-op）。 + (void)engine.release_idle_pool_blocks(); } + return static_cast(correct) / static_cast(N); +} } // namespace nn::cli diff --git a/include/neuralnet.cpp/cli/cli_precision_warning.hpp b/include/neuralnet.cpp/cli/cli_precision_warning.hpp index 679f7e28..8c335fa9 100644 --- a/include/neuralnet.cpp/cli/cli_precision_warning.hpp +++ b/include/neuralnet.cpp/cli/cli_precision_warning.hpp @@ -24,33 +24,32 @@ namespace nn::cli { - // 危险精度组合启动期告警文本（无危险组合 → 空串）。 - inline std::string dangerous_precision_warning(const nn::PrecisionProfile& p) - { - using nn::Precision; +// 危险精度组合启动期告警文本（无危险组合 → 空串）。 +inline std::string dangerous_precision_warning(const nn::PrecisionProfile& p) +{ + using nn::Precision; - // 唯一触发字段：优化器状态精度 = F16（Adam 的 m/v 下溢 → 更新爆炸）。 - if (p.optimizer != Precision::F16) - return {}; + // 唯一触发字段：优化器状态精度 = F16（Adam 的 m/v 下溢 → 更新爆炸）。 + if (p.optimizer != Precision::F16) return {}; - const bool all_f16 = p.param == Precision::F16 && p.compute == Precision::F16 && - p.stable == Precision::F16 && p.optimizer == Precision::F16; + const bool all_f16 = p.param == Precision::F16 && p.compute == Precision::F16 && + p.stable == Precision::F16 && p.optimizer == Precision::F16; - std::string w; - w += "[精度警告] 检测到危险精度组合（仅提醒，不阻断训练，详见 " - "docs/internal/24-f16-training-defects.md §2）：\n"; - w += std::string(" · optimizer=") + nn::precision_name(p.optimizer) + - "：Adam 二阶矩 v≈g²~1e-10 在 f16 下溢为 0 → 更新量 -lr·m̂/(√v̂+eps) " - "爆炸 → loss 发散/NaN。\n"; - if (p.stable == Precision::F16) - w += std::string(" · stable=") + nn::precision_name(p.stable) + - "：CE/softmax 归约在 f16 下范围(65504)/尾数(10 位)不足 → 大词表 " - "~200 步 NaN。\n"; - if (all_f16) - w += " · param/compute/stable/optimizer 全 f16：更新量 |delta| 可能小于" - "参数所在的 f16 网格步长 → 原地写回被舍入吃光（loss 恒定，数值性冻结）。\n"; - w += " 安全配方：`--f16`（param+compute=f16，stable/optimizer=f32）。" - "如需继续实验请自行留意数值稳定性。\n"; - return w; - } + std::string w; + w += "[精度警告] 检测到危险精度组合（仅提醒，不阻断训练，详见 " + "docs/internal/24-f16-training-defects.md §2）：\n"; + w += std::string(" · optimizer=") + nn::precision_name(p.optimizer) + + "：Adam 二阶矩 v≈g²~1e-10 在 f16 下溢为 0 → 更新量 -lr·m̂/(√v̂+eps) " + "爆炸 → loss 发散/NaN。\n"; + if (p.stable == Precision::F16) + w += std::string(" · stable=") + nn::precision_name(p.stable) + + "：CE/softmax 归约在 f16 下范围(65504)/尾数(10 位)不足 → 大词表 " + "~200 步 NaN。\n"; + if (all_f16) + w += " · param/compute/stable/optimizer 全 f16：更新量 |delta| 可能小于" + "参数所在的 f16 网格步长 → 原地写回被舍入吃光（loss 恒定，数值性冻结）。\n"; + w += " 安全配方：`--f16`（param+compute=f16，stable/optimizer=f32）。" + "如需继续实验请自行留意数值稳定性。\n"; + return w; +} } // namespace nn::cli diff --git a/include/neuralnet.cpp/cli/cli_train_common.hpp b/include/neuralnet.cpp/cli/cli_train_common.hpp index 51e7a546..ec945e23 100644 --- a/include/neuralnet.cpp/cli/cli_train_common.hpp +++ b/include/neuralnet.cpp/cli/cli_train_common.hpp @@ -41,125 +41,158 @@ namespace nn::cli { - // ── 通用训练参数 ─────────────────────────────────────────────────────── - // 默认值仅作为结构体占位；调用方应根据程序语义在解析前覆盖。 - struct TrainCommonArgs - { - int epochs = 10; - nn::Scalar lr = 0.001f; - std::size_t batch_size = 32; - std::string optimizer = "adam"; - nn::Scalar weight_decay = 0.0f; - nn::Scalar min_lr = 0.0f; - int warmup_epochs = 0; - std::string lr_schedule = "constant"; // "fixed"/"constant" 等价，"cosine" 退火 - std::vector lr_per_epoch; - bool use_gpu = false; - // --gpu 的可选设备选择子（索引或名称子串）；空 = 自动选卡 - std::string gpu_device; - }; +// ── 通用训练参数 ─────────────────────────────────────────────────────── +// 默认值仅作为结构体占位；调用方应根据程序语义在解析前覆盖。 +struct TrainCommonArgs +{ + int epochs = 10; + nn::Scalar lr = 0.001f; + std::size_t batch_size = 32; + std::string optimizer = "adam"; + nn::Scalar weight_decay = 0.0f; + nn::Scalar min_lr = 0.0f; + int warmup_epochs = 0; + std::string lr_schedule = "constant"; // "fixed"/"constant" 等价，"cosine" 退火 + std::vector lr_per_epoch; + bool use_gpu = false; + // --gpu 的可选设备选择子（索引或名称子串）；空 = 自动选卡 + std::string gpu_device; +}; - // ── 解析通用训练参数 ─────────────────────────────────────────────────── - // 在 argv[i] 处尝试解析一个通用 flag。 - // 成功消费时返回 true 并将 i 递增到本次消费的最后一个 argv 索引； - // 不是通用 flag 时返回 false（i 不变），由调用方继续判断。 - // 解析失败时打印错误并 std::exit(1)（与现有 parse_args 行为一致）。 - inline bool parse_train_common_args( - int argc, char *argv[], int &i, TrainCommonArgs &cfg) - { - std::string arg = argv[i]; +// ── 解析通用训练参数 ─────────────────────────────────────────────────── +// 在 argv[i] 处尝试解析一个通用 flag。 +// 成功消费时返回 true 并将 i 递增到本次消费的最后一个 argv 索引； +// 不是通用 flag 时返回 false（i 不变），由调用方继续判断。 +// 解析失败时打印错误并 std::exit(1)（与现有 parse_args 行为一致）。 +inline bool parse_train_common_args(int argc, char* argv[], int& i, TrainCommonArgs& cfg) +{ + std::string arg = argv[i]; - if (arg == "--epochs" && i + 1 < argc) + if (arg == "--epochs" && i + 1 < argc) + { + auto v = nn::parse_number(argv[++i]); + if (!v) { - auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --epochs: " << v.error().message << "\n"; std::exit(1); } - if (*v <= 0) { std::cerr << "--epochs 必须为正整数\n"; std::exit(1); } - cfg.epochs = *v; - return true; + std::cerr << "无效 --epochs: " << v.error().message << "\n"; + std::exit(1); } - if (arg == "--lr" && i + 1 < argc) + if (*v <= 0) { - auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --lr: " << v.error().message << "\n"; std::exit(1); } - cfg.lr = *v; - return true; + std::cerr << "--epochs 必须为正整数\n"; + std::exit(1); } - if (arg == "--batch-size" && i + 1 < argc) + cfg.epochs = *v; + return true; + } + if (arg == "--lr" && i + 1 < argc) + { + auto v = nn::parse_number(argv[++i]); + if (!v) { - auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --batch-size: " << v.error().message << "\n"; std::exit(1); } - if (*v == 0) { std::cerr << "--batch-size 必须为正整数\n"; std::exit(1); } - cfg.batch_size = *v; - return true; + std::cerr << "无效 --lr: " << v.error().message << "\n"; + std::exit(1); } - if (arg == "--optimizer" && i + 1 < argc) + cfg.lr = *v; + return true; + } + if (arg == "--batch-size" && i + 1 < argc) + { + auto v = nn::parse_number(argv[++i]); + if (!v) { - cfg.optimizer = argv[++i]; - if (cfg.optimizer != "sgd" && cfg.optimizer != "sgd_momentum" && - cfg.optimizer != "adam" && cfg.optimizer != "adamw" && - cfg.optimizer != "muon") - { - std::cerr << "未知优化器: " << cfg.optimizer - << "，可选: sgd, sgd_momentum, adam, adamw, muon\n"; - std::exit(1); - } - return true; + std::cerr << "无效 --batch-size: " << v.error().message << "\n"; + std::exit(1); } - if (arg == "--weight-decay" && i + 1 < argc) + if (*v == 0) + { + std::cerr << "--batch-size 必须为正整数\n"; + std::exit(1); + } + cfg.batch_size = *v; + return true; + } + if (arg == "--optimizer" && i + 1 < argc) + { + cfg.optimizer = argv[++i]; + if (cfg.optimizer != "sgd" && cfg.optimizer != "sgd_momentum" && cfg.optimizer != "adam" && + cfg.optimizer != "adamw" && cfg.optimizer != "muon") { - auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --weight-decay: " << v.error().message << "\n"; std::exit(1); } - cfg.weight_decay = *v; - return true; + std::cerr << "未知优化器: " << cfg.optimizer + << "，可选: sgd, sgd_momentum, adam, adamw, muon\n"; + std::exit(1); } - if (arg == "--min-lr" && i + 1 < argc) + return true; + } + if (arg == "--weight-decay" && i + 1 < argc) + { + auto v = nn::parse_number(argv[++i]); + if (!v) { - auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --min-lr: " << v.error().message << "\n"; std::exit(1); } - cfg.min_lr = *v; - return true; + std::cerr << "无效 --weight-decay: " << v.error().message << "\n"; + std::exit(1); } - if (arg == "--warmup-epochs" && i + 1 < argc) + cfg.weight_decay = *v; + return true; + } + if (arg == "--min-lr" && i + 1 < argc) + { + auto v = nn::parse_number(argv[++i]); + if (!v) { - auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --warmup-epochs: " << v.error().message << "\n"; std::exit(1); } - cfg.warmup_epochs = *v; - return true; + std::cerr << "无效 --min-lr: " << v.error().message << "\n"; + std::exit(1); } - if (arg == "--lr-per-epoch" && i + 1 < argc) + cfg.min_lr = *v; + return true; + } + if (arg == "--warmup-epochs" && i + 1 < argc) + { + auto v = nn::parse_number(argv[++i]); + if (!v) { - std::string dims_str = argv[++i]; - std::stringstream ss(dims_str); - std::string token; - cfg.lr_per_epoch.clear(); - while (std::getline(ss, token, ',')) - { - auto v = nn::parse_number(token); - if (!v) { std::cerr << "无效 --lr-per-epoch 值: " << v.error().message << "\n"; std::exit(1); } - cfg.lr_per_epoch.push_back(*v); - } - return true; + std::cerr << "无效 --warmup-epochs: " << v.error().message << "\n"; + std::exit(1); } - if (arg == "--lr-schedule" && i + 1 < argc) + cfg.warmup_epochs = *v; + return true; + } + if (arg == "--lr-per-epoch" && i + 1 < argc) + { + std::string dims_str = argv[++i]; + std::stringstream ss(dims_str); + std::string token; + cfg.lr_per_epoch.clear(); + while (std::getline(ss, token, ',')) { - cfg.lr_schedule = argv[++i]; - // "fixed" 与 "constant" 等价（同一调度分支）；"cosine" 为余弦退火 - if (cfg.lr_schedule != "fixed" && cfg.lr_schedule != "constant" && - cfg.lr_schedule != "cosine") + auto v = nn::parse_number(token); + if (!v) { - std::cerr << "未知 --lr-schedule: " << cfg.lr_schedule - << "，可选: fixed, constant, cosine\n"; + std::cerr << "无效 --lr-per-epoch 值: " << v.error().message << "\n"; std::exit(1); } - return true; + cfg.lr_per_epoch.push_back(*v); } - if (auto dev = parse_gpu_option(argc, argv, i)) + return true; + } + if (arg == "--lr-schedule" && i + 1 < argc) + { + cfg.lr_schedule = argv[++i]; + // "fixed" 与 "constant" 等价（同一调度分支）；"cosine" 为余弦退火 + if (cfg.lr_schedule != "fixed" && cfg.lr_schedule != "constant" && + cfg.lr_schedule != "cosine") { - cfg.use_gpu = true; - if (!dev->empty()) - cfg.gpu_device = *dev; - return true; + std::cerr << "未知 --lr-schedule: " << cfg.lr_schedule + << "，可选: fixed, constant, cosine\n"; + std::exit(1); } - return false; // 不是通用 flag + return true; + } + if (auto dev = parse_gpu_option(argc, argv, i)) + { + cfg.use_gpu = true; + if (!dev->empty()) cfg.gpu_device = *dev; + return true; } + return false; // 不是通用 flag +} } // namespace nn::cli diff --git a/include/neuralnet.cpp/cli/text_train_checkpoint.hpp b/include/neuralnet.cpp/cli/text_train_checkpoint.hpp index 30f8db59..ef4205ec 100644 --- a/include/neuralnet.cpp/cli/text_train_checkpoint.hpp +++ b/include/neuralnet.cpp/cli/text_train_checkpoint.hpp @@ -55,12 +55,11 @@ class TextTrainCheckpointLog std::span optimizer_blobs, std::ostream& err) { const nn::CheckpointState ckpt{&state, optimizer_blobs}; - auto r = nn::save_model(path, model, spec, tokenizer_bytes, - nn::ModelRole::Checkpoint, ckpt); + auto r = + nn::save_model(path, model, spec, tokenizer_bytes, nn::ModelRole::Checkpoint, ckpt); if (r) return; ++failures_; - err << "\n [ckpt] 保存失败（第 " << failures_ << " 次）: " << r.error().message - << "\n"; + err << "\n [ckpt] 保存失败（第 " << failures_ << " 次）: " << r.error().message << "\n"; } // 状态导出失败等"这一份断点没写出去"的事件：与写失败同一记账口径 @@ -90,10 +89,9 @@ class TextTrainCheckpointLog // ── 推理模型产出（Model 角色）────────────────────────────────────────── // 只含推理所需参数（三值层写打包形态 + τ）；失败即失败：写不出去 = 这一轮训练 // 对用户没有产出，调用方据此非零退出。 -[[nodiscard]] inline nn::Result save_text_train_model( - const std::string& path, nn::Model& model, const nn::ModelSpec& spec, - const std::string& tokenizer_bytes) -{ - return nn::save_model(path, model, spec, tokenizer_bytes, nn::ModelRole::Model); -} +[[nodiscard]] inline nn::Result save_text_train_model(const std::string& path, + nn::Model& model, + const nn::ModelSpec& spec, + const std::string& tokenizer_bytes) +{ return nn::save_model(path, model, spec, tokenizer_bytes, nn::ModelRole::Model); } } // namespace nn::cli diff --git a/include/neuralnet.cpp/cli/text_train_config.hpp b/include/neuralnet.cpp/cli/text_train_config.hpp index 7030b451..1e023a70 100644 --- a/include/neuralnet.cpp/cli/text_train_config.hpp +++ b/include/neuralnet.cpp/cli/text_train_config.hpp @@ -46,7 +46,7 @@ using nn::Scalar; // 不兼容组合由 `Layer::init` 的能力校验给出"层 X 不支持 param=t1_58"的报错， // 这正是设计意图（§4.1）—— CLI 不加模型白名单。 nn::Precision parse_precision_arg(const std::string& name, const char* flag, - bool allow_t1_58 = false) + bool allow_t1_58 = false) { if (name == "f16" || name == "half") return nn::Precision::F16; if (name == "f32" || name == "float") return nn::Precision::F32; @@ -57,83 +57,121 @@ nn::Precision parse_precision_arg(const std::string& name, const char* flag, << "，可选: f16, f32\n"; std::exit(1); } - std::cerr << "无效 --" << flag << ": " << name << "，可选: " - << (allow_t1_58 ? "f16, f32, t1_58" : "f16, f32") << "\n"; + std::cerr << "无效 --" << flag << ": " << name + << "，可选: " << (allow_t1_58 ? "f16, f32, t1_58" : "f16, f32") << "\n"; std::exit(1); } // ==================== 帮助信息 ==================== -void print_text_train_usage(const char *prog) +void print_text_train_usage(const char* prog) { nn::cli::Help help(std::cout, prog, "GPT 文本生成训练程序"); help.usage(" [选项]"); help.section("参数"); - help.item("", "数据集文件路径 (.nndataset，由 dataset_gen 生成；\n词表/loss 掩码/test 子集都在生成期定好)"); + help.item("", "数据集文件路径 (.nndataset，由 dataset_gen 生成；\n词表/loss " + "掩码/test 子集都在生成期定好)"); help.section("选项"); - help.opt("--save ", "推理模型保存路径 (默认: gpt_model.nnmodel；\n只含推理所需参数——三值层写打包形态，不含训练参数)"); - help.opt("--resume ", "从模型文件恢复训练：\n断点 .nnckpt = 断点续训（形状/配方/计数器由断点决定，二者不得再由命令行给出）；\n推理模型 .nnmodel = warm start（SFT：CLI 全权、优化器冷启动；\n三值模型的打包权重按文件 τ 反量化成 f32 latent 后可继续训练）"); + help.opt("--save ", + "推理模型保存路径 (默认: " + "gpt_model.nnmodel；\n只含推理所需参数——三值层写打包形态，不含训练参数)"); + help.opt("--resume ", + "从模型文件恢复训练：\n断点 .nnckpt = " + "断点续训（形状/配方/计数器由断点决定，二者不得再由命令行给出）；\n推理模型 .nnmodel " + "= warm start（SFT：CLI 全权、优化器冷启动；\n三值模型的打包权重按文件 τ 反量化成 f32 " + "latent 后可继续训练）"); help.opt("--resume-epoch ", "从第 n 个 epoch 继续 (0-based，需配合 --resume；默认 0)"); help.opt("--resume-step ", "从本 epoch 内第 n 步继续 (0-based，需配合 --resume；默认 0)"); help.opt("--epochs ", "训练轮数 (默认: 10)"); help.opt("--lr ", "学习率 (默认: 0.001)"); help.opt("--batch-size ", "理论 global batch / 每次参数更新的样本数 (默认: 32)"); - help.opt("--accum-steps ", "梯度累积次数 (默认: 1)\n实际 micro-batch = batch-size / accum-steps；必须整除"); + help.opt("--accum-steps ", + "梯度累积次数 (默认: 1)\n实际 micro-batch = batch-size / accum-steps；必须整除"); help.opt("--seq-len ", "序列长度 (默认: 256)"); - help.opt("--stride ", "滑动窗口步长 (默认: 等于 --seq-len，即不重叠)\n设小可产生重叠窗口，增加训练样本数"); + help.opt("--stride ", + "滑动窗口步长 (默认: 等于 --seq-len，即不重叠)\n设小可产生重叠窗口，增加训练样本数"); // RV-04：名单文本由 nn::optimizer_table()（唯一真源）生成，不再手抄 - help.opt("--optimizer ", - "优化器: " + nn::optimizer_names("/") + - " (默认: adam)\nadamwquant: 量化状态的 AdamW（状态 4 B/权重 vs 8 B；见 docs/internal/21 §4.13）"); + help.opt("--optimizer ", "优化器: " + nn::optimizer_names("/") + + " (默认: adam)\nadamwquant: 量化状态的 AdamW（状态 4 B/权重 " + "vs 8 B；见 docs/internal/21 §4.13）"); help.opt("--weight-decay ", "AdamW 权重衰减系数 (默认: 0.01)"); help.opt("--beta1 ", "Adam/AdamW 一阶动量衰减 β1 (默认: 0.9，sgd/muon 忽略)"); - help.opt("--beta2 ", "Adam/AdamW 二阶动量衰减 β2 (默认: 0.999，sgd/muon 忽略)\nLLM 预训练惯例常取 0.95 (GPT-3/LLaMA 配方)"); + help.opt("--beta2 ", "Adam/AdamW 二阶动量衰减 β2 (默认: 0.999，sgd/muon 忽略)\nLLM " + "预训练惯例常取 0.95 (GPT-3/LLaMA 配方)"); help.opt("--max-norm ", "梯度裁剪最大全局 L2 范数 (默认: 0=不裁剪)"); help.opt("--log-interval ", "每隔多少 step 显示进度 (默认: 50)"); - help.opt("--save-interval ", "每隔多少 update step 保存训练断点 .nnckpt (默认: 100；\n文件名 = --save 的路径换 .nnckpt 扩展名)"); - help.opt("--max-steps ", "本次运行最多执行多少 update step (默认: 0=不限)\n按本进程参数更新次数计，可与 --resume 组合分段跑；用于吞吐实测/分段训练"); + help.opt("--save-interval ", "每隔多少 update step 保存训练断点 .nnckpt (默认: " + "100；\n文件名 = --save 的路径换 .nnckpt 扩展名)"); + help.opt("--max-steps ", + "本次运行最多执行多少 update step (默认: 0=不限)\n按本进程参数更新次数计，可与 " + "--resume 组合分段跑；用于吞吐实测/分段训练"); help.opt("--grad-log", "显示梯度统计 (范数/最大值/均值)"); - help.opt("--gpu [索引|名称]", "启用 GPU 加速 (需要 Vulkan SDK)\n空格形式只收枚举索引；名称子串用 --gpu=<名称>，如 --gpu=NVIDIA / --gpu=40HX"); + help.opt("--gpu [索引|名称]", + "启用 GPU 加速 (需要 Vulkan SDK)\n空格形式只收枚举索引；名称子串用 --gpu=<名称>，如 " + "--gpu=NVIDIA / --gpu=40HX"); help.opt("--help, -h", "显示此帮助信息"); help.section("模型结构"); - help.opt("--model ", "模型架构: gpt/rapt (默认: gpt)\ngpt: 标准 Transformer 语言模型\nrapt: ReLU 线性注意力 (RLA) 语言模型，动态稀疏检索，\nO(L·d²) 复杂度；强制 RoPE (ReLU 前施加)"); - help.opt("--positional-encoding ", "位置编码类型: learned/sinusoidal/alibi/rope (默认: learned)\nlearned: 可学习位置嵌入 (GPT 原版)\nsinusoidal: 正弦波固定位置编码 (不参与训练)\nalibi: 线性偏置注意力 (无位置嵌入，支持长度外推)\nrope: 旋转位置编码 (现代方案，在注意力 Q/K 上施加)"); - help.opt("--activation ", "FFN 激活: gelu/swiglu/relu2 (默认: gelu)\ngelu: QuickGeLU (GPT-2 风格)\nswiglu: SwiGLU (LLaMA/Mistral 风格，每参数效率更高)\nrelu2: 门控平方 ReLU (BitNet b1.58 2B4T 的 hidden_act)"); - help.opt("--norm ", "归一化层: layernorm/rmsnorm (默认: layernorm)\nlayernorm: LayerNorm (GPT-2 风格)\nrmsnorm: RMSNorm (LLaMA/Mistral 风格，更快更稳)"); + help.opt("--model ", + "模型架构: gpt/rapt (默认: gpt)\ngpt: 标准 Transformer 语言模型\nrapt: ReLU " + "线性注意力 (RLA) 语言模型，动态稀疏检索，\nO(L·d²) 复杂度；强制 RoPE (ReLU 前施加)"); + help.opt("--positional-encoding ", + "位置编码类型: learned/sinusoidal/alibi/rope (默认: learned)\nlearned: 可学习位置嵌入 " + "(GPT 原版)\nsinusoidal: 正弦波固定位置编码 (不参与训练)\nalibi: 线性偏置注意力 " + "(无位置嵌入，支持长度外推)\nrope: 旋转位置编码 (现代方案，在注意力 Q/K 上施加)"); + help.opt("--activation ", + "FFN 激活: gelu/swiglu/relu2 (默认: gelu)\ngelu: QuickGeLU (GPT-2 风格)\nswiglu: " + "SwiGLU (LLaMA/Mistral 风格，每参数效率更高)\nrelu2: 门控平方 ReLU (BitNet b1.58 2B4T " + "的 hidden_act)"); + help.opt("--norm ", "归一化层: layernorm/rmsnorm (默认: layernorm)\nlayernorm: LayerNorm " + "(GPT-2 风格)\nrmsnorm: RMSNorm (LLaMA/Mistral 风格，更快更稳)"); help.opt("--d-model ", "模型维度 (默认: 128)"); help.opt("--num-heads ", "注意力头数 (默认: 4)"); help.opt("--num-layers ", "Transformer 层数 (默认: 4)"); help.opt("--d-ff ", "FFN 中间维度 (默认: 512)"); help.section("学习率调度"); - help.opt("--lr-schedule ", "学习率调度: fixed/cosine/step_cosine (默认: fixed)\ncosine: 余弦退火 (epoch 级)，lr 从初始值衰减到 min-lr\nstep_cosine: 余弦退火 (step 级)，按单个训练步预热+退火，\n适合每 epoch 步数很多的场景 (如大语料)"); + help.opt("--lr-schedule ", + "学习率调度: fixed/cosine/step_cosine (默认: fixed)\ncosine: 余弦退火 (epoch 级)，lr " + "从初始值衰减到 min-lr\nstep_cosine: 余弦退火 (step " + "级)，按单个训练步预热+退火，\n适合每 epoch 步数很多的场景 (如大语料)"); help.opt("--warmup-epochs ", "线性预热轮数 (默认: 0，仅 cosine)"); help.opt("--warmup-steps ", "线性预热步数 (默认: 0，仅 step_cosine)"); help.opt("--min-lr ", "余弦退火最低学习率 (默认: 1e-6)"); help.opt("--lr-per-epoch ", "手动指定每轮学习率 (逗号分隔，优先级最高)"); help.section("显存优化"); - help.opt("--checkpoint-every ", "每 N 个 Transformer block 重算一次 forward (激活重计算，默认: 0=不启用)\n1=每块都重算，显存收益最大，以约 1 次额外前向\nFLOPs 为代价省去整层激活驻留"); - help.opt("--activation-offload", "把每块激活搬 host-visible，backward 拷回\n不重算，FLOPs 保持 1.0×，代价是 PCIe 传输；可与\n--checkpoint-every 混合使用：checkpoint 块重算、其余块 offload；仅 GPU 有效；\n与 --f16 / --t1_58 / 四槽精度参数不兼容（slab 恒 f32，配置期拒绝）"); + help.opt( + "--checkpoint-every ", + "每 N 个 Transformer block 重算一次 forward (激活重计算，默认: " + "0=不启用)\n1=每块都重算，显存收益最大，以约 1 次额外前向\nFLOPs 为代价省去整层激活驻留"); + help.opt("--activation-offload", + "把每块激活搬 host-visible，backward 拷回\n不重算，FLOPs 保持 1.0×，代价是 PCIe " + "传输；可与\n--checkpoint-every 混合使用：checkpoint 块重算、其余块 offload；仅 GPU " + "有效；\n与 --f16 / --t1_58 / 四槽精度参数不兼容（slab 恒 f32，配置期拒绝）"); help.note("GPU 训练每 step 末尾自动归还完全空闲的内存池底材 (L2 整块释放)。"); help.section("Batch 录制粒度"); help.opt("--flush-interval ", "每 N 个 Transformer block flush 一次 (默认: 1)"); - help.note("按层切 batch 缩短 D1 延迟销毁锁窗、拆分大提交防 TDR；\n非阻塞提交后细粒度 flush 的额外 submit 代价不在关键路径，\n不影响 batch_size 和训练质量"); + help.note("按层切 batch 缩短 D1 延迟销毁锁窗、拆分大提交防 TDR；\n非阻塞提交后细粒度 flush " + "的额外 submit 代价不在关键路径，\n不影响 batch_size 和训练质量"); help.section("混合精度 (docs/development/05-mixed-precision.md)"); help.opt("--f16", "快捷方式: f16 存储 (param/compute=F16，stable/optimizer=F32)"); - help.opt("--t1_58", "快捷方式: 精度预设 t1_58 = f16 存储 + 三值权重\n(param=T1_58，compute=F16，stable/optimizer=F32；等价于 --f16 再把 param 换成 t1_58)。\n范围与前置拒绝同 --precision-param t1_58（docs/internal/21-quantized-weights.md §4.10）"); + help.opt("--t1_58", "快捷方式: 精度预设 t1_58 = f16 存储 + " + "三值权重\n(param=T1_58，compute=F16，stable/optimizer=F32；等价于 --f16 " + "再把 param 换成 t1_58)。\n范围与前置拒绝同 --precision-param " + "t1_58（docs/internal/21-quantized-weights.md §4.10）"); help.opt("--precision-param ", "权重/参数精度 (默认: f32)；t1_58 = 三值权重（1.58-bit）——\n" "GPT 已支持（注意力 4 投影 + fc1/fc2 + 未 tie 的 LM head 全部换成\n" "BitLinear；P1.5 落地）；--model rapt 选它会被前置拒绝\n" "(docs/internal/21-quantized-weights.md §4.8.3)"); help.opt("--precision-compute ", "常规算子计算精度 (matmul/逐元素/gather，默认: f32)"); - help.opt("--precision-stable ", "数值敏感算子精度 (softmax/LayerNorm/loss，默认: f32)"); + help.opt("--precision-stable ", + "数值敏感算子精度 (softmax/LayerNorm/loss，默认: f32)"); help.opt("--precision-optimizer ", "优化器状态精度 (Adam m/v，默认: f32)"); } struct TrainConfig : nn::TextTrainConfig @@ -184,20 +222,42 @@ struct TrainConfig : nn::TextTrainConfig // `--checkpoint-every` / `--activation-offload` / `--flush-interval` / `--resume` // 本身（含数据集路径 = 位置）、以及探针环境变量。 inline constexpr std::string_view kForbiddenResumeFlags[] = { - "--resume-epoch", "--resume-step", - "--epochs", "--lr", "--batch-size", "--accum-steps", "--seq-len", "--stride", - "--optimizer", "--weight-decay", "--beta1", "--beta2", - "--d-model", "--num-heads", "--num-layers", "--d-ff", "--model", - "--f16", "--t1_58", - "--precision-param", "--precision-compute", "--precision-stable", "--precision-optimizer", - "--lr-schedule", "--warmup-epochs", "--warmup-steps", "--min-lr", "--lr-per-epoch", - "--max-norm", "--positional-encoding", "--activation", "--norm", + "--resume-epoch", + "--resume-step", + "--epochs", + "--lr", + "--batch-size", + "--accum-steps", + "--seq-len", + "--stride", + "--optimizer", + "--weight-decay", + "--beta1", + "--beta2", + "--d-model", + "--num-heads", + "--num-layers", + "--d-ff", + "--model", + "--f16", + "--t1_58", + "--precision-param", + "--precision-compute", + "--precision-stable", + "--precision-optimizer", + "--lr-schedule", + "--warmup-epochs", + "--warmup-steps", + "--min-lr", + "--lr-per-epoch", + "--max-norm", + "--positional-encoding", + "--activation", + "--norm", }; [[nodiscard]] constexpr std::span forbidden_resume_flags() noexcept -{ - return kForbiddenResumeFlags; -} +{ return kForbiddenResumeFlags; } [[nodiscard]] inline bool is_forbidden_resume_flag(std::string_view arg) noexcept { @@ -230,22 +290,19 @@ inline constexpr std::string_view kForbiddenResumeFlags[] = { // 决定出口（argparse 层 `std::exit`；`run_text_train` 经 `NN_FAIL` 走 `Result`， // 保持"三模块不终止进程"不变量）。 [[nodiscard]] inline std::optional -precision_offload_combo_error(const TrainConfig &cfg) +precision_offload_combo_error(const TrainConfig& cfg) { if (!cfg.activation_offload || !cfg.gpu_enabled) return std::nullopt; - if (cfg.precision.param == nn::Precision::F32 && - cfg.precision.compute == nn::Precision::F32 && - cfg.precision.stable == nn::Precision::F32 && - cfg.precision.optimizer == nn::Precision::F32) + if (cfg.precision.param == nn::Precision::F32 && cfg.precision.compute == nn::Precision::F32 && + cfg.precision.stable == nn::Precision::F32 && cfg.precision.optimizer == nn::Precision::F32) return std::nullopt; - return std::string( - "--activation-offload 与混合精度/三值精度不兼容（--f16 / --t1_58 /" - " 四槽精度参数）：\n" - " offload slab 恒 f32，f16 激活 restore 后会变成 f32（精度语义漂移）。\n" - " 请改用 --checkpoint-every，或去掉 --activation-offload。\n"); + return std::string("--activation-offload 与混合精度/三值精度不兼容（--f16 / --t1_58 /" + " 四槽精度参数）：\n" + " offload slab 恒 f32，f16 激活 restore 后会变成 f32（精度语义漂移）。\n" + " 请改用 --checkpoint-every，或去掉 --activation-offload。\n"); } -TrainConfig parse_text_train_args(int argc, char *argv[]) +TrainConfig parse_text_train_args(int argc, char* argv[]) { TrainConfig cfg; // 配方默认值继承自 nn::TextTrainConfig；seq_len 的 CLI 默认是 GPT 序列长度， @@ -257,15 +314,13 @@ TrainConfig parse_text_train_args(int argc, char *argv[]) // ── A/B 类旗标登记（P1-2 断点续训禁列；按**旗标名**判定、不看值）── // 解析层的默认值与用户输入无法区分，故必须在解析期记录"本进程真出现过" // 的旗标；续训路径在断点带训练状态段时发现列表非空即报错。 - if (nn::cli::is_forbidden_resume_flag(arg)) - cfg.explicit_train_flags.push_back(arg); + if (nn::cli::is_forbidden_resume_flag(arg)) cfg.explicit_train_flags.push_back(arg); if (arg == "--help" || arg == "-h") { print_text_train_usage(argv[0]); std::exit(0); } - else if (arg == "--save" && i + 1 < argc) - cfg.save_path = argv[++i]; + else if (arg == "--save" && i + 1 < argc) cfg.save_path = argv[++i]; else if (arg == "--resume" && i + 1 < argc) { cfg.resume_path = argv[++i]; @@ -274,51 +329,91 @@ TrainConfig parse_text_train_args(int argc, char *argv[]) else if (arg == "--resume-epoch" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v || *v < 0) { std::cerr << "无效 --resume-epoch: " << argv[i] << "\n"; std::exit(1); } + if (!v || *v < 0) + { + std::cerr << "无效 --resume-epoch: " << argv[i] << "\n"; + std::exit(1); + } cfg.start_epoch = *v; } else if (arg == "--resume-step" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --resume-step: " << argv[i] << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --resume-step: " << argv[i] << "\n"; + std::exit(1); + } cfg.start_step = *v; } else if (arg == "--epochs" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --epochs: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --epochs: " << v.error().message << "\n"; + std::exit(1); + } cfg.epochs = *v; } else if (arg == "--lr" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --lr: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --lr: " << v.error().message << "\n"; + std::exit(1); + } cfg.lr = *v; } else if (arg == "--batch-size" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --batch-size: " << v.error().message << "\n"; std::exit(1); } - if (*v == 0) { std::cerr << "--batch-size 必须 >= 1\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --batch-size: " << v.error().message << "\n"; + std::exit(1); + } + if (*v == 0) + { + std::cerr << "--batch-size 必须 >= 1\n"; + std::exit(1); + } cfg.batch_size = *v; } else if (arg == "--accum-steps" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --accum-steps: " << v.error().message << "\n"; std::exit(1); } - if (*v == 0) { std::cerr << "--accum-steps 必须 >= 1\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --accum-steps: " << v.error().message << "\n"; + std::exit(1); + } + if (*v == 0) + { + std::cerr << "--accum-steps 必须 >= 1\n"; + std::exit(1); + } cfg.accum_steps = *v; } else if (arg == "--seq-len" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --seq-len: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --seq-len: " << v.error().message << "\n"; + std::exit(1); + } cfg.seq_len = *v; } else if (arg == "--stride" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --stride: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --stride: " << v.error().message << "\n"; + std::exit(1); + } cfg.stride = *v; } else if (arg == "--optimizer" && i + 1 < argc) @@ -335,43 +430,71 @@ TrainConfig parse_text_train_args(int argc, char *argv[]) else if (arg == "--weight-decay" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --weight-decay: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --weight-decay: " << v.error().message << "\n"; + std::exit(1); + } cfg.weight_decay = *v; } else if (arg == "--beta1" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v || *v <= 0 || *v >= 1) { std::cerr << "--beta1 须在 (0, 1) 区间内\n"; std::exit(1); } + if (!v || *v <= 0 || *v >= 1) + { + std::cerr << "--beta1 须在 (0, 1) 区间内\n"; + std::exit(1); + } cfg.beta1 = *v; } else if (arg == "--beta2" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v || *v <= 0 || *v >= 1) { std::cerr << "--beta2 须在 (0, 1) 区间内\n"; std::exit(1); } + if (!v || *v <= 0 || *v >= 1) + { + std::cerr << "--beta2 须在 (0, 1) 区间内\n"; + std::exit(1); + } cfg.beta2 = *v; } else if (arg == "--d-model" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --d-model: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --d-model: " << v.error().message << "\n"; + std::exit(1); + } cfg.d_model = *v; } else if (arg == "--num-heads" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --num-heads: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --num-heads: " << v.error().message << "\n"; + std::exit(1); + } cfg.num_heads = *v; } else if (arg == "--num-layers" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --num-layers: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --num-layers: " << v.error().message << "\n"; + std::exit(1); + } cfg.num_layers = *v; } else if (arg == "--d-ff" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --d-ff: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --d-ff: " << v.error().message << "\n"; + std::exit(1); + } cfg.d_ff = *v; } else if (arg == "--model" && i + 1 < argc) @@ -379,8 +502,7 @@ TrainConfig parse_text_train_args(int argc, char *argv[]) cfg.model_type = argv[++i]; if (cfg.model_type != "gpt" && cfg.model_type != "rapt") { - std::cerr << "未知模型架构: " << cfg.model_type - << ", 可选: gpt, rapt\n"; + std::cerr << "未知模型架构: " << cfg.model_type << ", 可选: gpt, rapt\n"; if (cfg.model_type == "zipt") std::cerr << "（ZiPT 已移除，代码保留在 legacy/zipt 分支；" "恢复前提见 docs/history/history.md）\n"; @@ -390,19 +512,31 @@ TrainConfig parse_text_train_args(int argc, char *argv[]) else if (arg == "--log-interval" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --log-interval: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --log-interval: " << v.error().message << "\n"; + std::exit(1); + } cfg.log_interval = std::max(1, *v); // 0 → 1，避免步进 % 0 除零 } else if (arg == "--save-interval" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --save-interval: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --save-interval: " << v.error().message << "\n"; + std::exit(1); + } cfg.save_interval = *v; // 0 = 禁用保存（下方 % 前有 >0 保护） } else if (arg == "--max-steps" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --max-steps: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --max-steps: " << v.error().message << "\n"; + std::exit(1); + } cfg.max_steps = *v; // 0 = 不限步数 } else if (auto gpu_dev = nn::cli::parse_gpu_option(argc, argv, i)) @@ -410,8 +544,7 @@ TrainConfig parse_text_train_args(int argc, char *argv[]) cfg.gpu_enabled = true; if (!gpu_dev->empty()) cfg.gpu_device = *gpu_dev; } - else if (arg == "--grad-log") - cfg.grad_log = true; + else if (arg == "--grad-log") cfg.grad_log = true; else if (arg == "--f16") { // 快捷方式：**f16 存储**（param=F16 + compute=F16；stable/optimizer @@ -437,7 +570,7 @@ TrainConfig parse_text_train_args(int argc, char *argv[]) else if (arg == "--precision-param" && i + 1 < argc) { cfg.precision.param = parse_precision_arg(argv[++i], "precision-param", - /*allow_t1_58=*/true); + /*allow_t1_58=*/true); } else if (arg == "--precision-compute" && i + 1 < argc) { @@ -465,19 +598,31 @@ TrainConfig parse_text_train_args(int argc, char *argv[]) else if (arg == "--warmup-epochs" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --warmup-epochs: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --warmup-epochs: " << v.error().message << "\n"; + std::exit(1); + } cfg.warmup_epochs = *v; } else if (arg == "--warmup-steps" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --warmup-steps: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --warmup-steps: " << v.error().message << "\n"; + std::exit(1); + } cfg.warmup_steps = *v; } else if (arg == "--min-lr" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --min-lr: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --min-lr: " << v.error().message << "\n"; + std::exit(1); + } cfg.min_lr = *v; } else if (arg == "--lr-per-epoch" && i + 1 < argc) @@ -488,27 +633,31 @@ TrainConfig parse_text_train_args(int argc, char *argv[]) while (std::getline(ss, token, ',')) { auto v = nn::parse_number(token); - if (!v) { std::cerr << "无效 --lr-per-epoch 值: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --lr-per-epoch 值: " << v.error().message << "\n"; + std::exit(1); + } cfg.lr_per_epoch.push_back(*v); } } else if (arg == "--max-norm" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --max-norm: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --max-norm: " << v.error().message << "\n"; + std::exit(1); + } cfg.max_norm = *v; } else if (arg == "--positional-encoding" && i + 1 < argc) { std::string v = argv[++i]; - if (v == "learned") - cfg.pos_encoding = nn::PosEncodingType::Learned; - else if (v == "sinusoidal") - cfg.pos_encoding = nn::PosEncodingType::Sinusoidal; - else if (v == "alibi") - cfg.pos_encoding = nn::PosEncodingType::ALiBi; - else if (v == "rope") - cfg.pos_encoding = nn::PosEncodingType::RoPE; + if (v == "learned") cfg.pos_encoding = nn::PosEncodingType::Learned; + else if (v == "sinusoidal") cfg.pos_encoding = nn::PosEncodingType::Sinusoidal; + else if (v == "alibi") cfg.pos_encoding = nn::PosEncodingType::ALiBi; + else if (v == "rope") cfg.pos_encoding = nn::PosEncodingType::RoPE; else { std::cerr << "未知位置编码类型: " << v @@ -519,51 +668,48 @@ TrainConfig parse_text_train_args(int argc, char *argv[]) else if (arg == "--activation" && i + 1 < argc) { std::string v = argv[++i]; - if (v == "gelu") - cfg.activation = nn::ActivationType::GeLU; - else if (v == "swiglu") - cfg.activation = nn::ActivationType::SwiGLU; - else if (v == "relu2") - cfg.activation = nn::ActivationType::ReLU2; + if (v == "gelu") cfg.activation = nn::ActivationType::GeLU; + else if (v == "swiglu") cfg.activation = nn::ActivationType::SwiGLU; + else if (v == "relu2") cfg.activation = nn::ActivationType::ReLU2; else { - std::cerr << "未知激活类型: " << v - << "，可选: gelu, swiglu, relu2\n"; + std::cerr << "未知激活类型: " << v << "，可选: gelu, swiglu, relu2\n"; std::exit(1); } } else if (arg == "--norm" && i + 1 < argc) { std::string v = argv[++i]; - if (v == "layernorm") - cfg.norm_type = nn::NormType::LayerNorm; - else if (v == "rmsnorm") - cfg.norm_type = nn::NormType::RMSNorm; + if (v == "layernorm") cfg.norm_type = nn::NormType::LayerNorm; + else if (v == "rmsnorm") cfg.norm_type = nn::NormType::RMSNorm; else { - std::cerr << "未知归一化层类型: " << v - << "，可选: layernorm, rmsnorm\n"; + std::cerr << "未知归一化层类型: " << v << "，可选: layernorm, rmsnorm\n"; std::exit(1); } } else if (arg == "--flush-interval" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --flush-interval: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --flush-interval: " << v.error().message << "\n"; + std::exit(1); + } cfg.flush_interval = *v; } else if (arg == "--checkpoint-every" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --checkpoint-every: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --checkpoint-every: " << v.error().message << "\n"; + std::exit(1); + } cfg.checkpoint_every = *v; } - else if (arg == "--activation-offload") - { - cfg.activation_offload = true; - } - else if (!arg.starts_with("--")) - cfg.text_path = arg; + else if (arg == "--activation-offload") { cfg.activation_offload = true; } + else if (!arg.starts_with("--")) cfg.text_path = arg; else { std::cerr << "未知参数: " << arg << "\n使用 --help 查看用法\n"; diff --git a/include/neuralnet.cpp/cli/text_train_loop.hpp b/include/neuralnet.cpp/cli/text_train_loop.hpp index 62fd5d65..2350a321 100644 --- a/include/neuralnet.cpp/cli/text_train_loop.hpp +++ b/include/neuralnet.cpp/cli/text_train_loop.hpp @@ -50,10 +50,10 @@ namespace nn::cli // ── 断点落盘（CLI 序列化策略；何时保存由 TextTrainer 决定）───────────────── // 状态段（计数器/配方/数据集身份）由 Trainer 装配好后经 TextCheckpointRequest // 传入；本函数只做"导出优化器张量 + 写 .nnckpt + 失败记账"。 -inline void save_ckpt_with_state( - const std::string& path, nn::Model& model, const nn::ModelSpec& spec, - const std::string& tokenizer_bytes, const nn::TrainState& state, - nn::Optimizer& optimizer, TextTrainCheckpointLog& log, std::ostream& err) +inline void save_ckpt_with_state(const std::string& path, nn::Model& model, + const nn::ModelSpec& spec, const std::string& tokenizer_bytes, + const nn::TrainState& state, nn::Optimizer& optimizer, + TextTrainCheckpointLog& log, std::ostream& err) { std::vector blobs; auto sr = optimizer.save_state(blobs); @@ -76,7 +76,8 @@ inline void save_ckpt_with_state( inline bool nn_nan_trace_on() { // MSVC CRT 弃用 getenv（-Werror 下必须走 _dupenv_s），与 NN_MEM_STATS 同款。 - static const bool on = [] { + static const bool on = [] + { #if defined(_MSC_VER) char* buf = nullptr; std::size_t len = 0; @@ -108,14 +109,14 @@ inline std::vector& nn_nan_trace_tokens() return toks; } -inline bool nn_nan_trace_step(nn::ComputeEngine& engine, nn::Model& model, - std::size_t step, std::size_t seq_len) +inline bool nn_nan_trace_step(nn::ComputeEngine& engine, nn::Model& model, std::size_t step, + std::size_t seq_len) { for (int pass = 0; pass < 2; ++pass) { const char* what = (pass == 0) ? "param" : "grad"; - std::vector refs = (pass == 0) ? model.parameters() - : model.param_gradients(); + std::vector refs = + (pass == 0) ? model.parameters() : model.param_gradients(); std::size_t idx = 0; for (auto& ref : refs) { @@ -140,8 +141,7 @@ inline bool nn_nan_trace_step(nn::ComputeEngine& engine, nn::Model& model, ++bad; row_bad = true; } - else - amax = std::max(amax, std::fabs(static_cast(v))); + else amax = std::max(amax, std::fabs(static_cast(v))); } if (row_bad && bad_rows.size() < 32) bad_rows.push_back(r); } @@ -149,8 +149,7 @@ inline bool nn_nan_trace_step(nn::ComputeEngine& engine, nn::Model& model, { std::printf("[NAN-TRACE] step %zu: %s[%zu] %zux%zu 非有限 %zu/%zu，" "max|有限|=%.4g\n", - step, what, idx, m.rows(), m.cols(), bad, - m.rows() * m.cols(), amax); + step, what, idx, m.rows(), m.cols(), bad, m.rows() * m.cols(), amax); // 同一 batch 的 doc_ids 快照（首样本）：文档感知掩码是 // "逐行可见列数"的源头，命中时把序列与统计一起打出来。 const auto& ids = nn_nan_trace_doc_ids(); @@ -169,33 +168,37 @@ inline bool nn_nan_trace_step(nn::ComputeEngine& engine, nn::Model& model, min_visible = std::min(min_visible, vis); if (vis == 0) ++zero_visible; } - std::printf("[NAN-TRACE] doc_ids 统计: 零可见列的行=%zu，最小可见列=%zu，" - "递减次数=%zu，唯一 id 数=%zu\n", - zero_visible, min_visible, - [&] { std::size_t d = 0; - for (std::size_t t = 1; t < seq_len; ++t) - if (ids[t] < ids[t - 1]) ++d; - return d; }(), - [&] { auto v = ids; - std::sort(v.begin(), v.begin() + seq_len); - return static_cast( - std::unique(v.begin(), v.begin() + seq_len) - - v.begin()); }()); + std::printf( + "[NAN-TRACE] doc_ids 统计: 零可见列的行=%zu，最小可见列=%zu，" + "递减次数=%zu，唯一 id 数=%zu\n", + zero_visible, min_visible, + [&] + { + std::size_t d = 0; + for (std::size_t t = 1; t < seq_len; ++t) + if (ids[t] < ids[t - 1]) ++d; + return d; + }(), + [&] + { + auto v = ids; + std::sort(v.begin(), v.begin() + seq_len); + return static_cast( + std::unique(v.begin(), v.begin() + seq_len) - v.begin()); + }()); } // 行级定位：`grad[0]` 是 token embedding 的梯度，行号 = token id； // 把它映射回 batch 内的 (t,b) 位置，可区分"某些 id 坏"（数据/查表） // 与"某些位置坏"（掩码/位置相关算子）。 if (pass == 1 && idx == 0 && !bad_rows.empty()) { - std::printf("[NAN-TRACE] 非有限行(token id, 前 %zu): ", - bad_rows.size()); + std::printf("[NAN-TRACE] 非有限行(token id, 前 %zu): ", bad_rows.size()); for (std::size_t k = 0; k < bad_rows.size(); ++k) - std::printf("%zu%s", bad_rows[k], - k + 1 < bad_rows.size() ? "," : ""); + std::printf("%zu%s", bad_rows[k], k + 1 < bad_rows.size() ? "," : ""); std::printf("\n"); const auto& toks = nn_nan_trace_tokens(); - const std::size_t bs = (seq_len > 0 && !toks.empty()) - ? toks.size() / seq_len : 0; + const std::size_t bs = + (seq_len > 0 && !toks.empty()) ? toks.size() / seq_len : 0; if (bs > 0) { std::printf("[NAN-TRACE] 命中位置 (t,b) 前 24: "); @@ -204,8 +207,8 @@ inline bool nn_nan_trace_step(nn::ComputeEngine& engine, nn::Model& model, for (std::size_t t = 0; t < seq_len && shown < 24; ++t) { const std::size_t id = toks[b * seq_len + t]; - if (std::find(bad_rows.begin(), bad_rows.end(), id) - != bad_rows.end()) + if (std::find(bad_rows.begin(), bad_rows.end(), id) != + bad_rows.end()) { std::printf("(%zu,%zu)=%zu ", t, b, id); ++shown; @@ -227,7 +230,7 @@ inline bool nn_nan_trace_step(nn::ComputeEngine& engine, nn::Model& model, // 计算并打印全局梯度统计：L2 范数、绝对值最大值、均值。 // GPU 模式下自动通过 engine.to_matrix() 下载张量到 CPU。 inline void log_text_train_gradient_stats(nn::ComputeEngine& engine, - const std::vector& grads) + const std::vector& grads) { Scalar global_sum_sq = 0.0; Scalar global_abs_max = 0.0; @@ -247,13 +250,12 @@ inline void log_text_train_gradient_stats(nn::ComputeEngine& engine, mat = std::move(*m); } - const Scalar sum_sq = mat.reduce(Scalar{0}, std::plus<>{}, - [](Scalar x) { return x * x; }); - const Scalar abs_max = mat.reduce(Scalar{0}, - [](Scalar a, Scalar b) { return std::max(a, b); }, - [](Scalar x) { return std::abs(x); }); - const Scalar abs_sum = mat.reduce(Scalar{0}, std::plus<>{}, + const Scalar sum_sq = mat.reduce(Scalar{0}, std::plus<>{}, [](Scalar x) { return x * x; }); + const Scalar abs_max = mat.reduce( + Scalar{0}, [](Scalar a, Scalar b) { return std::max(a, b); }, [](Scalar x) { return std::abs(x); }); + const Scalar abs_sum = + mat.reduce(Scalar{0}, std::plus<>{}, [](Scalar x) { return std::abs(x); }); const std::size_t n = mat.size(); global_sum_sq += sum_sq; @@ -266,9 +268,8 @@ inline void log_text_train_gradient_stats(nn::ComputeEngine& engine, { std::cout << " [" << tensor_idx << "] " << "(" << mat.rows() << "x" << mat.cols() << ") " - << "norm=" << std::scientific << std::setprecision(4) - << std::sqrt(sum_sq) << " max|g|=" << abs_max - << " mean|g|=" << (n > 0 ? abs_sum / n : Scalar{0}) + << "norm=" << std::scientific << std::setprecision(4) << std::sqrt(sum_sq) + << " max|g|=" << abs_max << " mean|g|=" << (n > 0 ? abs_sum / n : Scalar{0}) << std::endl; } ++tensor_idx; @@ -277,11 +278,8 @@ inline void log_text_train_gradient_stats(nn::ComputeEngine& engine, const Scalar global_norm = std::sqrt(global_sum_sq); const Scalar global_mean = global_count > 0 ? global_abs_sum / global_count : Scalar{0}; std::cout << " >> grad_norm=" << std::scientific << std::setprecision(4) << global_norm - << " max|g|=" << global_abs_max - << " mean|g|=" << global_mean - << " params=" << global_count - << " tensors=" << grads.size() - << std::endl; + << " max|g|=" << global_abs_max << " mean|g|=" << global_mean + << " params=" << global_count << " tensors=" << grads.size() << std::endl; } // ── 显存阶段采样（NN_MEM_STATS=1；探针门关时恒 false 且不读环境变量）───────── @@ -294,8 +292,8 @@ struct TextTrainMemProbe void mark(nn::ComputeEngine& engine, const char* tag) const { if (!on) return; - std::cout << "[mem] " << std::left << std::setw(16) << tag << std::right - << " " << engine.pool_stats() << std::endl; + std::cout << "[mem] " << std::left << std::setw(16) << tag << std::right << " " + << engine.pool_stats() << std::endl; } }; @@ -322,17 +320,17 @@ struct TextTrainMemProbe // ══════════════════════════════════════════════════════════════════════════ // 训练入口适配：给 TextTrainer 装 CLI 侧 hooks（诊断/断点），然后跑 fit() // ══════════════════════════════════════════════════════════════════════════ -[[nodiscard]] inline nn::Result run_text_train_loop( - nn::TextTrainer& trainer, const TrainConfig& cfg, const nn::ModelSpec& spec, - const std::string& tokenizer_bytes, const TextTrainMemProbe& mem_probe) +[[nodiscard]] inline nn::Result +run_text_train_loop(nn::TextTrainer& trainer, const TrainConfig& cfg, const nn::ModelSpec& spec, + const std::string& tokenizer_bytes, const TextTrainMemProbe& mem_probe) { nn::TextTrainerHooks hooks; hooks.on_phase = [&mem_probe, &trainer](const char* tag) { mem_probe.mark(trainer.session().engine(), tag); }; // 梯度统计（--grad-log；step 后、zero_grad 前） - hooks.on_update = [&cfg, &trainer](nn::TrainingSession& session, std::size_t, - std::size_t step, Scalar) + hooks.on_update = + [&cfg, &trainer](nn::TrainingSession& session, std::size_t, std::size_t step, Scalar) { if (!cfg.grad_log) return; const std::size_t steps_per_epoch = trainer.steps_per_epoch(); @@ -343,8 +341,7 @@ struct TextTrainMemProbe #if NN_PROBES_ENABLED // NN_NAN_TRACE：batch 装配后快照 doc_ids/token（batch-major），backward 后扫描 - hooks.on_batch_ready = [](nn::TrainingSession&, const nn::TextSampleBatch& batch, - std::size_t) + hooks.on_batch_ready = [](nn::TrainingSession&, const nn::TextSampleBatch& batch, std::size_t) { if (!nn_nan_trace_on()) return; nn_nan_trace_doc_ids() = batch.doc_ids; @@ -353,14 +350,12 @@ struct TextTrainMemProbe const auto in = batch.input.span(); for (std::size_t b = 0; b < batch.batch; ++b) for (std::size_t t = 0; t < batch.seq_len; ++t) - toks[b * batch.seq_len + t] = - static_cast(in[t * batch.batch + b]); + toks[b * batch.seq_len + t] = static_cast(in[t * batch.batch + b]); }; hooks.on_probe_check = [&cfg](nn::TrainingSession& session, std::size_t step) { if (!nn_nan_trace_on()) return false; - if (nn_nan_trace_step(session.engine(), session.model(), step, cfg.seq_len)) - return false; + if (nn_nan_trace_step(session.engine(), session.model(), step, cfg.seq_len)) return false; std::printf("[NAN-TRACE] 命中即停止（NN_NAN_TRACE=1）\n"); return true; }; @@ -368,15 +363,14 @@ struct TextTrainMemProbe // 定期断点（.nnckpt）：路径派生/记账/状态段序列化都在 CLI 侧。 auto ckpt_log = std::make_shared(); - hooks.on_checkpoint = [&cfg, &spec, &tokenizer_bytes, ckpt_log]( - const nn::TextCheckpointRequest& req) -> nn::Result + hooks.on_checkpoint = [&cfg, &spec, &tokenizer_bytes, + ckpt_log](const nn::TextCheckpointRequest& req) -> nn::Result { const std::size_t before = ckpt_log->failures(); save_ckpt_with_state(text_train_ckpt_path(cfg.save_path), req.session.model(), spec, tokenizer_bytes, req.state, req.session.optimizer(), *ckpt_log, std::cerr); - if (ckpt_log->failures() != before) - NN_FAIL("checkpoint 未落盘（记账见 stderr）"); + if (ckpt_log->failures() != before) NN_FAIL("checkpoint 未落盘（记账见 stderr）"); return {}; }; diff --git a/include/neuralnet.cpp/cli/text_train_recovery.hpp b/include/neuralnet.cpp/cli/text_train_recovery.hpp index 13abce15..57e87282 100644 --- a/include/neuralnet.cpp/cli/text_train_recovery.hpp +++ b/include/neuralnet.cpp/cli/text_train_recovery.hpp @@ -63,8 +63,8 @@ struct ResumeStateAdoption bool shape_adopted = false;// 形状/架构已按文件 spec 改写（两种角色都成立） }; -[[nodiscard]] inline nn::Result adopt_resume_state( - TrainConfig& cfg, std::ostream& out, std::ostream& err) +[[nodiscard]] inline nn::Result +adopt_resume_state(TrainConfig& cfg, std::ostream& out, std::ostream& err) { ResumeStateAdoption result; if (!cfg.load_existing) return result; @@ -107,8 +107,7 @@ struct ResumeStateAdoption uint32_t version = 0; auto role_r = nn::peek_model_role(cfg.resume_path, &version); - const bool is_checkpoint = - role_r && version >= 7 && *role_r == nn::ModelRole::Checkpoint; + const bool is_checkpoint = role_r && version >= 7 && *role_r == nn::ModelRole::Checkpoint; if (role_r && version >= 7 && *role_r == nn::ModelRole::Model) { out << "从推理模型 warm start（.nnmodel）：配方与步计数由命令行给出，优化器冷启动" @@ -142,14 +141,15 @@ struct ResumeStateAdoption if (!given.empty()) given += " "; given += f; } - NN_FAIL("错误: --resume " + cfg.resume_path + - " 是训练断点（.nnckpt，带训练状态段）⇒ 形状/配方/计数器一律由断点决定，" - "不得再由命令行给出。\n 检测到显式输入: " + given + - "\n · 禁止（A/B 类：形状 / 配方 / 计数器）: " + forbidden_resume_flags_text() + - "\n · 仍允许（不改变训练）: --save / --save-interval / --log-interval / " - "--max-steps / --gpu / --grad-log / --checkpoint-every / --activation-offload / " - "--flush-interval\n" - " · 要改配方 / 换数据集 ⇒ 从推理模型（.nnmodel）开新一轮（CLI 全权、优化器冷启动）。"); + NN_FAIL( + "错误: --resume " + cfg.resume_path + + " 是训练断点（.nnckpt，带训练状态段）⇒ 形状/配方/计数器一律由断点决定，" + "不得再由命令行给出。\n 检测到显式输入: " + + given + "\n · 禁止（A/B 类：形状 / 配方 / 计数器）: " + forbidden_resume_flags_text() + + "\n · 仍允许（不改变训练）: --save / --save-interval / --log-interval / " + "--max-steps / --gpu / --grad-log / --checkpoint-every / --activation-offload / " + "--flush-interval\n" + " · 要改配方 / 换数据集 ⇒ 从推理模型（.nnmodel）开新一轮（CLI 全权、优化器冷启动）。"); } // ── 按状态段改写 cfg（A/B 类全量）──────────────────────────────────── @@ -174,20 +174,16 @@ struct ResumeStateAdoption // docs/internal/30：旧状态的 batch_size 是 micro-batch；还原成等效 global。 // 旧断点若落在累积中间，pending gradients 不在文件里，只能从最近的 // 已完成 update 边界继续。 - if (s.recipe.accum_steps == 0) - NN_FAIL("旧断点的 accum_steps=0，无法转换 batch_size 语义"); - cfg.batch_size = static_cast( - s.recipe.batch_size * s.recipe.accum_steps); - if (s.recipe.accum_steps > 0) - s.counters.step_in_epoch /= s.recipe.accum_steps; + if (s.recipe.accum_steps == 0) NN_FAIL("旧断点的 accum_steps=0，无法转换 batch_size 语义"); + cfg.batch_size = static_cast(s.recipe.batch_size * s.recipe.accum_steps); + if (s.recipe.accum_steps > 0) s.counters.step_in_epoch /= s.recipe.accum_steps; s.recipe.batch_size = static_cast(cfg.batch_size); - s.recipe.micro_batch_size = static_cast(s.recipe.batch_size / - s.recipe.accum_steps); + s.recipe.micro_batch_size = + static_cast(s.recipe.batch_size / s.recipe.accum_steps); s.recipe.batch_size_is_global = true; if (s.dataset.num_samples > 0) s.dataset.steps_per_epoch = - (s.dataset.num_samples + s.recipe.batch_size - 1) / - s.recipe.batch_size; + (s.dataset.num_samples + s.recipe.batch_size - 1) / s.recipe.batch_size; } cfg.stride = static_cast(s.stride); cfg.lr_per_epoch.clear(); @@ -208,12 +204,9 @@ struct ResumeStateAdoption << " 精度=" << nn::precision_name(s.recipe.precision.param) << "/" << nn::precision_name(s.recipe.precision.compute) << "/" << nn::precision_name(s.recipe.precision.stable) << "/" - << nn::precision_name(s.recipe.precision.optimizer) - << " 优化器=" << s.recipe.optimizer - << " global_batch=" << s.recipe.batch_size - << " micro_batch=" << s.recipe.micro_batch_size - << " accum=" << s.recipe.accum_steps - << " epochs=" << s.recipe.epochs << "\n"; + << nn::precision_name(s.recipe.precision.optimizer) << " 优化器=" << s.recipe.optimizer + << " global_batch=" << s.recipe.batch_size << " micro_batch=" << s.recipe.micro_batch_size + << " accum=" << s.recipe.accum_steps << " epochs=" << s.recipe.epochs << "\n"; return result; } @@ -221,15 +214,15 @@ struct ResumeStateAdoption // 数据集一变 ⇒ 窗口数变 ⇒ steps_per_epoch 变 ⇒ step 级调度的 total_steps、 // `global_step` 的语义、epoch 边界**同时**改变 ⇒ "同一断点 + 不同数据集" // 不存在正确的续训语义。合法出口 = 从 `.nnmodel` 开新一轮。 -[[nodiscard]] inline nn::Result verify_dataset_identity( - const nn::DatasetIdentity& saved, const nn::DatasetIdentity& current, - const std::string& resume_path) +[[nodiscard]] inline nn::Result verify_dataset_identity(const nn::DatasetIdentity& saved, + const nn::DatasetIdentity& current, + const std::string& resume_path) { const std::string why = nn::dataset_mismatch_reason(saved, current); if (why.empty()) return {}; NN_FAIL("错误: --resume 的数据集与断点不一致（内容身份变了 ⇒ 窗口网格 / step 号语义 / " - "lr 曲线形状全变，续训不存在正确语义）:\n" + why + - "\n 断点记录的路径: " + saved.path + "\n 本次使用的路径: " + current.path + + "lr 曲线形状全变，续训不存在正确语义）:\n" + + why + "\n 断点记录的路径: " + saved.path + "\n 本次使用的路径: " + current.path + "\n 断点文件: " + resume_path + "\n · 换数据集 = 新任务 ⇒ 请从推理模型（.nnmodel）重开一轮（CLI 全权、优化器冷启动）；" "\n · 仍要用这个断点 ⇒ 请提供与断点**内容完全一致**的 .nndataset。"); @@ -239,16 +232,17 @@ struct ResumeStateAdoption // ── 载入优化器状态（续训路径；建优化器之后调用）────────────────────────── // 优化器身份（种类/超参）已由状态段写回 cfg ⇒ 张量数量与形状必须对上， // 对不上由 `Optimizer::load_state` 明确报错（不静默冷启动）。 -[[nodiscard]] inline nn::Result load_optimizer_state( - nn::ComputeEngine& engine, const std::string& path, nn::Optimizer& optimizer, - std::ostream& out) +[[nodiscard]] inline nn::Result load_optimizer_state(nn::ComputeEngine& engine, + const std::string& path, + nn::Optimizer& optimizer, + std::ostream& out) { NN_TRY(loaded, nn::load_train_state(engine, path)); optimizer.set_state_step(loaded->state->optimizer_step); NN_TRY_CHECK(optimizer.load_state(loaded->optimizer_blobs)); out << "已恢复优化器状态: " << loaded->optimizer_blobs.size() << " 个张量" - << "，步计数 " << loaded->state->optimizer_step - << "（" << loaded->state->recipe.optimizer << "）\n"; + << "，步计数 " << loaded->state->optimizer_step << "（" << loaded->state->recipe.optimizer + << "）\n"; return {}; } @@ -262,16 +256,15 @@ struct TextTrainResume // ── 恢复流程：角色门禁 → 读规格 → 按规格重建模型 → 载入参数 ────────────── // out/err 语义与旧 main 一致（进度走 out，提示与警告走 err）。 -[[nodiscard]] inline nn::Result resume_text_train( - nn::ComputeEngine& engine, const TrainConfig& cfg, nn::Model& model, - nn::ModelSpec& spec, std::ostream& out, std::ostream& err) +[[nodiscard]] inline nn::Result +resume_text_train(nn::ComputeEngine& engine, const TrainConfig& cfg, nn::Model& model, + nn::ModelSpec& spec, std::ostream& out, std::ostream& err) { TextTrainResume result; auto spec_result = nn::peek_model_spec(cfg.resume_path); if (!spec_result) { - err << "加载模型失败: " << spec_result.error().message << "，将从头训练。\n" - << std::endl; + err << "加载模型失败: " << spec_result.error().message << "，将从头训练。\n" << std::endl; return result; } @@ -281,8 +274,7 @@ struct TextTrainResume // （P1-7），dense / `.nnckpt` 路径与逐位不变（意图只影响打包载荷那一条分支）。 uint32_t file_version = 0; const auto role_r = nn::peek_model_role(cfg.resume_path, &file_version); - const bool model_role = - role_r && file_version >= 7 && *role_r == nn::ModelRole::Model; + const bool model_role = role_r && file_version >= 7 && *role_r == nn::ModelRole::Model; // RAPT 独立构建路径 if (file_spec.is_rapt()) { @@ -302,37 +294,33 @@ struct TextTrainResume out << "从模型文件读取 ALiBi GPT 规格\n"; else if (fg->pos_encoding == nn::PosEncodingType::RoPE) out << "从模型文件读取 RoPE GPT 规格\n"; - else - out << "从模型文件读取 GPT 规格\n"; + else out << "从模型文件读取 GPT 规格\n"; NN_TRY_MSG(build_result, "从断点重建 GPT 模型失败", nn::build_gpt_model_from_spec(engine, file_spec, cfg.precision)); model = std::move(*build_result); spec = file_spec; result.spec_restored = true; } - else - out << "旧格式模型文件，使用命令行参数\n"; + else out << "旧格式模型文件，使用命令行参数\n"; - auto load_result = nn::load_model(cfg.resume_path, model, nullptr, nullptr, - nn::WeightLoadIntent::Trainable); + auto load_result = + nn::load_model(cfg.resume_path, model, nullptr, nullptr, nn::WeightLoadIntent::Trainable); if (!load_result) { - err << "加载模型失败: " << load_result.error().message << "，将从头训练。\n" - << std::endl; + err << "加载模型失败: " << load_result.error().message << "，将从头训练。\n" << std::endl; return result; } result.weights_loaded = true; result.vocab_bytes = load_result->size(); - const nn::WeightQuant file_wq = - file_spec.is_gpt() ? file_spec.as()->weight_quant - : file_spec.is_rapt() ? file_spec.as()->weight_quant - : nn::WeightQuant::None; + const nn::WeightQuant file_wq = file_spec.is_gpt() ? file_spec.as()->weight_quant + : file_spec.is_rapt() + ? file_spec.as()->weight_quant + : nn::WeightQuant::None; if (model_role && file_wq == nn::WeightQuant::T1_58) out << "三值（T1_58）推理模型：按文件 τ 反量化打包权重为 f32 latent ⇒ 可续训" "（首次 forward 与推理文件逐位一致）\n"; out << "已加载模型: " << cfg.resume_path; - if (!load_result->empty()) - out << " (含嵌入词表 " << load_result->size() << " 字节)"; + if (!load_result->empty()) out << " (含嵌入词表 " << load_result->size() << " 字节)"; out << "\n" << std::endl; return result; } diff --git a/include/neuralnet.cpp/compute_cpu_engine.hpp b/include/neuralnet.cpp/compute_cpu_engine.hpp index 3c528691..92fc1657 100644 --- a/include/neuralnet.cpp/compute_cpu_engine.hpp +++ b/include/neuralnet.cpp/compute_cpu_engine.hpp @@ -15,12 +15,12 @@ #include #include "compute_engine.hpp" -#include // nn::for_each / parallel_for_blocks / parallel_for_samples +#include // nn::for_each / parallel_for_blocks / parallel_for_samples #include "expr_opt.hpp" #include "expr_op_table.hpp" // 算子元数据单一真源（CPU 标量语义 / GLSL 形态，RV-07） #include "expr_registry.hpp" // scan 模式融合 kernel 登记（NN_EXPR_SCAN） -#include "expr_dsl.hpp" // matmul_with_bias 经 DSL 融合（单一事实源） -#include "ternary_pack.hpp" // 三值打包布局（trit_pack / matmul_q，P2） +#include "expr_dsl.hpp" // matmul_with_bias 经 DSL 融合（单一事实源） +#include "ternary_pack.hpp" // 三值打包布局（trit_pack / matmul_q，P2） namespace nn { @@ -44,9 +44,7 @@ class CpuEngine final : public ComputeEngine public: // 引擎必须由 shared_ptr 管理：Layer / Model / Optimizer 会保存强引用。 [[nodiscard]] static std::shared_ptr create() - { - return std::shared_ptr(new CpuEngine()); - } + { return std::shared_ptr(new CpuEngine()); } [[nodiscard]] Device device() const noexcept override { return Device::CPU; } @@ -65,8 +63,7 @@ class CpuEngine final : public ComputeEngine { #ifdef NN_EXPR_SCAN auto& owner = fused::recording_graph_owner(); - if (owner) - NN_FAIL("begin_expr: 嵌套录制（已有未结束的 begin_expr）"); + if (owner) NN_FAIL("begin_expr: 嵌套录制（已有未结束的 begin_expr）"); owner = std::make_unique(); owner->owner = weak_from_this().lock(); #endif @@ -77,8 +74,7 @@ class CpuEngine final : public ComputeEngine { #ifdef NN_EXPR_SCAN auto& owner = fused::recording_graph_owner(); - if (!owner) - return {}; + if (!owner) return {}; ExprGraph g = std::move(*owner); owner.reset(); auto kernels = fuse_expr_graph(g); @@ -101,15 +97,13 @@ class CpuEngine final : public ComputeEngine // 内部依赖输入（node>=0）此时未绑定张量 → 该视图按 f32 记； // 纯逐元素复合的成员引用经寄存器直替、不进 views，故正常构图 // 下这里的每个槽都是外部输入。 - if (slot < 32u && in.external.valid() && - in.external.precision() == Precision::F16) + if (slot < 32u && in.external.valid() && in.external.precision() == Precision::F16) in_bits |= (1u << slot); ++slot; } if (k.members.size() > 1) fused::global_registry().add_variant( - k.spec, expr_prec_sig_make(in_bits, - k.prec == Precision::F16)); + k.spec, expr_prec_sig_make(in_bits, k.prec == Precision::F16)); } #endif return {}; @@ -120,13 +114,15 @@ class CpuEngine final : public ComputeEngine // dsl::compute 的 scan_register，必须在引擎入口补记——注意力 fold、CE 的 // online-softmax fold、归约链都走这里；漏记 = 这些结构拿不到特化 shader， // 运行时退回带 `prec` 三元的 `#x`（实测注意力 kernel 慢 18%）。 - static void scan_record_variant_(const ExprSpec& spec, - std::span inputs, Precision P) + static void scan_record_variant_(const ExprSpec& spec, std::span inputs, + Precision P) { #ifdef NN_EXPR_SCAN fused::global_registry().add_variant(spec, expr_prec_sig_of(inputs, P)); #else - (void)spec; (void)inputs; (void)P; + (void)spec; + (void)inputs; + (void)P; #endif } @@ -135,10 +131,9 @@ class CpuEngine final : public ComputeEngine // RV-02：分配在宿主直配上不会失败（失败 = bad_alloc → abort），故恒成功； // 仍按 Result 形态返回以统一协议。 [[nodiscard]] Result create_tensor_impl(std::size_t rows, std::size_t cols, - Precision P) override + Precision P) override { - if (P == Precision::F16) - return Tensor::cpu(rows, cols); + if (P == Precision::F16) return Tensor::cpu(rows, cols); return Tensor::cpu(rows, cols); } @@ -151,17 +146,16 @@ class CpuEngine final : public ComputeEngine const auto src = m.span(); auto dst = m16.span(); for (std::size_t i = 0; i < src.size(); ++i) - dst[i] = src[i]; // RHE 舍入 + dst[i] = src[i]; // RHE 舍入 return Tensor::from_matrix(std::move(m16)); } - return Tensor::from_matrix(Matrix(m)); // f32 拷贝 + return Tensor::from_matrix(Matrix(m)); // f32 拷贝 } [[nodiscard]] Result to_matrix(const Tensor& t, Precision P = Precision::F32) override { NN_PROF_OP("to_matrix"); - if (!t.is_cpu()) - NN_FAIL("to_matrix: tensor is not CPU"); + if (!t.is_cpu()) NN_FAIL("to_matrix: tensor is not CPU"); if (P == Precision::F16) { // 下载为 f16 → cast 到 f32（升 cast，精确无损） @@ -202,17 +196,15 @@ class CpuEngine final : public ComputeEngine dst[i] = static_cast(src[i]); return m32; } - return Matrix(t.cpu_matrix()); // f32 拷贝 + return Matrix(t.cpu_matrix()); // f32 拷贝 } // ── cast 原语（§7.5，唯一"变精度"算子）────────────────────────────── [[nodiscard]] Result cast_impl(const Tensor& src, Precision dst) override { - if (src.precision() == dst) - return src; + if (src.precision() == dst) return src; - if (!src.is_cpu()) - NN_FAIL("cast: CPU engine only supports CPU tensors"); + if (!src.is_cpu()) NN_FAIL("cast: CPU engine only supports CPU tensors"); if (src.precision() == Precision::F16 && dst == Precision::F32) { const auto& m16 = src.cpu_matrix(); @@ -249,19 +241,20 @@ class CpuEngine final : public ComputeEngine NN_FAIL("copy_into: CPU engine only supports CPU tensors"); if (dst.rows() != src.rows() || dst.cols() != src.cols()) NN_FAIL("copy_into: shape mismatch"); - if (dst.precision() != src.precision()) - NN_FAIL("copy_into: precision mismatch"); + if (dst.precision() != src.precision()) NN_FAIL("copy_into: precision mismatch"); if (dst.precision() == Precision::F16) { const auto s = src.cpu_matrix().span(); auto d = dst.cpu_matrix().span(); - for (std::size_t i = 0; i < s.size(); ++i) d[i] = s[i]; + for (std::size_t i = 0; i < s.size(); ++i) + d[i] = s[i]; } else { const auto s = src.cpu_matrix().span(); auto d = dst.cpu_matrix().span(); - for (std::size_t i = 0; i < s.size(); ++i) d[i] = s[i]; + for (std::size_t i = 0; i < s.size(); ++i) + d[i] = s[i]; } return {}; } @@ -274,14 +267,13 @@ class CpuEngine final : public ComputeEngine NN_FAIL("cast_into: CPU engine only supports CPU tensors"); if (dst.rows() != src.rows() || dst.cols() != src.cols()) NN_FAIL("cast_into: shape mismatch"); - if (src.precision() == dst.precision()) - return copy_into(dst, src); + if (src.precision() == dst.precision()) return copy_into(dst, src); if (src.precision() == Precision::F16 && dst.precision() == Precision::F32) { const auto s = src.cpu_matrix().span(); auto d = dst.cpu_matrix().span(); for (std::size_t i = 0; i < s.size(); ++i) - d[i] = static_cast(s[i]); // 升 cast：精确无损 + d[i] = static_cast(s[i]); // 升 cast：精确无损 return {}; } if (src.precision() == Precision::F32 && dst.precision() == Precision::F16) @@ -289,7 +281,7 @@ class CpuEngine final : public ComputeEngine const auto s = src.cpu_matrix().span(); auto d = dst.cpu_matrix().span(); for (std::size_t i = 0; i < s.size(); ++i) - d[i] = s[i]; // 降 cast：round-half-to-even + d[i] = s[i]; // 降 cast：round-half-to-even return {}; } NN_FAIL("cast_into: unsupported precision conversion"); @@ -299,60 +291,56 @@ class CpuEngine final : public ComputeEngine bool /*fresh_dst*/) override { // fresh_dst 是 GPU 录制窗口排空的跳过声明，CPU 直拷无此概念。 - if (!dst.is_cpu()) - NN_FAIL("copy_from: dst tensor is not CPU"); + if (!dst.is_cpu()) NN_FAIL("copy_from: dst tensor is not CPU"); if (dst.rows() != src.rows() || dst.cols() != src.cols()) NN_FAIL("copy_from: shape mismatch"); - dst.cpu_matrix() = src; // 拷贝赋值 + dst.cpu_matrix() = src; // 拷贝赋值 return {}; } [[nodiscard]] Result clone_impl(const Tensor& src) override { - if (!src.is_cpu()) - NN_FAIL("clone: src tensor is not CPU"); + if (!src.is_cpu()) NN_FAIL("clone: src tensor is not CPU"); // CPU 原语基于 f32 Matrix → 非 f32 输入走通用边界 cast（抬 f32 拷贝 → // 落回源精度）；f32 直通。精度策略在此实现，基类入口不做能力分支。 - return move_(src, [](const Tensor& s) -> Result - { - return Tensor::from_matrix(Matrix(s.cpu_matrix())); // 深拷贝 - }); + return move_(src, + [](const Tensor& s) -> Result + { + return Tensor::from_matrix(Matrix(s.cpu_matrix())); // 深拷贝 + }); } // ── 行切片：拷贝 src 的行 [start_row, start_row + count) 到新 Tensor ── - [[nodiscard]] Result slice_rows_impl( - const Tensor& src, std::size_t start_row, std::size_t count) override + [[nodiscard]] Result slice_rows_impl(const Tensor& src, std::size_t start_row, + std::size_t count) override { - if (!src.is_cpu()) - NN_FAIL("slice_rows: src tensor is not CPU"); - return move_(src, [start_row, count](const Tensor& s) -> Result - { - const Matrix& m = s.cpu_matrix(); - if (start_row + count > m.rows()) - NN_FAIL("slice_rows: range out of bounds"); - Matrix result(count, m.cols()); - const auto src_span = m.span(); - auto dst_span = result.span(); - const std::size_t cols = m.cols(); - for (std::size_t r = 0; r < count; ++r) - std::copy_n(src_span.begin() + (start_row + r) * cols, cols, - dst_span.begin() + r * cols); - return Tensor::from_matrix(std::move(result)); - }); + if (!src.is_cpu()) NN_FAIL("slice_rows: src tensor is not CPU"); + return move_(src, + [start_row, count](const Tensor& s) -> Result + { + const Matrix& m = s.cpu_matrix(); + if (start_row + count > m.rows()) + NN_FAIL("slice_rows: range out of bounds"); + Matrix result(count, m.cols()); + const auto src_span = m.span(); + auto dst_span = result.span(); + const std::size_t cols = m.cols(); + for (std::size_t r = 0; r < count; ++r) + std::copy_n(src_span.begin() + (start_row + r) * cols, cols, + dst_span.begin() + r * cols); + return Tensor::from_matrix(std::move(result)); + }); } // ── 行插入：将 src 的所有行写入 dst 的行 [dst_start_row, ...) ── - [[nodiscard]] Result insert_rows_impl( - Tensor& dst, std::size_t dst_start_row, const Tensor& src) override + [[nodiscard]] Result insert_rows_impl(Tensor& dst, std::size_t dst_start_row, + const Tensor& src) override { - if (!dst.is_cpu() || !src.is_cpu()) - NN_FAIL("insert_rows: tensors are not CPU"); + if (!dst.is_cpu() || !src.is_cpu()) NN_FAIL("insert_rows: tensors are not CPU"); if (dst.precision() != Precision::F32 || src.precision() != Precision::F32) return insert_rows_cast_(dst, dst_start_row, src); - if (dst.cols() != src.cols()) - NN_FAIL("insert_rows: column count mismatch"); - if (dst_start_row + src.rows() > dst.rows()) - NN_FAIL("insert_rows: range out of bounds"); + if (dst.cols() != src.cols()) NN_FAIL("insert_rows: column count mismatch"); + if (dst_start_row + src.rows() > dst.rows()) NN_FAIL("insert_rows: range out of bounds"); Matrix& d = dst.cpu_matrix(); const Matrix& s = src.cpu_matrix(); const auto dst_span = d.span(); @@ -373,27 +361,23 @@ class CpuEngine final : public ComputeEngine // - num * D >= PARALLEL_THRESHOLD 时启用行块并行 // - 否则串行（小批量查表的 D 通常较小，并行调度开销不划算） // - 行块并行避免跨行数据竞争（每行独立写入） - [[nodiscard]] Result gather_rows_impl( - const Tensor& table, const Tensor& indices) override + [[nodiscard]] Result gather_rows_impl(const Tensor& table, + const Tensor& indices) override { // 输出精度 = 源精度（表）→ 非 f32 表走边界 cast 助手。 return move_(table, [this, &indices](const Tensor& t) -> Result - { - return gather_rows_f32_(t, indices); - }); + { return gather_rows_f32_(t, indices); }); } private: - [[nodiscard]] Result gather_rows_f32_( - const Tensor& table, const Tensor& indices) + [[nodiscard]] Result gather_rows_f32_(const Tensor& table, const Tensor& indices) { - if (!table.is_cpu() || !indices.is_cpu()) - NN_FAIL("gather_rows: tensors are not CPU"); + if (!table.is_cpu() || !indices.is_cpu()) NN_FAIL("gather_rows: tensors are not CPU"); const Matrix& tbl = table.cpu_matrix(); const Matrix& idx = indices.cpu_matrix(); const std::size_t vocab = tbl.rows(); const std::size_t D = tbl.cols(); - const std::size_t num = idx.size(); // 遍历 indices 所有元素（支持任意形状） + const std::size_t num = idx.size(); // 遍历 indices 所有元素（支持任意形状） Matrix result(num, D); auto dst_span = result.span(); @@ -406,18 +390,19 @@ class CpuEngine final : public ComputeEngine // 行块并行：每行独立查表，无数据竞争 auto row_indices = std::views::iota(std::size_t{0}, num); nn::for_each(row_indices.begin(), row_indices.end(), - [dst_span, tbl_span, idx_span, vocab, D](std::size_t i) noexcept { - const auto row_idx = static_cast(idx_span[i]); - if (row_idx < vocab) - { - std::copy_n(tbl_span.begin() + row_idx * D, D, - dst_span.begin() + i * D); - } - else - { - std::fill_n(dst_span.begin() + i * D, D, Scalar{0}); - } - }); + [dst_span, tbl_span, idx_span, vocab, D](std::size_t i) noexcept + { + const auto row_idx = static_cast(idx_span[i]); + if (row_idx < vocab) + { + std::copy_n(tbl_span.begin() + row_idx * D, D, + dst_span.begin() + i * D); + } + else + { + std::fill_n(dst_span.begin() + i * D, D, Scalar{0}); + } + }); } else { @@ -427,8 +412,7 @@ class CpuEngine final : public ComputeEngine const auto row_idx = static_cast(idx_span[i]); if (row_idx < vocab) { - std::copy_n(tbl_span.begin() + row_idx * D, D, - dst_span.begin() + i * D); + std::copy_n(tbl_span.begin() + row_idx * D, D, dst_span.begin() + i * D); } else { @@ -443,8 +427,8 @@ class CpuEngine final : public ComputeEngine // ── scatter_add_rows: 按 indices 把 grad 的行原子累加到 dst ── // dst: (vocab, D) 原地修改, indices: 任意形状（按 flat 遍历）, grad: (num_indices, D) // 语义: dst[flat_indices[i]] += grad[i] - [[nodiscard]] Result scatter_add_rows_impl( - Tensor& dst, const Tensor& indices, const Tensor& grad) override + [[nodiscard]] Result scatter_add_rows_impl(Tensor& dst, const Tensor& indices, + const Tensor& grad) override { if (dst.precision() != Precision::F32 || grad.precision() != Precision::F32) return scatter_add_rows_cast_(dst, indices, grad); @@ -452,19 +436,18 @@ class CpuEngine final : public ComputeEngine } private: - [[nodiscard]] Result scatter_add_rows_f32_( - Tensor& dst, const Tensor& indices, const Tensor& grad) + [[nodiscard]] Result scatter_add_rows_f32_(Tensor& dst, const Tensor& indices, + const Tensor& grad) { if (!dst.is_cpu() || !indices.is_cpu() || !grad.is_cpu()) NN_FAIL("scatter_add_rows: tensors are not CPU"); - if (dst.cols() != grad.cols()) - NN_FAIL("scatter_add_rows: column count mismatch"); + if (dst.cols() != grad.cols()) NN_FAIL("scatter_add_rows: column count mismatch"); Matrix& d = dst.cpu_matrix(); const Matrix& idx = indices.cpu_matrix(); const Matrix& g = grad.cpu_matrix(); const std::size_t vocab = d.rows(); const std::size_t D = d.cols(); - const std::size_t num = idx.size(); // 遍历 indices 所有元素（支持任意形状） + const std::size_t num = idx.size(); // 遍历 indices 所有元素（支持任意形状） auto dst_span = d.span(); const auto idx_span = idx.span(); @@ -494,22 +477,17 @@ class CpuEngine final : public ComputeEngine // - total >= PARALLEL_THRESHOLD 时启用块级并行 // - 每块仅一次 std::copy_n(N)，cache 友好 // - 典型场景：MHA 中 d_model=128, batch=32, seq=256 → 1M 元素，明显受益 - [[nodiscard]] Result rearrange_3d_impl( - const Tensor& x, std::size_t M, std::size_t B, std::size_t N, - bool inverse) override + [[nodiscard]] Result rearrange_3d_impl(const Tensor& x, std::size_t M, std::size_t B, + std::size_t N, bool inverse) override { // 非 f32：边界 cast（抬 f32 → 重入本函数 → 落回源精度） if (x.precision() != Precision::F32) return move_(x, [this, M, B, N, inverse](const Tensor& t) -> Result - { - return rearrange_3d_impl(t, M, B, N, inverse); - }); - if (x.is_gpu()) - NN_FAIL("CpuEngine: GPU tensor on CPU engine"); + { return rearrange_3d_impl(t, M, B, N, inverse); }); + if (x.is_gpu()) NN_FAIL("CpuEngine: GPU tensor on CPU engine"); const Matrix& m = x.cpu_matrix(); const std::size_t total = M * B * N; - if (m.size() != total) - NN_FAIL("rearrange_3d: element count mismatch"); + if (m.size() != total) NN_FAIL("rearrange_3d: element count mismatch"); Matrix out(inverse ? M : (B * M), inverse ? (B * N) : N); const auto src = m.span(); const auto dst = out.span(); @@ -517,7 +495,8 @@ class CpuEngine final : public ComputeEngine // (b, m) 组合索引：bm_idx = b * M + m，总块数 = B * M // 每块独立处理 N 个元素 const std::size_t n_blocks = B * M; - auto block_kernel = [src, dst, M, B, N, inverse](std::size_t bm_idx) noexcept { + auto block_kernel = [src, dst, M, B, N, inverse](std::size_t bm_idx) noexcept + { const std::size_t b = bm_idx / M; const std::size_t mi = bm_idx % M; if (!inverse) @@ -525,24 +504,21 @@ class CpuEngine final : public ComputeEngine // out[b*M + m, n] = in[m, b*N + n] const std::size_t dst_off = (b * M + mi) * N; const std::size_t src_off = mi * (B * N) + b * N; - std::copy_n(src.begin() + src_off, N, - dst.begin() + dst_off); + std::copy_n(src.begin() + src_off, N, dst.begin() + dst_off); } else { // out[m, b*N + n] = in[b*M + m, n] const std::size_t src_off = (b * M + mi) * N; const std::size_t dst_off = mi * (B * N) + b * N; - std::copy_n(src.begin() + src_off, N, - dst.begin() + dst_off); + std::copy_n(src.begin() + src_off, N, dst.begin() + dst_off); } }; if (total >= PARALLEL_THRESHOLD && n_blocks > 1) { auto block_indices = std::views::iota(std::size_t{0}, n_blocks); - nn::parallel_for_blocks(block_indices.begin(), block_indices.end(), - block_kernel); + nn::parallel_for_blocks(block_indices.begin(), block_indices.end(), block_kernel); } else { @@ -557,36 +533,29 @@ class CpuEngine final : public ComputeEngine [[nodiscard]] Result transpose_impl(const Tensor& A) override { if (A.precision() != Precision::F32) - return move_(A, [this](const Tensor& t) -> Result - { - return transpose_impl(t); - }); - if (A.is_gpu()) - NN_FAIL("CpuEngine: GPU tensor on CPU engine"); + return move_(A, + [this](const Tensor& t) -> Result { return transpose_impl(t); }); + if (A.is_gpu()) NN_FAIL("CpuEngine: GPU tensor on CPU engine"); return Tensor::from_matrix(A.cpu_matrix().transpose()); } // ── 卷积/池化窗口展开（详见 compute_engine.hpp 契约）─────────────── // 并行策略：按输出行（im2col 的 r / col2im 的输入像素行）切块，块间无竞争； // 两个方向的写入都是**连续行**，读取按窗口散布。 - [[nodiscard]] Result im2col_impl( - const Tensor& x, - std::size_t C, std::size_t H, std::size_t W, - std::size_t k, std::size_t stride, std::size_t pad, - std::size_t OH, std::size_t OW) override + [[nodiscard]] Result im2col_impl(const Tensor& x, std::size_t C, std::size_t H, + std::size_t W, std::size_t k, std::size_t stride, + std::size_t pad, std::size_t OH, + std::size_t OW) override { if (x.precision() != Precision::F32) - return move_(x, [this, C, H, W, k, stride, pad, OH, OW](const Tensor& t) -> Result - { - return im2col_impl(t, C, H, W, k, stride, pad, OH, OW); - }); - if (x.is_gpu()) - NN_FAIL("CpuEngine: GPU tensor on CPU engine"); + return move_(x, + [this, C, H, W, k, stride, pad, OH, OW](const Tensor& t) -> Result + { return im2col_impl(t, C, H, W, k, stride, pad, OH, OW); }); + if (x.is_gpu()) NN_FAIL("CpuEngine: GPU tensor on CPU engine"); if (C == 0 || H == 0 || W == 0 || k == 0 || stride == 0 || OH == 0 || OW == 0) NN_FAIL("im2col: C/H/W/k/stride/OH/OW must be > 0"); const std::size_t B = x.cols(); - if (x.rows() != C * H * W) - NN_FAIL("im2col: x must be (C*H*W, B)"); + if (x.rows() != C * H * W) NN_FAIL("im2col: x must be (C*H*W, B)"); const std::size_t kk = k * k; const std::size_t P = OH * OW; const std::size_t rows = C * kk; @@ -597,7 +566,8 @@ class CpuEngine final : public ComputeEngine const auto dst = out.span(); // 每行 r=(ci,kh,kw) 独立；列序为位置优先 (p, b)，内层对 b 连续 - auto row_kernel = [&, src, dst](std::size_t r) noexcept { + auto row_kernel = [&, src, dst](std::size_t r) noexcept + { const std::size_t ci = r / kk; const std::size_t kh = (r / k) % k; const std::size_t kw = r % k; @@ -605,21 +575,19 @@ class CpuEngine final : public ComputeEngine Scalar* outrow = dst.data() + r * cols; for (std::size_t oh = 0; oh < OH; ++oh) { - const long ih = static_cast(oh * stride + kh) - - static_cast(pad); + const long ih = static_cast(oh * stride + kh) - static_cast(pad); const bool h_ok = (ih >= 0 && ih < static_cast(H)); for (std::size_t ow = 0; ow < OW; ++ow) { - const long iw = static_cast(ow * stride + kw) - - static_cast(pad); + const long iw = static_cast(ow * stride + kw) - static_cast(pad); const bool ok = h_ok && iw >= 0 && iw < static_cast(W); Scalar* outp = outrow + (oh * OW + ow) * B; if (ok) { // 源 (row, b) 行主序 → row*B + b - const Scalar* inp = src.data() + - (base + static_cast(ih) * W + - static_cast(iw)) * B; + const Scalar* inp = src.data() + (base + static_cast(ih) * W + + static_cast(iw)) * + B; std::copy_n(inp, B, outp); } else @@ -637,24 +605,22 @@ class CpuEngine final : public ComputeEngine } else { - for (std::size_t r = 0; r < rows; ++r) row_kernel(r); + for (std::size_t r = 0; r < rows; ++r) + row_kernel(r); } return Tensor::from_matrix(std::move(out)); } - [[nodiscard]] Result col2im_impl( - const Tensor& col, - std::size_t C, std::size_t H, std::size_t W, - std::size_t k, std::size_t stride, std::size_t pad, - std::size_t OH, std::size_t OW) override + [[nodiscard]] Result col2im_impl(const Tensor& col, std::size_t C, std::size_t H, + std::size_t W, std::size_t k, std::size_t stride, + std::size_t pad, std::size_t OH, + std::size_t OW) override { if (col.precision() != Precision::F32) - return move_(col, [this, C, H, W, k, stride, pad, OH, OW](const Tensor& t) -> Result - { - return col2im_impl(t, C, H, W, k, stride, pad, OH, OW); - }); - if (col.is_gpu()) - NN_FAIL("CpuEngine: GPU tensor on CPU engine"); + return move_(col, + [this, C, H, W, k, stride, pad, OH, OW](const Tensor& t) -> Result + { return col2im_impl(t, C, H, W, k, stride, pad, OH, OW); }); + if (col.is_gpu()) NN_FAIL("CpuEngine: GPU tensor on CPU engine"); if (C == 0 || H == 0 || W == 0 || k == 0 || stride == 0 || OH == 0 || OW == 0) NN_FAIL("col2im: C/H/W/k/stride/OH/OW must be > 0"); const std::size_t kk = k * k; @@ -670,7 +636,8 @@ class CpuEngine final : public ComputeEngine // 每个输入像素 (ci, ih, iw) 独立： // 对每个 tap (kh,kw)，至多存在一个窗口 (oh,ow) 覆盖该像素；整除性与 // 范围检查通过则累加 col 对应行 r=(ci*k+kh)*k+kw 的 (oh,ow) 位置（内层 b 连续）。 - auto pixel_kernel = [&, src, dst](std::size_t row) noexcept { + auto pixel_kernel = [&, src, dst](std::size_t row) noexcept + { const std::size_t ci = row / (H * W); const std::size_t rem = row % (H * W); const std::size_t ih = rem / W; @@ -690,10 +657,11 @@ class CpuEngine final : public ComputeEngine const long ow = s / static_cast(stride); if (ow >= static_cast(OW)) continue; const std::size_t r = (ci * k + kh) * k + kw; - const Scalar* inp = src.data() + r * (B * P) + - (static_cast(oh) * OW + - static_cast(ow)) * B; - for (std::size_t b = 0; b < B; ++b) outp[b] += inp[b]; + const Scalar* inp = + src.data() + r * (B * P) + + (static_cast(oh) * OW + static_cast(ow)) * B; + for (std::size_t b = 0; b < B; ++b) + outp[b] += inp[b]; } } }; @@ -706,7 +674,8 @@ class CpuEngine final : public ComputeEngine } else { - for (std::size_t r = 0; r < in_rows; ++r) pixel_kernel(r); + for (std::size_t r = 0; r < in_rows; ++r) + pixel_kernel(r); } return Tensor::from_matrix(std::move(out)); } @@ -737,11 +706,9 @@ class CpuEngine final : public ComputeEngine // 解码只在"逐存储行"上做一次（一行 16 个权重共享一个 32-bit 字，解码一遍 // 进小缓冲，之后按标量 FMA 复用）—— 既不物化 4 B/权重的 f32 副本，也不 // 让内层循环反复移位。 - [[nodiscard]] Result trit_pack_impl(const Tensor& w, - const Tensor& scale) override + [[nodiscard]] Result trit_pack_impl(const Tensor& w, const Tensor& scale) override { - if (!w.is_cpu() || !scale.is_cpu()) - NN_FAIL("trit_pack: tensors are not CPU"); + if (!w.is_cpu() || !scale.is_cpu()) NN_FAIL("trit_pack: tensors are not CPU"); if (w.precision() != Precision::F32 || scale.precision() != Precision::F32) NN_FAIL("trit_pack: 需要 f32 权重 + f32 尺度"); if (scale.rows() != w.rows() || scale.cols() != 1) @@ -767,9 +734,8 @@ class CpuEngine final : public ComputeEngine return out; } - [[nodiscard]] Result matmul_q_impl(const Tensor& packed, - const Tensor& scale, const Tensor& B, - bool transA, std::size_t out_rows, + [[nodiscard]] Result matmul_q_impl(const Tensor& packed, const Tensor& scale, + const Tensor& B, bool transA, std::size_t out_rows, Precision P) override { if (!packed.is_cpu() || !scale.is_cpu() || !B.is_cpu()) @@ -780,11 +746,10 @@ class CpuEngine final : public ComputeEngine { // 注：to_f32/to_prec 是基类**私有**助手（只有基类入口可用）→ 这里 // 用公共 cast（与 matmul_impl 的 CPU 边界 cast 同款）。 - auto b32 = (B.precision() == Precision::F32) ? Result{B} - : cast(B, Precision::F32); + auto b32 = + (B.precision() == Precision::F32) ? Result{B} : cast(B, Precision::F32); NN_TRY_CHECK(b32); - NN_TRY(r, matmul_q_impl(packed, scale, *b32, transA, out_rows, - Precision::F32)); + NN_TRY(r, matmul_q_impl(packed, scale, *b32, transA, out_rows, Precision::F32)); if (P == Precision::F32) return r; // 递归产物是 `*_impl` 直构；v3 起设备域从存储派生，进公共 cast // （内含 bind_check_）无需补任何身份。 @@ -805,8 +770,7 @@ class CpuEngine final : public ComputeEngine if (out_rows == 0) NN_FAIL("matmul_q_t: out_rows 必须显式给出"); if (B.rows() != K) NN_FAIL("matmul_q_t: B 行数 != 打包 A 的行数 K"); } - else if (B.rows() != K) - NN_FAIL("matmul_q: B 行数 != A 的列数 K"); + else if (B.rows() != K) NN_FAIL("matmul_q: B 行数 != A 的列数 K"); if (nn::trit_pack::words_per_row(a_len) != packed.cols()) NN_FAIL("matmul_q: 打包容器列数与 (K/M) 不匹配"); if (scale.rows() != packed.rows() || scale.cols() != 1) @@ -829,7 +793,7 @@ class CpuEngine final : public ComputeEngine // （与稠密路径读一遍权重同量级），故正确做法是把它与乘加**分开**： // 解码一个行块到 f32 scratch，再交给代数层的分块微内核 // （`multiply_to` / `transpose_multiply_to`：SIMD + 寄存器分块）。 - constexpr std::size_t RB = 64; // 行块（scratch = RB × a_len × 4B） + constexpr std::size_t RB = 64; // 行块（scratch = RB × a_len × 4B） if (!transA) { // C[m,n] = τ_m · Σ_k trit(m,k)·B[k,n] @@ -845,11 +809,10 @@ class CpuEngine final : public ComputeEngine float* dst = ab.data() + i * a_len; for (std::size_t j = 0; j < a_len; ++j) { - const std::uint32_t word = std::bit_cast( - pw[j / nn::trit_pack::kPerWord]); + const std::uint32_t word = + std::bit_cast(pw[j / nn::trit_pack::kPerWord]); const std::uint32_t code = - (word >> (nn::trit_pack::kBits * - (j % nn::trit_pack::kPerWord))) & + (word >> (nn::trit_pack::kBits * (j % nn::trit_pack::kPerWord))) & nn::trit_pack::kMask; dst[j] = tau * nn::trit_pack::value_of(code); } @@ -877,11 +840,10 @@ class CpuEngine final : public ComputeEngine for (std::size_t i = 0; i < nr; ++i) { const std::size_t m = m0 + i; - const std::uint32_t word = std::bit_cast( - pw[m / nn::trit_pack::kPerWord]); + const std::uint32_t word = + std::bit_cast(pw[m / nn::trit_pack::kPerWord]); const std::uint32_t code = - (word >> (nn::trit_pack::kBits * - (m % nn::trit_pack::kPerWord))) & + (word >> (nn::trit_pack::kBits * (m % nn::trit_pack::kPerWord))) & nn::trit_pack::kMask; dst[i] = tau * nn::trit_pack::value_of(code); } @@ -895,21 +857,17 @@ class CpuEngine final : public ComputeEngine return out; } - [[nodiscard]] Result matmul_impl( - const Tensor& A, const Tensor& B, - bool transA, bool transB, - Precision P = Precision::F32) override + [[nodiscard]] Result matmul_impl(const Tensor& A, const Tensor& B, bool transA, + bool transB, Precision P = Precision::F32) override { - if (A.is_gpu() || B.is_gpu()) - NN_FAIL("CpuEngine: GPU tensor on CPU engine"); + if (A.is_gpu() || B.is_gpu()) NN_FAIL("CpuEngine: GPU tensor on CPU engine"); // ── 精度与存储不匹配时不可直读（cpu_matrix

对错位存储返回空 → UB； // 实测 CPU f16 训练 step0 即 NaN）→ 先统一到 f32 空间计算，输出按 P // 舍入（§7.2 f32 参考语义）。两端同 f16 才走下方原生 f16 GEMM。 { const bool a16 = A.precision() == Precision::F16; const bool b16 = B.precision() == Precision::F16; - if ((P == Precision::F32 && (a16 || b16)) || - (P == Precision::F16 && !(a16 && b16))) + if ((P == Precision::F32 && (a16 || b16)) || (P == Precision::F16 && !(a16 && b16))) { auto a32 = a16 ? cast(A, Precision::F32) : Result{A}; NN_TRY_CHECK(a32); @@ -928,26 +886,23 @@ class CpuEngine final : public ComputeEngine const Matrix& a = A.cpu_matrix(); const Matrix& b = B.cpu_matrix(); - const std::size_t M = transA ? a.cols() : a.rows(); - const std::size_t K = transA ? a.rows() : a.cols(); + const std::size_t M = transA ? a.cols() : a.rows(); + const std::size_t K = transA ? a.rows() : a.cols(); const std::size_t K2 = transB ? b.cols() : b.rows(); - const std::size_t N = transB ? b.rows() : b.cols(); + const std::size_t N = transB ? b.rows() : b.cols(); if (K != K2) - NN_FAIL("matmul: dimension mismatch A=" + - std::to_string(a.rows()) + "x" + std::to_string(a.cols()) + - " transA=" + (transA ? "1" : "0") + - " B=" + std::to_string(b.rows()) + "x" + std::to_string(b.cols()) + - " transB=" + (transB ? "1" : "0") + - " K=" + std::to_string(K) + " K2=" + std::to_string(K2)); + NN_FAIL("matmul: dimension mismatch A=" + std::to_string(a.rows()) + "x" + + std::to_string(a.cols()) + " transA=" + (transA ? "1" : "0") + + " B=" + std::to_string(b.rows()) + "x" + std::to_string(b.cols()) + + " transB=" + (transB ? "1" : "0") + " K=" + std::to_string(K) + + " K2=" + std::to_string(K2)); // 构造函数里的全尺寸零填充是**重复**的（本机单线程写满一遍 6MB ~2ms） Matrix result = Matrix::make_uninitialized(M, N); - if (!transA && !transB) { - a.multiply_to(result, b); - } else if (!transA && transB) { - a.multiply_transposed_to(result, b); - } else if (transA && !transB) { - a.transpose_multiply_to(result, b); - } else { + if (!transA && !transB) { a.multiply_to(result, b); } + else if (!transA && transB) { a.multiply_transposed_to(result, b); } + else if (transA && !transB) { a.transpose_multiply_to(result, b); } + else + { Matrix a_t = a.transpose(); a_t.multiply_transposed_to(result, b); } @@ -961,21 +916,18 @@ class CpuEngine final : public ComputeEngine const auto& a = A.cpu_matrix(); const auto& b = B.cpu_matrix(); - const std::size_t M = transA ? a.cols() : a.rows(); - const std::size_t K = transA ? a.rows() : a.cols(); + const std::size_t M = transA ? a.cols() : a.rows(); + const std::size_t K = transA ? a.rows() : a.cols(); const std::size_t K2 = transB ? b.cols() : b.rows(); - const std::size_t N = transB ? b.rows() : b.cols(); - if (K != K2) - NN_FAIL("matmul(f16): dimension mismatch"); + const std::size_t N = transB ? b.rows() : b.cols(); + if (K != K2) NN_FAIL("matmul(f16): dimension mismatch"); MatrixT result(M, N); // f16 无原生 x86 SIMD → 向量化失败是预期行为（§7.2 兼容路径） - if (!transA && !transB) { - a.multiply_to(result, b); - } else if (!transA && transB) { - a.multiply_transposed_to(result, b); - } else if (transA && !transB) { - a.transpose_multiply_to(result, b); - } else { + if (!transA && !transB) { a.multiply_to(result, b); } + else if (!transA && transB) { a.multiply_transposed_to(result, b); } + else if (transA && !transB) { a.transpose_multiply_to(result, b); } + else + { auto a_t = a.transpose(); a_t.multiply_transposed_to(result, b); } @@ -989,52 +941,44 @@ class CpuEngine final : public ComputeEngine // 覆盖基类默认的"matmul + host 端 to_matrix 往返加 bias"：CPU 走模板求值， // GPU 走 AOT 融合 shader；scan_exprs dry-run 经本路径收集该结构 → gen_fused // 生成融合 kernel（结构没被收集 = GPU 闭合世界里没有对应 key）。 - [[nodiscard]] Result matmul_with_bias_impl( - const Tensor& A, const Tensor& B, const Tensor& bias, - bool transA = false, bool transB = false, - Precision P = Precision::F32) override + [[nodiscard]] Result matmul_with_bias_impl(const Tensor& A, const Tensor& B, + const Tensor& bias, bool transA = false, + bool transB = false, + Precision P = Precision::F32) override { - if (A.is_gpu() || B.is_gpu()) - NN_FAIL("CpuEngine: GPU tensor on CPU engine"); + if (A.is_gpu() || B.is_gpu()) NN_FAIL("CpuEngine: GPU tensor on CPU engine"); const std::size_t rows = transA ? A.cols() : A.rows(); const std::size_t cols = transB ? B.rows() : B.cols(); // P 必须下传（不可丢弃）：scan 的 f16 dry-run 靠 dsl::compute 的 // NN_EXPR_SCAN 钩子按**真实操作数精度 + P** 登记 (结构, 签名)——丢掉 P // 就会只登记全 f32 签名，GPU 侧带类型 matmul 段变体永远发现不到 // （Linear::forward 每个 Linear 一次，是 cast 临时量最大的单一来源）。 - return nn::dsl::compute(*this, - nn::dsl::matmul(A, B, transA, transB, 1) - + nn::dsl::row_broadcast(bias), - rows, cols, P); + return nn::dsl::compute( + *this, nn::dsl::matmul(A, B, transA, transB, 1) + nn::dsl::row_broadcast(bias), rows, + cols, P); } // ── 批量矩阵乘法：按 batch 切分行块，逐 batch 矩阵乘 ── // C_b = alpha * op(A_b, B_b)（alpha 为 cuBLAS sgemm 语义的输出缩放系数） // P: 计算精度（§8.1，同 matmul） - [[nodiscard]] Result batched_matmul_impl( - const Tensor& A, const Tensor& B, - std::size_t batch, - bool transA, bool transB, - Scalar alpha, - Precision P = Precision::F32) override + [[nodiscard]] Result batched_matmul_impl(const Tensor& A, const Tensor& B, + std::size_t batch, bool transA, bool transB, + Scalar alpha, + Precision P = Precision::F32) override { - if (A.is_gpu() || B.is_gpu()) - NN_FAIL("CpuEngine: GPU tensor on CPU engine"); - if (batch == 0) - NN_FAIL("batched_matmul: batch must be > 0"); + if (A.is_gpu() || B.is_gpu()) NN_FAIL("CpuEngine: GPU tensor on CPU engine"); + if (batch == 0) NN_FAIL("batched_matmul: batch must be > 0"); // ── 精度与存储不匹配：统一到 f32 空间计算 + 按 P 舍入（同 matmul）── { const bool a16 = A.precision() == Precision::F16; const bool b16 = B.precision() == Precision::F16; - if ((P == Precision::F32 && (a16 || b16)) || - (P == Precision::F16 && !(a16 && b16))) + if ((P == Precision::F32 && (a16 || b16)) || (P == Precision::F16 && !(a16 && b16))) { auto a32 = a16 ? cast(A, Precision::F32) : Result{A}; NN_TRY_CHECK(a32); auto b32 = b16 ? cast(B, Precision::F32) : Result{B}; NN_TRY_CHECK(b32); - auto r = batched_matmul(*a32, *b32, batch, transA, transB, alpha, - Precision::F32); + auto r = batched_matmul(*a32, *b32, batch, transA, transB, alpha, Precision::F32); NN_TRY_CHECK(r); if (P == Precision::F32) return r; return cast(*r, Precision::F16); @@ -1066,20 +1010,17 @@ class CpuEngine final : public ComputeEngine MatrixT c_view(M, N); // f16 无原生 x86 SIMD → 向量化失败是预期行为（§7.2 兼容路径） - if (!transA && !transB) - a_view.multiply_to(c_view, b_view); - else if (!transA && transB) - a_view.multiply_transposed_to(c_view, b_view); - else if (transA && !transB) - a_view.transpose_multiply_to(c_view, b_view); - else { + if (!transA && !transB) a_view.multiply_to(c_view, b_view); + else if (!transA && transB) a_view.multiply_transposed_to(c_view, b_view); + else if (transA && !transB) a_view.transpose_multiply_to(c_view, b_view); + else + { auto a_t = a_view.transpose(); a_t.multiply_transposed_to(c_view, b_view); } // 写入结果的对应 batch 行块 - std::copy_n(c_view.span().begin(), M * N, - result.span().data() + bi * M * N); + std::copy_n(c_view.span().begin(), M * N, result.span().data() + bi * M * N); } // alpha 缩放 @@ -1087,7 +1028,7 @@ class CpuEngine final : public ComputeEngine { auto dst = result.span(); for (std::size_t i = 0; i < dst.size(); ++i) - dst[i] = dst[i] * alpha; // f16 * float → f16 + dst[i] = dst[i] * alpha; // f16 * float → f16 } return Tensor::from_matrix(std::move(result)); } @@ -1107,12 +1048,11 @@ class CpuEngine final : public ComputeEngine // transA=true: A_b 视为 (a.cols(), a_rows_per)（按 A_b^T 使用），M=a.cols(), K=a_rows_per // transB=false: B_b 视为 (b_rows_per, b.cols())，K2=b_rows_per, N=b.cols() // transB=true: B_b 视为 (b.cols(), b_rows_per)（按 B_b^T 使用），K2=b.cols(), N=b_rows_per - const std::size_t M = transA ? a.cols() : a_rows_per; - const std::size_t K = transA ? a_rows_per : a.cols(); + const std::size_t M = transA ? a.cols() : a_rows_per; + const std::size_t K = transA ? a_rows_per : a.cols(); const std::size_t K2 = transB ? b.cols() : b_rows_per; - const std::size_t N = transB ? b_rows_per : b.cols(); - if (K != K2) - NN_FAIL("batched_matmul: K dimension mismatch"); + const std::size_t N = transB ? b_rows_per : b.cols(); + if (K != K2) NN_FAIL("batched_matmul: K dimension mismatch"); Matrix result(batch * M, N); auto result_span = result.span(); const auto a_span = a.span(); @@ -1127,19 +1067,23 @@ class CpuEngine final : public ComputeEngine auto b_sub = std::span(b_span.data() + b_off, b_rows_per * b.cols()); auto c_sub = std::span(result_span.data() + bi * M * N, M * N); - if (!transA && !transB) { - Matrix::multiply_to_span(c_sub, M, N, - a_sub, a_rows_per, a.cols(), - b_sub, b_rows_per, b.cols()); - } else if (!transA && transB) { - Matrix::multiply_transposed_to_span(c_sub, M, N, - a_sub, a_rows_per, a.cols(), - b_sub, b_rows_per, b.cols()); - } else if (transA && !transB) { - Matrix::transpose_multiply_to_span(c_sub, M, N, - a_sub, a_rows_per, a.cols(), - b_sub, b_rows_per, b.cols()); - } else { + if (!transA && !transB) + { + Matrix::multiply_to_span(c_sub, M, N, a_sub, a_rows_per, a.cols(), b_sub, + b_rows_per, b.cols()); + } + else if (!transA && transB) + { + Matrix::multiply_transposed_to_span(c_sub, M, N, a_sub, a_rows_per, a.cols(), b_sub, + b_rows_per, b.cols()); + } + else if (transA && !transB) + { + Matrix::transpose_multiply_to_span(c_sub, M, N, a_sub, a_rows_per, a.cols(), b_sub, + b_rows_per, b.cols()); + } + else + { // 双转置：先转置 A 子块，再做 A^T × B^T // 等价于 (B × A)^T，用 span 暂不支持，fallback 到临时矩阵 Matrix a_view(a_rows_per, a.cols()); @@ -1170,11 +1114,8 @@ class CpuEngine final : public ComputeEngine { if (A.precision() != Precision::F32 || B.precision() != Precision::F32) return inplace2_(A, B, [this](Tensor& a, const Tensor& b) -> Result - { - return add_inplace_impl(a, b); - }); - if (A.rows() != B.rows() || A.cols() != B.cols()) - NN_FAIL("add_inplace: shape mismatch"); + { return add_inplace_impl(a, b); }); + if (A.rows() != B.rows() || A.cols() != B.cols()) NN_FAIL("add_inplace: shape mismatch"); A.cpu_matrix().add_inplace(B.cpu_matrix()); return {}; } @@ -1183,9 +1124,7 @@ class CpuEngine final : public ComputeEngine { if (A.precision() != Precision::F32) return inplace1_(A, [this, s](Tensor& a) -> Result - { - return scale_inplace_impl(a, s); - }); + { return scale_inplace_impl(a, s); }); A.cpu_matrix().scale_inplace(s); return {}; } @@ -1193,10 +1132,7 @@ class CpuEngine final : public ComputeEngine [[nodiscard]] Result zero_impl(Tensor& A) override { if (A.precision() != Precision::F32) - return inplace1_(A, [this](Tensor& a) -> Result - { - return zero_impl(a); - }); + return inplace1_(A, [this](Tensor& a) -> Result { return zero_impl(a); }); A.cpu_matrix().zero(); return {}; } @@ -1210,44 +1146,38 @@ class CpuEngine final : public ComputeEngine // ══════════════════════════════════════════════════════════════════════ // ── 前缀扫描（forward / backward pass 1 / forward_step）────────────── - [[nodiscard]] Result scan_prefix_outer_impl( - const Tensor& K, const Tensor& V, const Tensor& P, const Tensor& R, - const Tensor& A0, const Tensor& B0, bool has_state, - std::size_t dk, std::size_t heads, bool causal, - const Tensor& boundary, bool has_bnd, - Precision prec = Precision::F32) override + [[nodiscard]] Result + scan_prefix_outer_impl(const Tensor& K, const Tensor& V, const Tensor& P, const Tensor& R, + const Tensor& A0, const Tensor& B0, bool has_state, std::size_t dk, + std::size_t heads, bool causal, const Tensor& boundary, bool has_bnd, + Precision prec = Precision::F32) override { const std::vector ts{K, V, P, R, A0, B0, boundary}; if (prec == Precision::F32 && all_f32(ts)) - return scan_prefix_outer_f32_(K, V, P, R, A0, B0, has_state, dk, heads, - causal, boundary, has_bnd); + return scan_prefix_outer_f32_(K, V, P, R, A0, B0, has_state, dk, heads, causal, + boundary, has_bnd); // 边界 cast：全部操作数抬 f32 → 算 → 输出按 prec 落回 auto c = to_f32_all(ts); NN_TRY_CHECK(c); - auto r = scan_prefix_outer_f32_((*c)[0], (*c)[1], (*c)[2], (*c)[3], - (*c)[4], (*c)[5], has_state, dk, heads, - causal, (*c)[6], has_bnd); + auto r = scan_prefix_outer_f32_((*c)[0], (*c)[1], (*c)[2], (*c)[3], (*c)[4], (*c)[5], + has_state, dk, heads, causal, (*c)[6], has_bnd); NN_TRY_CHECK(r); return to_prec(std::move(*r), prec); } private: - [[nodiscard]] Result scan_prefix_outer_f32_( - const Tensor& K, const Tensor& V, const Tensor& P, const Tensor& R, - const Tensor& A0, const Tensor& B0, bool has_state, - std::size_t dk, std::size_t heads, bool causal, - const Tensor& boundary, bool has_bnd) + [[nodiscard]] Result + scan_prefix_outer_f32_(const Tensor& K, const Tensor& V, const Tensor& P, const Tensor& R, + const Tensor& A0, const Tensor& B0, bool has_state, std::size_t dk, + std::size_t heads, bool causal, const Tensor& boundary, bool has_bnd) { for (const auto& t : {K, V, P, R, A0, B0, boundary}) - if (t.is_gpu()) - NN_FAIL("CpuEngine: GPU tensor on CPU engine"); - if (dk == 0 || heads == 0) - NN_FAIL("scan_prefix_outer: dk/heads must be > 0"); + if (t.is_gpu()) NN_FAIL("CpuEngine: GPU tensor on CPU engine"); + if (dk == 0 || heads == 0) NN_FAIL("scan_prefix_outer: dk/heads must be > 0"); const Matrix& Km = K.cpu_matrix(); const std::size_t rows = Km.rows(); - const std::size_t seq = Km.cols(); - if (rows % (dk * heads) != 0) - NN_FAIL("scan_prefix_outer: K rows not divisible by H*dk"); + const std::size_t seq = Km.cols(); + if (rows % (dk * heads) != 0) NN_FAIL("scan_prefix_outer: K rows not divisible by H*dk"); for (const auto* m : {&V.cpu_matrix(), &P.cpu_matrix(), &R.cpu_matrix()}) if (m->rows() != rows || m->cols() != seq) NN_FAIL("scan_prefix_outer: K/V/P/R shape mismatch"); @@ -1276,13 +1206,12 @@ class CpuEngine final : public ComputeEngine { const std::size_t r0 = bh * dk; const std::size_t batch = bh / heads; - const std::size_t h = bh % heads; - const auto doc_reset = [bdm, batch, seq, has_bnd](std::size_t t) { - return has_bnd && bdm.at_unchecked(0, batch * seq + t) != Scalar{0}; - }; + const std::size_t h = bh % heads; + const auto doc_reset = [bdm, batch, seq, has_bnd](std::size_t t) + { return has_bnd && bdm.at_unchecked(0, batch * seq + t) != Scalar{0}; }; - std::vector A(dk * dk, Scalar{0}), B(dk * dk, Scalar{0}), - qv(dk), kv(dk), vv(dk), num(dk), Aq(dk); + std::vector A(dk * dk, Scalar{0}), B(dk * dk, Scalar{0}), qv(dk), kv(dk), + vv(dk), num(dk), Aq(dk); if (has_state) { const Matrix& a0m = A0.cpu_matrix(); @@ -1294,7 +1223,7 @@ class CpuEngine final : public ComputeEngine B[a * dk + b2] = b0m.at_unchecked(h * dk + a, b2); } } - if (!causal) // 双向：先求全集 A, B + if (!causal) // 双向：先求全集 A, B { for (std::size_t t = 0; t < seq; ++t) { @@ -1319,7 +1248,7 @@ class CpuEngine final : public ComputeEngine kv[j] = Km.at_unchecked(r0 + j, t); vv[j] = Vm.at_unchecked(r0 + j, t); } - if (causal) // 前缀和含自身 i<=t；文档边界处重置 + if (causal) // 前缀和含自身 i<=t；文档边界处重置 { if (doc_reset(t)) { @@ -1339,11 +1268,11 @@ class CpuEngine final : public ComputeEngine Scalar acc{0}, acc2{0}; for (std::size_t j = 0; j < dk; ++j) { - acc += B[i * dk + j] * qv[j]; + acc += B[i * dk + j] * qv[j]; acc2 += A[i * dk + j] * qv[j]; } num[i] = acc; - Aq[i] = acc2; + Aq[i] = acc2; } Scalar s{0}; for (std::size_t j = 0; j < dk; ++j) @@ -1380,39 +1309,36 @@ class CpuEngine final : public ComputeEngine public: // ── 后缀扫描（backward pass 2：gK/gV）──────────────────────────────── - [[nodiscard]] Result scan_suffix_outer_impl( - const Tensor& D, const Tensor& X, const Tensor& Y, - std::size_t dk, std::size_t heads, bool causal, - const Tensor& boundary, bool has_bnd, - Precision prec = Precision::F32) override + [[nodiscard]] Result scan_suffix_outer_impl(const Tensor& D, const Tensor& X, + const Tensor& Y, std::size_t dk, + std::size_t heads, bool causal, + const Tensor& boundary, bool has_bnd, + Precision prec = Precision::F32) override { const std::vector ts{D, X, Y, boundary}; if (prec == Precision::F32 && all_f32(ts)) return scan_suffix_outer_f32_(D, X, Y, dk, heads, causal, boundary, has_bnd); auto c = to_f32_all(ts); NN_TRY_CHECK(c); - auto r = scan_suffix_outer_f32_((*c)[0], (*c)[1], (*c)[2], dk, heads, - causal, (*c)[3], has_bnd); + auto r = + scan_suffix_outer_f32_((*c)[0], (*c)[1], (*c)[2], dk, heads, causal, (*c)[3], has_bnd); NN_TRY_CHECK(r); return to_prec(std::move(*r), prec); } private: - [[nodiscard]] Result scan_suffix_outer_f32_( - const Tensor& D, const Tensor& X, const Tensor& Y, - std::size_t dk, std::size_t heads, bool causal, - const Tensor& boundary, bool has_bnd) + [[nodiscard]] Result scan_suffix_outer_f32_(const Tensor& D, const Tensor& X, + const Tensor& Y, std::size_t dk, + std::size_t heads, bool causal, + const Tensor& boundary, bool has_bnd) { for (const auto& t : {D, X, Y, boundary}) - if (t.is_gpu()) - NN_FAIL("CpuEngine: GPU tensor on CPU engine"); - if (dk == 0 || heads == 0) - NN_FAIL("scan_suffix_outer: dk/heads must be > 0"); + if (t.is_gpu()) NN_FAIL("CpuEngine: GPU tensor on CPU engine"); + if (dk == 0 || heads == 0) NN_FAIL("scan_suffix_outer: dk/heads must be > 0"); const Matrix& Xm = X.cpu_matrix(); const std::size_t rows = Xm.rows(); - const std::size_t seq = Xm.cols(); - if (rows % (dk * heads) != 0) - NN_FAIL("scan_suffix_outer: X rows not divisible by H*dk"); + const std::size_t seq = Xm.cols(); + if (rows % (dk * heads) != 0) NN_FAIL("scan_suffix_outer: X rows not divisible by H*dk"); const Matrix& Dm = D.cpu_matrix(); if (Dm.rows() != rows * dk || Dm.cols() != seq) NN_FAIL("scan_suffix_outer: D must be (B*H*dk*dk, seq)"); @@ -1436,9 +1362,8 @@ class CpuEngine final : public ComputeEngine const std::size_t r0 = bh * dk; const std::size_t dbase = bh * dk * dk; const std::size_t batch = bh / heads; - const auto doc_reset = [bdm, batch, seq, has_bnd](std::size_t t) { - return has_bnd && bdm.at_unchecked(0, batch * seq + t) != Scalar{0}; - }; + const auto doc_reset = [bdm, batch, seq, has_bnd](std::size_t t) + { return has_bnd && bdm.at_unchecked(0, batch * seq + t) != Scalar{0}; }; std::vector S(dk * dk, Scalar{0}), xv(dk), yv(dk); for (std::size_t i = seq; i-- > 0;) @@ -1446,8 +1371,7 @@ class CpuEngine final : public ComputeEngine if (causal) { // 后缀和：S = Σ_{t>=i} D_t；跨入前一文档时重置 - if (i + 1 < seq && doc_reset(i + 1)) - std::fill(S.begin(), S.end(), Scalar{0}); + if (i + 1 < seq && doc_reset(i + 1)) std::fill(S.begin(), S.end(), Scalar{0}); for (std::size_t idx = 0; idx < dk * dk; ++idx) S[idx] += Dm.at_unchecked(dbase + idx, i); } @@ -1491,14 +1415,12 @@ class CpuEngine final : public ComputeEngine public: // ── 逐列外积（backward 的 dL/dA、dL/dB 物化）───────────────────────── - [[nodiscard]] Result outer_col_impl( - const Tensor& P, const Tensor& R, const Tensor& S, - std::size_t dk, bool has_scale, - Precision prec = Precision::F32) override + [[nodiscard]] Result outer_col_impl(const Tensor& P, const Tensor& R, const Tensor& S, + std::size_t dk, bool has_scale, + Precision prec = Precision::F32) override { const std::vector ts{P, R, S}; - if (prec == Precision::F32 && all_f32(ts)) - return outer_col_f32_(P, R, S, dk, has_scale); + if (prec == Precision::F32 && all_f32(ts)) return outer_col_f32_(P, R, S, dk, has_scale); auto c = to_f32_all(ts); NN_TRY_CHECK(c); auto r = outer_col_f32_((*c)[0], (*c)[1], (*c)[2], dk, has_scale); @@ -1507,24 +1429,19 @@ class CpuEngine final : public ComputeEngine } private: - [[nodiscard]] Result outer_col_f32_( - const Tensor& P, const Tensor& R, const Tensor& S, - std::size_t dk, bool has_scale) + [[nodiscard]] Result outer_col_f32_(const Tensor& P, const Tensor& R, const Tensor& S, + std::size_t dk, bool has_scale) { for (const auto& t : {P, R, S}) - if (t.is_gpu()) - NN_FAIL("CpuEngine: GPU tensor on CPU engine"); - if (dk == 0) - NN_FAIL("outer_col: dk must be > 0"); + if (t.is_gpu()) NN_FAIL("CpuEngine: GPU tensor on CPU engine"); + if (dk == 0) NN_FAIL("outer_col: dk must be > 0"); const Matrix& Pm = P.cpu_matrix(); const std::size_t rows = Pm.rows(); - const std::size_t seq = Pm.cols(); - if (rows % dk != 0) - NN_FAIL("outer_col: rows not divisible by dk"); + const std::size_t seq = Pm.cols(); + if (rows % dk != 0) NN_FAIL("outer_col: rows not divisible by dk"); if (R.cpu_matrix().rows() != rows || R.cpu_matrix().cols() != seq) NN_FAIL("outer_col: P/R shape mismatch"); - if (has_scale && - (S.cpu_matrix().rows() != rows || S.cpu_matrix().cols() != seq)) + if (has_scale && (S.cpu_matrix().rows() != rows || S.cpu_matrix().cols() != seq)) NN_FAIL("outer_col: S must be (B*H*dk, seq)"); const Matrix& Rm = R.cpu_matrix(); const Matrix& Sm = S.cpu_matrix(); @@ -1557,83 +1474,78 @@ class CpuEngine final : public ComputeEngine // 归约原语 // ══════════════════════════════════════════════════════════════════════ - [[nodiscard]] Result row_reduce_sum_impl(const Tensor& A, Precision P = Precision::F32) override + [[nodiscard]] Result row_reduce_sum_impl(const Tensor& A, + Precision P = Precision::F32) override { // 非 f32 输入 / 非 f32 输出精度 → 通用边界 cast（归约恒 f32 累加） if (A.precision() != Precision::F32 || P != Precision::F32) return unary_(A, P, [this](const Tensor& a) -> Result - { - return row_reduce_sum_impl(a, Precision::F32); - }); + { return row_reduce_sum_impl(a, Precision::F32); }); const Matrix& m = A.cpu_matrix(); - Matrix result = m.row_reduce(Scalar{0}, - [](Scalar a, Scalar b) noexcept { return a + b; }, + Matrix result = m.row_reduce( + Scalar{0}, [](Scalar a, Scalar b) noexcept { return a + b; }, [](Scalar x) noexcept { return x; }); return Tensor::from_matrix(std::move(result)); } - [[nodiscard]] Result col_reduce_sum_impl(const Tensor& A, Precision P = Precision::F32) override + [[nodiscard]] Result col_reduce_sum_impl(const Tensor& A, + Precision P = Precision::F32) override { if (A.precision() != Precision::F32 || P != Precision::F32) return unary_(A, P, [this](const Tensor& a) -> Result - { - return col_reduce_sum_impl(a, Precision::F32); - }); + { return col_reduce_sum_impl(a, Precision::F32); }); const Matrix& m = A.cpu_matrix(); - Matrix result = m.col_reduce(Scalar{0}, - [](Scalar a, Scalar b) noexcept { return a + b; }, + Matrix result = m.col_reduce( + Scalar{0}, [](Scalar a, Scalar b) noexcept { return a + b; }, [](Scalar x) noexcept { return x; }); return Tensor::from_matrix(std::move(result)); } - [[nodiscard]] Result col_reduce_max_impl(const Tensor& A, Precision P = Precision::F32) override + [[nodiscard]] Result col_reduce_max_impl(const Tensor& A, + Precision P = Precision::F32) override { if (A.precision() != Precision::F32 || P != Precision::F32) return unary_(A, P, [this](const Tensor& a) -> Result - { - return col_reduce_max_impl(a, Precision::F32); - }); + { return col_reduce_max_impl(a, Precision::F32); }); const Matrix& m = A.cpu_matrix(); Matrix result = m.col_reduce( - std::numeric_limits::lowest(), - [](Scalar a, Scalar b) noexcept { return std::max(a, b); }, - [](Scalar x) noexcept { return x; }); + std::numeric_limits::lowest(), [](Scalar a, Scalar b) noexcept + { return std::max(a, b); }, [](Scalar x) noexcept { return x; }); return Tensor::from_matrix(std::move(result)); } // ── 分组归约（契约见 compute_engine.hpp）─────────────────────────── // 按组并行（组间无竞争），组内沿 R 行累加、内层对列连续（可向量化）。 template - [[nodiscard]] Result grouped_reduce_cpu_( - const Tensor& x, std::size_t G, std::size_t R) + [[nodiscard]] Result grouped_reduce_cpu_(const Tensor& x, std::size_t G, std::size_t R) { - if (x.is_gpu()) - NN_FAIL("CpuEngine: GPU tensor on CPU engine"); - if (G == 0 || R == 0) - NN_FAIL("grouped_reduce: G/R must be > 0"); - if (x.rows() != G * R) - NN_FAIL("grouped_reduce: x must be (G*R, N)"); + if (x.is_gpu()) NN_FAIL("CpuEngine: GPU tensor on CPU engine"); + if (G == 0 || R == 0) NN_FAIL("grouped_reduce: G/R must be > 0"); + if (x.rows() != G * R) NN_FAIL("grouped_reduce: x must be (G*R, N)"); const std::size_t N = x.cols(); - if (N == 0) - NN_FAIL("grouped_reduce: N must be > 0"); + if (N == 0) NN_FAIL("grouped_reduce: N must be > 0"); Matrix out(G, N); const auto src = x.cpu_matrix().span(); const auto dst = out.span(); const Scalar init = IsMax ? std::numeric_limits::lowest() : Scalar{0}; - auto group_kernel = [&, src, dst, init](std::size_t g) noexcept { + auto group_kernel = [&, src, dst, init](std::size_t g) noexcept + { Scalar* o = dst.data() + g * N; - for (std::size_t n = 0; n < N; ++n) o[n] = init; + for (std::size_t n = 0; n < N; ++n) + o[n] = init; for (std::size_t i = 0; i < R; ++i) { const Scalar* row = src.data() + (g * R + i) * N; if constexpr (IsMax) { - for (std::size_t n = 0; n < N; ++n) o[n] = std::max(o[n], row[n]); + for (std::size_t n = 0; n < N; ++n) + o[n] = std::max(o[n], row[n]); } else { - for (std::size_t n = 0; n < N; ++n) o[n] += row[n]; + for (std::size_t n = 0; n < N; ++n) + o[n] += row[n]; } } }; @@ -1645,32 +1557,29 @@ class CpuEngine final : public ComputeEngine } else { - for (std::size_t g = 0; g < G; ++g) group_kernel(g); + for (std::size_t g = 0; g < G; ++g) + group_kernel(g); } return Tensor::from_matrix(std::move(out)); } - [[nodiscard]] Result grouped_reduce_sum_impl( - const Tensor& x, std::size_t G, std::size_t R, - Precision P = Precision::F32) override + [[nodiscard]] Result grouped_reduce_sum_impl(const Tensor& x, std::size_t G, + std::size_t R, + Precision P = Precision::F32) override { if (x.precision() != Precision::F32 || P != Precision::F32) return unary_(x, P, [this, G, R](const Tensor& t) -> Result - { - return grouped_reduce_sum_impl(t, G, R, Precision::F32); - }); + { return grouped_reduce_sum_impl(t, G, R, Precision::F32); }); return grouped_reduce_cpu_(x, G, R); } - [[nodiscard]] Result grouped_reduce_max_impl( - const Tensor& x, std::size_t G, std::size_t R, - Precision P = Precision::F32) override + [[nodiscard]] Result grouped_reduce_max_impl(const Tensor& x, std::size_t G, + std::size_t R, + Precision P = Precision::F32) override { if (x.precision() != Precision::F32 || P != Precision::F32) return unary_(x, P, [this, G, R](const Tensor& t) -> Result - { - return grouped_reduce_max_impl(t, G, R, Precision::F32); - }); + { return grouped_reduce_max_impl(t, G, R, Precision::F32); }); return grouped_reduce_cpu_(x, G, R); } @@ -1682,19 +1591,17 @@ class CpuEngine final : public ComputeEngine // 纯索引映射，无物化。性能与等价手写循环一致。 // ══════════════════════════════════════════════════════════════════════ - [[nodiscard]] Result eval_expr_impl( - const ExprSpec& spec, - std::span inputs, - std::size_t rows, std::size_t cols, - Precision P = Precision::F32) override + [[nodiscard]] Result eval_expr_impl(const ExprSpec& spec, + std::span inputs, std::size_t rows, + std::size_t cols, + Precision P = Precision::F32) override { scan_record_variant_(spec, inputs, P); trace_variant_(spec, inputs, P); // CPU 无带类型融合变体：任何非 f32（输入或输出精度）都走边界 cast if (P != Precision::F32 || !all_f32(inputs)) { - if (prec_trace_enabled_()) - trace_miss_(spec, inputs, P); + if (prec_trace_enabled_()) trace_miss_(spec, inputs, P); auto in32 = to_f32_all(inputs); NN_TRY_CHECK(in32); auto r = eval_expr_impl(spec, *in32, rows, cols, Precision::F32); @@ -1714,29 +1621,26 @@ class CpuEngine final : public ComputeEngine // 行归约轴 → (rows,1)，列归约轴 → (1,cols)（而非广播到 (rows,cols)）。 // 要求：表达式归约轴为 0/1；末指令为归约时直接取归约向量，否则按代表 // 元素求值（所有 Input 须经归约/广播视图访问，保证沿归约轴恒定）。 - [[nodiscard]] Result eval_expr_reduce_impl( - const ExprSpec& spec, - std::span inputs, - std::size_t rows, std::size_t cols, - Precision P = Precision::F32) override + [[nodiscard]] Result eval_expr_reduce_impl(const ExprSpec& spec, + std::span inputs, + std::size_t rows, std::size_t cols, + Precision P = Precision::F32) override { scan_record_variant_(spec, inputs, P); trace_variant_(spec, inputs, P); if (P != Precision::F32 || !all_f32(inputs)) { - if (prec_trace_enabled_()) - trace_miss_(spec, inputs, P); + if (prec_trace_enabled_()) trace_miss_(spec, inputs, P); auto in32 = to_f32_all(inputs); NN_TRY_CHECK(in32); auto r = eval_expr_reduce_impl(spec, *in32, rows, cols, Precision::F32); NN_TRY_CHECK(r); return to_prec(std::move(*r), P); } - const bool vec_is_row = - (expr_spec_reduce_axis(canonicalize_expr_spec(spec)) == 0); + const bool vec_is_row = (expr_spec_reduce_axis(canonicalize_expr_spec(spec)) == 0); // 归约向量输出同样被完整写满 → 未初始化构造 - Tensor out = vec_is_row ? Tensor::cpu_uninitialized(rows, 1) - : Tensor::cpu_uninitialized(1, cols); + Tensor out = + vec_is_row ? Tensor::cpu_uninitialized(rows, 1) : Tensor::cpu_uninitialized(1, cols); auto r = eval_expr_impl(spec, inputs, rows, cols, /*vector_out=*/true, out); NN_TRY_CHECK(r); return out; @@ -1751,10 +1655,9 @@ class CpuEngine final : public ComputeEngine // （就地安全前提：每个元素先读完全部输入再写 out[i]；GPU 同一 buffer 绑定为 // readonly 输入 + writeonly 输出的逐元素同索引读写亦无跨调用危害）。 // 仅支持逐元素表达式（无归约）；归约向量输出用 eval_expr_reduce。 - [[nodiscard]] Result eval_expr_into_impl( - const ExprSpec& spec, - std::span inputs, - std::size_t rows, std::size_t cols, Tensor& dst) override + [[nodiscard]] Result eval_expr_into_impl(const ExprSpec& spec, + std::span inputs, std::size_t rows, + std::size_t cols, Tensor& dst) override { scan_record_variant_(spec, inputs, dst.precision()); // 目标传递：out 可能是 f16 —— 在 f32 空间求值后 cast_into 写回原存储 @@ -1762,8 +1665,7 @@ class CpuEngine final : public ComputeEngine if (dst.precision() != Precision::F32 || !all_f32(inputs)) { trace_variant_(spec, inputs, dst.precision()); - if (prec_trace_enabled_()) - trace_miss_(spec, inputs, dst.precision()); + if (prec_trace_enabled_()) trace_miss_(spec, inputs, dst.precision()); auto in32 = to_f32_all(inputs); NN_TRY_CHECK(in32); auto tmp = create_tensor(rows, cols, Precision::F32); @@ -1772,8 +1674,7 @@ class CpuEngine final : public ComputeEngine return cast_into(tmp, dst); } trace_variant_(spec, inputs, Precision::F32); - if (dst.rows() != rows || dst.cols() != cols) - NN_FAIL("eval_expr_into: dst shape mismatch"); + if (dst.rows() != rows || dst.cols() != cols) NN_FAIL("eval_expr_into: dst shape mismatch"); if (expr_spec_reduce_axis(canonicalize_expr_spec(spec)) != -1) NN_FAIL("eval_expr_into: 仅支持逐元素表达式（无归约）"); return eval_expr_impl(spec, inputs, rows, cols, /*vector_out=*/false, dst); @@ -1784,14 +1685,14 @@ class CpuEngine final : public ComputeEngine // 把目标寄存器提升为主输出（解释器单输出写回；Neg∘Neg 逐位恒等——含 -0 // 与 NaN；末指令是输出 ⇒ simplify 的输出指令豁免不折叠该链）。CPU 无 // 运行时图录制，本路径仅测试/工具使用，重复求值可接受。 - [[nodiscard]] Result eval_expr_multi_into_impl( - const ExprSpec& raw_spec, std::span inputs, - std::size_t rows, std::size_t cols, std::span outs) override + [[nodiscard]] Result eval_expr_multi_into_impl(const ExprSpec& raw_spec, + std::span inputs, + std::size_t rows, std::size_t cols, + std::span outs) override { - if (raw_spec.instrs.empty()) - NN_FAIL("eval_expr_multi_into: 空指令表"); + if (raw_spec.instrs.empty()) NN_FAIL("eval_expr_multi_into: 空指令表"); auto r0 = eval_expr_impl(raw_spec, inputs, rows, cols, - /*vector_out=*/false, outs[0]); + /*vector_out=*/false, outs[0]); NN_TRY_CHECK(r0); for (std::size_t k = 0; k < raw_spec.extras.size(); ++k) { @@ -1801,18 +1702,18 @@ class CpuEngine final : public ComputeEngine NN_FAIL("eval_expr_multi_into: extras 复制链超出寄存器上限"); ExprSpec s = raw_spec; ExprInstr n1; - n1.op = static_cast(ExprOp::Neg); + n1.op = static_cast(ExprOp::Neg); n1.dst = static_cast(s.num_regs); - n1.a = expr::reg(raw_spec.extras[k]); + n1.a = expr::reg(raw_spec.extras[k]); s.instrs.push_back(n1); ExprInstr n2; - n2.op = static_cast(ExprOp::Neg); + n2.op = static_cast(ExprOp::Neg); n2.dst = static_cast(s.num_regs + 1); - n2.a = expr::reg(static_cast(s.num_regs)); + n2.a = expr::reg(static_cast(s.num_regs)); s.instrs.push_back(n2); s.num_regs += 2; auto r = eval_expr_impl(s, inputs, rows, cols, - /*vector_out=*/false, outs[k + 1]); + /*vector_out=*/false, outs[k + 1]); NN_TRY_CHECK(r); } return {}; @@ -1828,17 +1729,14 @@ class CpuEngine final : public ComputeEngine // - 键域每块：[mm 段预计算 tile] → body 指令 →（块尾）vecacc 更新行向量态 // - 块内有效列 valid = min(BLOCK, K-k0) 天然 gate 尾块越界 // - 向量域：finalize 逐 (row, d∈[0,out_cols)) 求值 → 输出 - [[nodiscard]] Result eval_fold_impl( - const ExprSpec& spec, - std::span inputs, - std::size_t rows, std::size_t cols, bool vector_out, Tensor& output) + [[nodiscard]] Result eval_fold_impl(const ExprSpec& spec, std::span inputs, + std::size_t rows, std::size_t cols, bool vector_out, + Tensor& output) { - if (vector_out) - NN_FAIL("eval_fold: vector_out not expected (fold outputs via eval_expr)"); + if (vector_out) NN_FAIL("eval_fold: vector_out not expected (fold outputs via eval_expr)"); const FoldSpec& f = *spec.fold; const std::size_t K = f.k; - const std::size_t out_cols = - (f.vec_state_len > 0) ? f.vec_state_len : 1; + const std::size_t out_cols = (f.vec_state_len > 0) ? f.vec_state_len : 1; if (cols != out_cols) NN_FAIL("eval_fold: cols must equal fold output width (vec_state_len or 1)"); // 输入：分 fold.matmul 的 A/B（按 trans 布局）与普通输入（视图校验） @@ -1847,40 +1745,33 @@ class CpuEngine final : public ComputeEngine for (std::size_t k = 0; k < inputs.size(); ++k) { const Tensor& t = inputs[k]; - if (!t.is_cpu()) - NN_FAIL("eval_fold: input not CPU"); + if (!t.is_cpu()) NN_FAIL("eval_fold: input not CPU"); const ExprViewKind vk = static_cast(spec.views[k].kind); // vecacc.b（(K, vec) 独立校验于下方）与掩码小表（BatchMod/BatchCol // 为 (1, *) 标签/斜率向量，行主序平坦读、不占 (rows,K) 网格） // 豁免普通输入的 (rows, K) 形状假设 - const bool is_vecb = f.vecacc && - (k == static_cast(f.vecacc->b_input)); + const bool is_vecb = f.vecacc && (k == static_cast(f.vecacc->b_input)); const bool is_tag_table = (vk == ExprViewKind::BatchMod || vk == ExprViewKind::BatchCol); - const bool is_rowvec = (vk == ExprViewKind::RowBroadcast); // (rows,1) 行参数 + const bool is_rowvec = (vk == ExprViewKind::RowBroadcast); // (rows,1) 行参数 if (!is_vecb && !is_tag_table && !is_rowvec && f.matmul && (k == static_cast(f.matmul->a_input) || k == static_cast(f.matmul->b_input))) { const bool is_a = (k == static_cast(f.matmul->a_input)); - const bool tr = is_a ? (f.matmul->transA != 0) - : (f.matmul->transB != 0); + const bool tr = is_a ? (f.matmul->transA != 0) : (f.matmul->transB != 0); const std::size_t batch = f.matmul->batch; const std::size_t mm_k = f.matmul->k; - const std::size_t m_per = - (batch > 0 && rows % batch == 0) ? rows / batch : rows; + const std::size_t m_per = (batch > 0 && rows % batch == 0) ? rows / batch : rows; // 形状按 MatmulSpec 布局核对（同 eval_expr_impl 的 matmul 分支）： // A: trans0 → (rows, mm_k)；trans1 → (batch*mm_k, m_per) // B: trans0 → (batch*mm_k, K)；trans1 → (batch*K, mm_k) // 其中 mm_k=内层收缩、K=fold 收缩轴=网格列数 const std::size_t nb_a = batch / gqa_ratio_of(spec, f.matmul->a_input); const std::size_t nb_b = batch / gqa_ratio_of(spec, f.matmul->b_input); - const std::size_t er = is_a - ? (tr ? nb_a * mm_k : rows) - : (tr ? nb_b * K : nb_b * mm_k); - const std::size_t ec = is_a - ? (tr ? m_per : mm_k) - : (tr ? mm_k : K); + const std::size_t er = + is_a ? (tr ? nb_a * mm_k : rows) : (tr ? nb_b * K : nb_b * mm_k); + const std::size_t ec = is_a ? (tr ? m_per : mm_k) : (tr ? mm_k : K); if (t.rows() != er || t.cols() != ec) NN_FAIL("eval_fold: fold matmul input shape mismatch"); spans.push_back(t.cpu_matrix().span()); @@ -1889,11 +1780,9 @@ class CpuEngine final : public ComputeEngine // 视图白名单（与 expr_spec 的 fold validate **同步维护**——两处清单 // 必须一致，否则合法视图在一侧被拒、fold forward 直接报错） if (vk != ExprViewKind::Linear && vk != ExprViewKind::HeadGroup && - vk != ExprViewKind::RowMod && - vk != ExprViewKind::RowBroadcast && + vk != ExprViewKind::RowMod && vk != ExprViewKind::RowBroadcast && vk != ExprViewKind::Transpose && - !((vk == ExprViewKind::BatchMod || vk == ExprViewKind::BatchCol) && - f.matmul)) + !((vk == ExprViewKind::BatchMod || vk == ExprViewKind::BatchCol) && f.matmul)) NN_FAIL("eval_fold: view outside supported scope"); const bool is_trans = (vk == ExprViewKind::Transpose); if (!is_vecb && !is_tag_table && !is_rowvec && !is_trans && @@ -1916,12 +1805,11 @@ class CpuEngine final : public ComputeEngine if (is_tag_table) { const std::uint32_t vparam = spec.views[k].param; - if (vparam == 0) - NN_FAIL("eval_fold: BatchMod/BatchCol param must be > 0"); + if (vparam == 0) NN_FAIL("eval_fold: BatchMod/BatchCol param must be > 0"); const std::size_t batch_n = f.matmul ? f.matmul->batch : 1; const std::size_t need = vk == ExprViewKind::BatchMod - ? static_cast(vparam) - : (batch_n - 1) * static_cast(vparam) + K; + ? static_cast(vparam) + : (batch_n - 1) * static_cast(vparam) + K; if (t.cols() < need) NN_FAIL("eval_fold: tag table cols out of range (BatchMod/BatchCol)"); } @@ -1936,12 +1824,11 @@ class CpuEngine final : public ComputeEngine const Tensor& tb = inputs[f.vecacc->b_input]; // GQA：V_t 只存 n_head_kv 个头 → 行数按头分组比缩小 const std::size_t vb = - (f.matmul ? f.matmul->batch : 1) / - gqa_ratio_of(spec, f.vecacc->b_input); - const std::size_t want_rows = - (vb ? vb : 1) * K; + (f.matmul ? f.matmul->batch : 1) / gqa_ratio_of(spec, f.vecacc->b_input); + const std::size_t want_rows = (vb ? vb : 1) * K; if (tb.rows() != want_rows || tb.cols() != f.vec_state_len) - NN_FAIL("eval_fold: vecacc b_input shape mismatch (expect (batch*k, vec_state_len))"); + NN_FAIL( + "eval_fold: vecacc b_input shape mismatch (expect (batch*k, vec_state_len))"); } auto cls_r = expr_fold_classify(f, spec.num_regs); @@ -1951,34 +1838,32 @@ class CpuEngine final : public ComputeEngine // fold matmul 上下文（掩码网格索引与 mm 段求值共用） const MatmulSpec* fmm = f.matmul ? &*f.matmul : nullptr; const std::size_t mm_batch = fmm ? fmm->batch : 1; - const std::size_t m_per = (fmm && mm_batch > 0 && rows % mm_batch == 0) - ? rows / mm_batch : rows; + const std::size_t m_per = + (fmm && mm_batch > 0 && rows % mm_batch == 0) ? rows / mm_batch : rows; // 输入视图读（列 = 全局收缩下标 gk；双域 fold 加 BatchMod/BatchCol 掩码钩子） const auto view_read = [&](std::size_t k, std::size_t r, std::size_t gk) -> Scalar { const ExprView& v = spec.views[k]; const auto vk = static_cast(v.kind); - if (vk == ExprViewKind::RowMod && v.param > 0) - return spans[k][(r % v.param) * K + gk]; + if (vk == ExprViewKind::RowMod && v.param > 0) return spans[k][(r % v.param) * K + gk]; if (vk == ExprViewKind::BatchMod) - return spans[k][(r / m_per) % v.param]; // b[batch % param] + return spans[k][(r / m_per) % v.param]; // b[batch % param] if (vk == ExprViewKind::BatchCol) - return spans[k][(r / m_per) * v.param + gk]; // b[batch*param + col] - if (vk == ExprViewKind::RowBroadcast) - return spans[k][r]; // (rows,1)：b[row] + return spans[k][(r / m_per) * v.param + gk]; // b[batch*param + col] + if (vk == ExprViewKind::RowBroadcast) return spans[k][r]; // (rows,1)：b[row] if (vk == ExprViewKind::Transpose) - return spans[k][gk * static_cast(v.param) + r]; // b[col*stride+row] + return spans[k][gk * static_cast(v.param) + r]; // b[col*stride+row] return spans[k][r * K + gk]; }; - std::vector reg(spec.num_regs, Scalar{0}); // 状态 + 标量类临时 + std::vector reg(spec.num_regs, Scalar{0}); // 状态 + 标量类临时 std::vector> elem(spec.num_regs); // 元素类临时（BLOCK 槽） for (std::uint32_t r = 0; r < spec.num_regs; ++r) if (is_elem[r]) elem[r].resize(EXPR_FOLD_BLOCK); - std::vector red(spec.num_regs, Scalar{0}); // 块归约 dst（行标量） + std::vector red(spec.num_regs, Scalar{0}); // 块归约 dst（行标量） std::vector vecd(f.vec_state_len, Scalar{0}); // 行向量态（累加初值 0） - std::size_t cur_d = 0; // 向量域当前输出列（eval_op 的 VecState 读） + std::size_t cur_d = 0; // 向量域当前输出列（eval_op 的 VecState 读） // fold matmul 段：网格 (row, gk) 上 C = Σ_d opA·opB（惰性求值—— // CPU 正确性基准不追性能；GPU 侧由生成器块内预计算 tile）。 @@ -1997,20 +1882,18 @@ class CpuEngine final : public ComputeEngine Scalar acc = 0; for (std::size_t d = 0; d < mm_k; ++d) { - const Scalar av = fmm->transA - ? as[(bA * mm_k + d) * m_per + row_in] - : as[(bA * m_per + row_in) * mm_k + d]; - const Scalar bv = fmm->transB - ? bs[(bB * K + gk) * mm_k + d] - : bs[(bB * mm_k + d) * K + gk]; + const Scalar av = fmm->transA ? as[(bA * mm_k + d) * m_per + row_in] + : as[(bA * m_per + row_in) * mm_k + d]; + const Scalar bv = + fmm->transB ? bs[(bB * K + gk) * mm_k + d] : bs[(bB * mm_k + d) * K + gk]; acc += av * bv; } return acc; }; // 操作数求值：kb = 块内下标（元素类临时槽），gk = 全局列（Input/Col 读） - const auto eval_op = [&](const ExprOperand& op, std::size_t r, - std::size_t kb, std::size_t gk) -> Scalar + const auto eval_op = [&](const ExprOperand& op, std::size_t r, std::size_t kb, + std::size_t gk) -> Scalar { switch (op.kind) { @@ -2018,11 +1901,16 @@ class CpuEngine final : public ComputeEngine case static_cast(ExprOperandKind::Reg): case static_cast(ExprOperandKind::Fanout): return is_elem[op.idx] ? elem[op.idx][kb] : reg[op.idx]; - case static_cast(ExprOperandKind::Const): return spec.consts[op.idx]; - case static_cast(ExprOperandKind::RParam): return spec.rparams[op.idx]; - case static_cast(ExprOperandKind::Reduce): return red[op.idx]; - case static_cast(ExprOperandKind::Input): return view_read(op.idx, r, gk); - case static_cast(ExprOperandKind::Matmul): return mm_at(r, gk); + case static_cast(ExprOperandKind::Const): + return spec.consts[op.idx]; + case static_cast(ExprOperandKind::RParam): + return spec.rparams[op.idx]; + case static_cast(ExprOperandKind::Reduce): + return red[op.idx]; + case static_cast(ExprOperandKind::Input): + return view_read(op.idx, r, gk); + case static_cast(ExprOperandKind::Matmul): + return mm_at(r, gk); case static_cast(ExprOperandKind::Row): return fmm ? static_cast(r % m_per) : Scalar{0}; case static_cast(ExprOperandKind::Col): @@ -2037,12 +1925,10 @@ class CpuEngine final : public ComputeEngine // 非归约指令求值（标量语义唯一来源 = expr_op_table 的 CPU 函数指针对； // va 已取自 a、vb 取自 b；Select 由 eval_full 单独处理，归约不进本函数） const auto exec_alu = [](ExprOp op, Scalar va, Scalar vb) -> Scalar - { - return detail::expr_alu_scalar(op, va, vb, Scalar{0}); - }; + { return detail::expr_alu_scalar(op, va, vb, Scalar{0}); }; // Select：三元（fold body 允许，validate 未禁——按主链语义） - const auto eval_full = [&](const ExprInstr& ins, std::size_t r, - std::size_t kb, std::size_t gk) -> Scalar + const auto eval_full = [&](const ExprInstr& ins, std::size_t r, std::size_t kb, + std::size_t gk) -> Scalar { const ExprOp op = static_cast(ins.op); const Scalar va = eval_op(ins.a, r, kb, gk); @@ -2079,8 +1965,7 @@ class CpuEngine final : public ComputeEngine if (expr_op_is_reduce(op)) { const bool is_max = (op == ExprOp::RowMax || op == ExprOp::ColMax); - Scalar acc = is_max - ? std::numeric_limits::lowest() : Scalar{0}; + Scalar acc = is_max ? std::numeric_limits::lowest() : Scalar{0}; for (std::size_t kb = 0; kb < valid; ++kb) { const Scalar v = eval_op(ins.a, r, kb, k0 + kb); @@ -2111,14 +1996,13 @@ class CpuEngine final : public ComputeEngine const ConstSpan& bs = spans[va.b_input]; // GQA：V_t 的行块号按其视图的头分组比取 const std::size_t b_of_row = - fmm ? (static_cast(r / m_per) / - gqa_ratio_of(spec, va.b_input)) - : 0; + fmm ? (static_cast(r / m_per) / gqa_ratio_of(spec, va.b_input)) + : 0; for (std::size_t kb = 0; kb < valid; ++kb) { - const Scalar w = is_elem[va.weight_reg] - ? elem[va.weight_reg][kb] : reg[va.weight_reg]; - const std::size_t jr = b_of_row * K + (k0 + kb); // 全局键行 + const Scalar w = + is_elem[va.weight_reg] ? elem[va.weight_reg][kb] : reg[va.weight_reg]; + const std::size_t jr = b_of_row * K + (k0 + kb); // 全局键行 for (std::size_t d = 0; d < f.vec_state_len; ++d) vecd[d] += w * bs[jr * f.vec_state_len + d]; } @@ -2141,23 +2025,20 @@ class CpuEngine final : public ComputeEngine // ── eval_expr / eval_expr_reduce / eval_expr_into 共用实现 ──────────── // output：调用方预分配的输出张量，形状须与输出一致 // （vector_out=false → (rows,cols)；true → (rows,1)/(1,cols)）。 - [[nodiscard]] Result eval_expr_impl( - const ExprSpec& raw_spec, - std::span inputs, - std::size_t rows, std::size_t cols, bool vector_out, Tensor& output) + [[nodiscard]] Result eval_expr_impl(const ExprSpec& raw_spec, + std::span inputs, std::size_t rows, + std::size_t cols, bool vector_out, Tensor& output) { // canonical IR：canonicalize 为引擎内部优化（IR-A/IR-B），Layer 无感知； // 结构统一在 canonical 形态上，与 scan/gen_fused 两端一致。 NN_TRY(v, validate_expr_spec(raw_spec, inputs.size())); // fold 段：独立执行路径——跳过 canonicalize（IR-A/B 不作用于 // body/finalize 序列）与空指令表检查（fold 的 finalize 即尾链） - if (raw_spec.fold) - return eval_fold_impl(raw_spec, inputs, rows, cols, vector_out, output); + if (raw_spec.fold) return eval_fold_impl(raw_spec, inputs, rows, cols, vector_out, output); const ExprSpec spec = canonicalize_expr_spec(raw_spec); NN_TRY(v2, validate_expr_spec(spec, inputs.size())); // 空指令表仅允许"纯 matmul"（输出 = matmul 结果本身） - if (spec.instrs.empty() && !spec.matmul) - NN_FAIL("eval_expr: empty instruction list"); + if (spec.instrs.empty() && !spec.matmul) NN_FAIL("eval_expr: empty instruction list"); // vector_out：要求表达式为纯行/列归约（归约轴 0/1），输出取向量形状 const int raxis = vector_out ? expr_spec_reduce_axis(spec) : -1; if (vector_out && raxis != 0 && raxis != 1) @@ -2174,17 +2055,17 @@ class CpuEngine final : public ComputeEngine const MatmulSpec* mm = spec.matmul ? &*spec.matmul : nullptr; const std::size_t mm_k = mm ? mm->k : 0; const std::size_t mm_batch = mm ? mm->batch : 1; - const std::size_t m_per = (mm && mm->batch > 0 && rows % mm->batch == 0) - ? rows / mm->batch : rows; + const std::size_t m_per = + (mm && mm->batch > 0 && rows % mm->batch == 0) ? rows / mm->batch : rows; for (std::size_t k = 0; k < inputs.size(); ++k) { const Tensor& t = inputs[k]; - if (!t.is_cpu()) - NN_FAIL("eval_expr: input not CPU"); + if (!t.is_cpu()) NN_FAIL("eval_expr: input not CPU"); // CPU 解释器按 f32 存储读取 span：f16 输入直读 = 空指针 UB → // 响亮报错，让误用立刻暴露。 if (t.precision() != Precision::F32) - NN_FAIL("eval_expr: CPU 解释器仅支持 f32 输入（f16 须由基类 NVI 入口先抬到 f32（直调 eval_expr_impl 才会到这里））"); + NN_FAIL("eval_expr: CPU 解释器仅支持 f32 输入（f16 须由基类 NVI 入口先抬到 " + "f32（直调 eval_expr_impl 才会到这里））"); const ExprView& v = spec.views[k]; if (mm && (k == static_cast(mm->a_input) || k == static_cast(mm->b_input))) @@ -2195,12 +2076,10 @@ class CpuEngine final : public ComputeEngine // GQA：操作数只存 mm_batch/kv_ratio 份行块（1 = 恒等） const std::size_t nb = mm_batch / gqa_ratio_of(spec, is_a ? mm->a_input : mm->b_input); - const std::size_t er = is_a - ? (tr ? nb * mm_k : nb * m_per) // A 存储行数期望 - : (tr ? nb * cols : nb * mm_k); - const std::size_t ec = is_a - ? (tr ? m_per : mm_k) // A 存储列数期望 - : (tr ? mm_k : cols); + const std::size_t er = is_a ? (tr ? nb * mm_k : nb * m_per) // A 存储行数期望 + : (tr ? nb * cols : nb * mm_k); + const std::size_t ec = is_a ? (tr ? m_per : mm_k) // A 存储列数期望 + : (tr ? mm_k : cols); if (t.rows() != er || t.cols() != ec) NN_FAIL("eval_expr: matmul input shape mismatch"); spans.push_back(t.cpu_matrix().span()); @@ -2218,12 +2097,11 @@ class CpuEngine final : public ComputeEngine case static_cast(ExprViewKind::RowReduceMax): case static_cast(ExprViewKind::RowGather): // 按标签行收集：主输入行数任意（读取经 label 索引），列数须一致 - if (t.cols() != cols) - NN_FAIL("eval_expr: RowGather input cols mismatch"); + if (t.cols() != cols) NN_FAIL("eval_expr: RowGather input cols mismatch"); break; case static_cast(ExprViewKind::RotateHalf): - if (t.rows() != rows || t.cols() != cols || v.param == 0 || - rows % v.param != 0 || v.param % 2 != 0) + if (t.rows() != rows || t.cols() != cols || v.param == 0 || rows % v.param != 0 || + v.param % 2 != 0) NN_FAIL("eval_expr: RotateHalf shape/param invalid"); break; case static_cast(ExprViewKind::RowMod): @@ -2254,8 +2132,7 @@ class CpuEngine final : public ComputeEngine break; case static_cast(ExprViewKind::BatchMod): // 输入 (1,modulo 相关)：按批次取模索引的小向量 - if (t.rows() != 1) - NN_FAIL("eval_expr: BatchMod input must be (1,cols)"); + if (t.rows() != 1) NN_FAIL("eval_expr: BatchMod input must be (1,cols)"); break; case static_cast(ExprViewKind::BatchCol): // 输入 (1, batch*param)：按 (batch, col) 切片（doc_ids 等） @@ -2278,8 +2155,7 @@ class CpuEngine final : public ComputeEngine NN_FAIL("eval_expr: CPU 解释器输出须为 f32 存储（f16 输出须由基类 NVI 入口落回）"); Span out = output.cpu_matrix().span(); const std::size_t n = rows * cols; - if (n == 0) - return {}; + if (n == 0) return {}; // ── 归约视图预计算：每行/每列一个标量，供广播读取 ────────────── // 仅在 views[k] 为归约视图时填充 view_reduce[k]（长度 rows 或 cols）。 @@ -2290,12 +2166,10 @@ class CpuEngine final : public ComputeEngine for (std::size_t k = 0; k < inputs.size(); ++k) { const ExprView& v = spec.views[k]; - if (!expr_view_is_reduce(static_cast(v.kind))) - continue; + if (!expr_view_is_reduce(static_cast(v.kind))) continue; const bool is_row = expr_view_reduces_rows(static_cast(v.kind)); - const bool is_max = - (v.kind == static_cast(ExprViewKind::RowReduceMax) || - v.kind == static_cast(ExprViewKind::ColReduceMax)); + const bool is_max = (v.kind == static_cast(ExprViewKind::RowReduceMax) || + v.kind == static_cast(ExprViewKind::ColReduceMax)); const ConstSpan& s = spans[k]; const std::size_t len = is_row ? rows : cols; std::vector& acc = view_reduce[k]; @@ -2346,12 +2220,12 @@ class CpuEngine final : public ComputeEngine const std::size_t mm_batch = mm->batch; const bool trA = (mm->transA != 0); const bool trB = (mm->transB != 0); - const std::size_t M = m_per; // 每批输出行 + const std::size_t M = m_per; // 每批输出行 const std::size_t N = cols; const std::size_t K = mm_k; const ConstSpan& as = spans[mm->a_input]; const ConstSpan& bs = spans[mm->b_input]; - matmul_out = Matrix(rows, N); // 零初始化：span 内核为累加语义（r += A*B） + matmul_out = Matrix(rows, N); // 零初始化：span 内核为累加语义（r += A*B） auto out = matmul_out.span(); for (std::size_t b = 0; b < mm_batch; ++b) { @@ -2368,21 +2242,15 @@ class CpuEngine final : public ComputeEngine auto c_sub = std::span(out.data() + b * M * N, M * N); if (!trA && !trB) { - Matrix::multiply_to_span(c_sub, M, N, - a_sub, M, K, - b_sub, K, N); + Matrix::multiply_to_span(c_sub, M, N, a_sub, M, K, b_sub, K, N); } else if (!trA && trB) { - Matrix::multiply_transposed_to_span(c_sub, M, N, - a_sub, M, K, - b_sub, N, K); + Matrix::multiply_transposed_to_span(c_sub, M, N, a_sub, M, K, b_sub, N, K); } else if (trA && !trB) { - Matrix::transpose_multiply_to_span(c_sub, M, N, - a_sub, K, M, - b_sub, K, N); + Matrix::transpose_multiply_to_span(c_sub, M, N, a_sub, K, M, b_sub, K, N); } else { @@ -2404,8 +2272,7 @@ class CpuEngine final : public ComputeEngine // ── 纯 matmul（无逐元素链）：输出 = matmul 结果本身 ───────────── if (spec.instrs.empty()) { - if (!mm) - NN_FAIL("eval_expr: empty instruction list without matmul"); + if (!mm) NN_FAIL("eval_expr: empty instruction list without matmul"); const auto ms = matmul_out.span(); for (std::size_t i = 0; i < n; ++i) out[i] = ms[i]; @@ -2427,11 +2294,10 @@ class CpuEngine final : public ComputeEngine { const std::size_t block = v.param; const std::size_t rl = r % block; - const std::size_t rr = (r / block) * block - + ((rl < block / 2) ? (rl + block / 2) : (rl - block / 2)); + const std::size_t rr = + (r / block) * block + ((rl < block / 2) ? (rl + block / 2) : (rl - block / 2)); Scalar val = s[rr * cols + c]; - if (v.negate_first_half && rl < block / 2) - return -val; + if (v.negate_first_half && rl < block / 2) return -val; return val; } case static_cast(ExprViewKind::RowMod): @@ -2449,22 +2315,21 @@ class CpuEngine final : public ComputeEngine case static_cast(ExprViewKind::ColReduceMax): return view_reduce[k][c]; case static_cast(ExprViewKind::RowBroadcast): - return s[r]; // 输入 (rows,1) + return s[r]; // 输入 (rows,1) case static_cast(ExprViewKind::ColBroadcast): - return s[c]; // 输入 (1,cols) + return s[c]; // 输入 (1,cols) case static_cast(ExprViewKind::RowGather): { // 按标签行收集：data[label[c] * cols + c]（标签槽 = v.param）； // 越界守卫按主输入行数（s.size()/cols），label 越界返回 0 const ConstSpan& ls = spans[v.param]; - const std::size_t lbl = (ls[c] >= Scalar{0}) - ? static_cast(ls[c]) : 0; + const std::size_t lbl = (ls[c] >= Scalar{0}) ? static_cast(ls[c]) : 0; return (lbl < (s.size() / cols)) ? s[lbl * cols + c] : Scalar{0}; } case static_cast(ExprViewKind::BatchMod): - return s[b_idx % v.param]; // 按批次取模索引（ALiBi 斜率等） + return s[b_idx % v.param]; // 按批次取模索引（ALiBi 斜率等） case static_cast(ExprViewKind::BatchCol): - return s[b_idx * v.param + c]; // 按 (batch, col) 切片（doc_ids 等） + return s[b_idx * v.param + c]; // 按 (batch, col) 切片（doc_ids 等） case static_cast(ExprViewKind::GroupedReduceSum): case static_cast(ExprViewKind::GroupedReduceMax): { @@ -2502,8 +2367,7 @@ class CpuEngine final : public ComputeEngine { const ExprInstr& R = spec.instrs[ri]; const ExprOp rop = static_cast(R.op); - if (!expr_op_is_reduce(rop)) - continue; + if (!expr_op_is_reduce(rop)) continue; const bool is_col = expr_op_reduces_cols(rop); const bool is_max = (rop == ExprOp::ColMax || rop == ExprOp::RowMax); @@ -2528,16 +2392,20 @@ class CpuEngine final : public ComputeEngine { default: case static_cast(ExprOperandKind::Reg): - case static_cast(ExprOperandKind::Fanout): return regs[op.idx]; - case static_cast(ExprOperandKind::Const): return spec.consts[op.idx]; - case static_cast(ExprOperandKind::RParam): return spec.rparams[op.idx]; - case static_cast(ExprOperandKind::Input): return read_input(op.idx, r, c); + case static_cast(ExprOperandKind::Fanout): + return regs[op.idx]; + case static_cast(ExprOperandKind::Const): + return spec.consts[op.idx]; + case static_cast(ExprOperandKind::RParam): + return spec.rparams[op.idx]; + case static_cast(ExprOperandKind::Input): + return read_input(op.idx, r, c); case static_cast(ExprOperandKind::Reduce): return reduce_vec[op.idx][reduce_axis[op.idx] ? c : r]; case static_cast(ExprOperandKind::Matmul): - return matmul_out.span()[r * cols + c]; // matmul 段输出（虚拟寄存器 0） + return matmul_out.span()[r * cols + c]; // matmul 段输出（虚拟寄存器 0） case static_cast(ExprOperandKind::Row): - return static_cast(r % m_per); // batch 内行号 + return static_cast(r % m_per); // batch 内行号 case static_cast(ExprOperandKind::Col): return static_cast(c); case static_cast(ExprOperandKind::Batch): @@ -2549,17 +2417,15 @@ class CpuEngine final : public ComputeEngine { const ExprInstr& J = spec.instrs[j]; if (expr_op_is_reduce(static_cast(J.op))) - continue; // 归约指令已在上方处理中完整计算 + continue; // 归约指令已在上方处理中完整计算 const Scalar va = eval_op(J.a); const ExprOp op = static_cast(J.op); // 算子语义查表（expr_op_table）——一元/二元/三元分流 const detail::ExprOpInfo& info = detail::expr_op_info(op); if (info.arity == 3) regs[J.dst] = (va != Scalar{0}) ? eval_op(J.b) : eval_op(J.c); - else if (info.arity == 2) - regs[J.dst] = info.binary(va, eval_op(J.b)); - else - regs[J.dst] = info.unary(va); + else if (info.arity == 2) regs[J.dst] = info.binary(va, eval_op(J.b)); + else regs[J.dst] = info.unary(va); } return eval_op(R.a); @@ -2617,16 +2483,20 @@ class CpuEngine final : public ComputeEngine { default: case static_cast(ExprOperandKind::Reg): - case static_cast(ExprOperandKind::Fanout): return regs[op.idx]; - case static_cast(ExprOperandKind::Const): return spec.consts[op.idx]; - case static_cast(ExprOperandKind::RParam): return spec.rparams[op.idx]; - case static_cast(ExprOperandKind::Input): return read_input(op.idx, r, c); + case static_cast(ExprOperandKind::Fanout): + return regs[op.idx]; + case static_cast(ExprOperandKind::Const): + return spec.consts[op.idx]; + case static_cast(ExprOperandKind::RParam): + return spec.rparams[op.idx]; + case static_cast(ExprOperandKind::Input): + return read_input(op.idx, r, c); case static_cast(ExprOperandKind::Reduce): return reduce_vec[op.idx][reduce_axis[op.idx] ? c : r]; case static_cast(ExprOperandKind::Matmul): - return matmul_out.span()[r * cols + c]; // matmul 段输出（虚拟寄存器 0） + return matmul_out.span()[r * cols + c]; // matmul 段输出（虚拟寄存器 0） case static_cast(ExprOperandKind::Row): - return static_cast(r % m_per); // batch 内行号 + return static_cast(r % m_per); // batch 内行号 case static_cast(ExprOperandKind::Col): return static_cast(c); case static_cast(ExprOperandKind::Batch): @@ -2636,17 +2506,15 @@ class CpuEngine final : public ComputeEngine for (const auto& ins : spec.instrs) { if (expr_op_is_reduce(static_cast(ins.op))) - continue; // 归约指令已预计算（隐式向量，不经寄存器） + continue; // 归约指令已预计算（隐式向量，不经寄存器） const Scalar va = eval_op(ins.a); const ExprOp op = static_cast(ins.op); // 算子语义查表（expr_op_table）——一元/二元/三元分流 const detail::ExprOpInfo& info = detail::expr_op_info(op); if (info.arity == 3) regs[ins.dst] = (va != Scalar{0}) ? eval_op(ins.b) : eval_op(ins.c); - else if (info.arity == 2) - regs[ins.dst] = info.binary(va, eval_op(ins.b)); - else - regs[ins.dst] = info.unary(va); + else if (info.arity == 2) regs[ins.dst] = info.binary(va, eval_op(ins.b)); + else regs[ins.dst] = info.unary(va); } return regs[last.dst]; }; @@ -2655,7 +2523,7 @@ class CpuEngine final : public ComputeEngine { // 归约向量原生形状输出：(rows,1)（行归约轴）或 (1,cols)（列归约轴） const std::size_t len = vec_is_row ? rows : cols; - Span o = out; // 调用方已按归约向量形状分配（见 eval_expr_reduce） + Span o = out; // 调用方已按归约向量形状分配（见 eval_expr_reduce） if (last_is_reduce) { // 末指令为归约 → 直接取归约向量 @@ -2679,17 +2547,15 @@ class CpuEngine final : public ComputeEngine // 与寄存器复用无关（execution 本身按指令序写 regs[]，复用是正确的）。 { std::vector instr_needed(spec.instrs.size(), 0); - instr_needed[spec.instrs.size() - 1] = 1; // 输出 = instrs.back().dst + instr_needed[spec.instrs.size() - 1] = 1; // 输出 = instrs.back().dst for (std::size_t ii = spec.instrs.size(); ii-- > 0;) { - if (!instr_needed[ii]) - continue; + if (!instr_needed[ii]) continue; const ExprInstr& ins = spec.instrs[ii]; if (expr_op_is_reduce(static_cast(ins.op))) - continue; // 归约指令：源全网格求值，不经输出链 + continue; // 归约指令：源全网格求值，不经输出链 // 只遍历该算子实际使用的操作数（c 默认 {0,0} 会被误当作 Reg(0)） - const std::size_t nops = - expr_instr_num_operands(static_cast(ins.op)); + const std::size_t nops = expr_instr_num_operands(static_cast(ins.op)); const ExprOperand* ops[3] = {&ins.a, &ins.b, &ins.c}; for (std::size_t oi = 0; oi < nops; ++oi) { @@ -2699,7 +2565,8 @@ class CpuEngine final : public ComputeEngine const ExprViewKind vk = static_cast(spec.views[op.idx].kind); if (!expr_view_is_reduce(vk) && !expr_view_is_broadcast(vk)) - NN_FAIL("eval_expr_reduce: 输出链经 Linear/视图直接访问输入，" "输出不沿归约轴恒定"); + NN_FAIL("eval_expr_reduce: 输出链经 Linear/视图直接访问输入，" + "输出不沿归约轴恒定"); } else if (op.kind == static_cast(ExprOperandKind::Matmul)) { @@ -2712,7 +2579,8 @@ class CpuEngine final : public ComputeEngine { // 索引操作数沿归约轴的恒定性由生成器/引擎按轴判定； // 这里保守拒绝（实际表达式不触发此路径） - NN_FAIL("eval_expr_reduce: 输出链经 Row/Col/Batch 索引，输出不沿归约轴恒定"); + NN_FAIL("eval_expr_reduce: 输出链经 Row/Col/Batch " + "索引，输出不沿归约轴恒定"); } else if (op.kind == static_cast(ExprOperandKind::Reg) || op.kind == static_cast(ExprOperandKind::Fanout) || @@ -2791,7 +2659,6 @@ class CpuEngine final : public ComputeEngine return {}; } - }; } // namespace nn diff --git a/include/neuralnet.cpp/compute_engine.hpp b/include/neuralnet.cpp/compute_engine.hpp index a9005039..4b8079d8 100644 --- a/include/neuralnet.cpp/compute_engine.hpp +++ b/include/neuralnet.cpp/compute_engine.hpp @@ -73,19 +73,17 @@ namespace nn namespace dsl { - // 前置声明（compute_engine.hpp 不能反向 include expr_dsl.hpp）： - // 默认实参（P = F32）在 expr_dsl.hpp 的定义处。 - // 生命周期 v3 起不再有库内 stamp/adopt 通道：DSL 静态工厂产出的 CPU - // 张量天然属于宿主域，设备域从存储派生。 - template - [[nodiscard]] Result compute(ComputeEngine& eng, const E& e, - std::size_t rows, std::size_t cols, - Precision P = Precision::F32); - template - [[nodiscard]] Result compute_reduce(ComputeEngine& eng, const E& e, - std::size_t rows, std::size_t cols, - Precision P = Precision::F32); -} // namespace nn::dsl +// 前置声明（compute_engine.hpp 不能反向 include expr_dsl.hpp）： +// 默认实参（P = F32）在 expr_dsl.hpp 的定义处。 +// 生命周期 v3 起不再有库内 stamp/adopt 通道：DSL 静态工厂产出的 CPU +// 张量天然属于宿主域，设备域从存储派生。 +template +[[nodiscard]] Result compute(ComputeEngine& eng, const E& e, std::size_t rows, + std::size_t cols, Precision P = Precision::F32); +template +[[nodiscard]] Result compute_reduce(ComputeEngine& eng, const E& e, std::size_t rows, + std::size_t cols, Precision P = Precision::F32); +} // namespace dsl // ══════════════════════════════════════════════════════════════════════════ // 算子枚举（op-level，不含算法语义） @@ -104,14 +102,13 @@ enum class ReduceOp : uint32_t // bit16 = 输出是 f16。全 0 = 全 f32：不追加变体后缀，与基础结构共用同一条目 // （融合算术保持 f32）。运行时用它选 (结构 key, 签名) 对应的带类型 shader； // 构建期 scan/gen 两端同源。 -[[nodiscard]] inline ExprPrecSig expr_prec_sig_of( - std::span inputs, Precision P = Precision::F32) noexcept +[[nodiscard]] inline ExprPrecSig expr_prec_sig_of(std::span inputs, + Precision P = Precision::F32) noexcept { std::uint32_t bits = 0; const std::size_t n = inputs.size() < 32u ? inputs.size() : 32u; for (std::size_t i = 0; i < n; ++i) - if (inputs[i].precision() == Precision::F16) - bits |= (1u << i); + if (inputs[i].precision() == Precision::F16) bits |= (1u << i); return expr_prec_sig_make(bits, P == Precision::F16); } @@ -253,15 +250,12 @@ class ComputeEngine : public std::enable_shared_from_this // 取得本引擎的强引用。未由 shared_ptr 管理时返回空指针（调用方必须 // 把它当错误处理，不能继续解引用 this 之外的任何生命周期假设）。 [[nodiscard]] std::shared_ptr keepalive() noexcept - { - return weak_from_this().lock(); - } + { return weak_from_this().lock(); } protected: explicit ComputeEngine(DomainId domain) : domain_(domain) {} public: - // ── 设备查询 ────────────────────────────────────────────────────────── [[nodiscard]] virtual Device device() const noexcept = 0; @@ -317,16 +311,12 @@ class ComputeEngine : public std::enable_shared_from_this [[nodiscard]] virtual std::string pool_stats() const { const auto& led = host_alloc_ledger(); - return "direct{ blocks=" + - std::to_string(led.live_blocks.load(std::memory_order_relaxed)) + - " live_bytes=" + - std::to_string(led.live_bytes.load(std::memory_order_relaxed)) + - " peak_bytes=" + - std::to_string(led.peak_bytes.load(std::memory_order_relaxed)) + - " total_blocks=" + - std::to_string(led.total_blocks.load(std::memory_order_relaxed)) + - " total_bytes=" + - std::to_string(led.total_bytes.load(std::memory_order_relaxed)) + " }"; + return "direct{ blocks=" + std::to_string(led.live_blocks.load(std::memory_order_relaxed)) + + " live_bytes=" + std::to_string(led.live_bytes.load(std::memory_order_relaxed)) + + " peak_bytes=" + std::to_string(led.peak_bytes.load(std::memory_order_relaxed)) + + " total_blocks=" + std::to_string(led.total_blocks.load(std::memory_order_relaxed)) + + " total_bytes=" + std::to_string(led.total_bytes.load(std::memory_order_relaxed)) + + " }"; } // ── 激活 offload（L1-offload）─────────────────────────────────── @@ -343,8 +333,8 @@ class ComputeEngine : public std::enable_shared_from_this // ── 边界 cast 入口：f16 激活在写入 slab 前抬到 f32（slab 恒 f32 存）── // restore 出来是 f32，后续运算由入口层适配（正确性不受影响，代价是该份 // 激活按 f32 存、不享 f16 存储折半）。 - [[nodiscard]] Result offload_save( - const Tensor& buffer, std::size_t offset, const Tensor& src) + [[nodiscard]] Result offload_save(const Tensor& buffer, std::size_t offset, + const Tensor& src) { NN_PROF_OP("offload_save"); NN_TRY(ec, bind_check_({&buffer, &src})); @@ -353,9 +343,8 @@ class ComputeEngine : public std::enable_shared_from_this return offload_save_impl(buffer, offset, *s); } // 从 buffer 的 offset（float 单位）处复制 rows×cols 到新 GPU tensor - [[nodiscard]] Result offload_restore( - const Tensor& buffer, std::size_t offset, - std::size_t rows, std::size_t cols) + [[nodiscard]] Result offload_restore(const Tensor& buffer, std::size_t offset, + std::size_t rows, std::size_t cols) { NN_PROF_OP("offload_restore"); NN_TRY(ec, bind_check_({&buffer})); @@ -383,12 +372,10 @@ class ComputeEngine : public std::enable_shared_from_this return submit_scalar_readback_impl(slot, *s); } - [[nodiscard]] virtual Result poll_scalar_readback( - std::size_t slot, Scalar& out) + [[nodiscard]] virtual Result poll_scalar_readback(std::size_t slot, Scalar& out) { NN_PROF_OP("poll_scalar_readback"); - if (slot >= sync_readback_slots_.size()) - return false; + if (slot >= sync_readback_slots_.size()) return false; out = sync_readback_slots_[slot]; return true; } @@ -419,17 +406,15 @@ class ComputeEngine : public std::enable_shared_from_this } [[nodiscard]] Tensor create_tensor(std::size_t rows, std::size_t cols, Precision P = Precision::F32) - { - return NN_UNWRAP(try_create_tensor(rows, cols, P), "engine.create_tensor"); - } + { return NN_UNWRAP(try_create_tensor(rows, cols, P), "engine.create_tensor"); } // ── 声明式创建/初始化（M2，17 §4.3/§4.4）──────────────────────────── // InitSpec 指定初值；策略 = 引擎 host 生成后上传（D6：两设备同序列、 // 调用方不可见；将来设备端原生生成是引擎内部优化，不改本签名）。 // 失败返回 Result 错误（与 3 参 try_create_tensor 同协议）。 // 语义分工：3 参 = 纯分配（沿用现状：CPU Matrix 零填充 / GPU 分配未初始 // 化，需要零的调用方显式 zero）；"要什么初值"一律走本重载显式声明。 - [[nodiscard]] Result try_create_tensor(std::size_t rows, std::size_t cols, - Precision P, const InitSpec& spec) + [[nodiscard]] Result try_create_tensor(std::size_t rows, std::size_t cols, Precision P, + const InitSpec& spec) { NN_PROF_OP("create_tensor"); NN_TRY(pc, check_precision_supported(P)); @@ -457,10 +442,7 @@ class ComputeEngine : public std::enable_shared_from_this } [[nodiscard]] Tensor create_tensor(std::size_t rows, std::size_t cols, Precision P, const InitSpec& spec) - { - return NN_UNWRAP(try_create_tensor(rows, cols, P, spec), - "engine.create_tensor(spec)"); - } + { return NN_UNWRAP(try_create_tensor(rows, cols, P, spec), "engine.create_tensor(spec)"); } [[nodiscard]] Result from_matrix(const Matrix& m, Precision P = Precision::F32) { NN_PROF_OP("from_matrix"); @@ -468,7 +450,8 @@ class ComputeEngine : public std::enable_shared_from_this NN_TRY(pc, check_precision_supported(P)); return from_matrix_impl(m, P); } - [[nodiscard]] virtual Result to_matrix(const Tensor& t, Precision P = Precision::F32) = 0; + [[nodiscard]] virtual Result to_matrix(const Tensor& t, + Precision P = Precision::F32) = 0; // ── 跨设备拉取（M6，17 §4.3 / 15 §3.2 P3；v3 起是唯一跨域入口）──────── // 目标侧拉取语义（15 §3.2）： @@ -477,32 +460,27 @@ class ComputeEngine : public std::enable_shared_from_this // GPU → CPU **不走这里**（用 to_matrix 落宿主）；GPU 异 backend = 报错。 // 返回**新句柄**，绝不改写调用方的 `src`——需要就地重绑定必须在调用点 // 显式写 `t = std::move(*import(t))`（与 `cast_into` 的不替换红线同族）。 - [[nodiscard]] Result import( - const Tensor& src, - std::source_location loc = std::source_location::current()) + [[nodiscard]] Result import(const Tensor& src, + std::source_location loc = std::source_location::current()) { NN_PROF_OP("import"); return import(src, src.precision(), loc); } - [[nodiscard]] Result import( - const Tensor& src, Precision P, - std::source_location loc = std::source_location::current()) + [[nodiscard]] Result import(const Tensor& src, Precision P, + std::source_location loc = std::source_location::current()) { NN_PROF_OP("import"); // 热路径快路：同一设备域且精度一致时，import_impl 只会原样返回共享 // 句柄。optimizer 每 step 会对每个 DSL leaf 反复走此入口（profile 约 // 3.4 万次/20 steps）；在这里提前返回可跳过下方检查与 virtual dispatch。 // 域从存储派生（只读、无原子操作）。 - if (src.valid() && src.precision() == P && src.device_domain_() == domain_) - return src; - if (!src.valid()) - NN_FAIL("import: invalid tensor"); + if (src.valid() && src.precision() == P && src.device_domain_() == domain_) return src; + if (!src.valid()) NN_FAIL("import: invalid tensor"); const DomainId d = src.device_domain_(); // 唯一允许的跨域方向：宿主域 → 本引擎（GPU 上载）。CPU 引擎收到 GPU // 张量、或 GPU 引擎收到另一个 backend 的张量，一律明确报错。 - if (d != domain_ && d != kHostDomain) - NN_FAIL(domain_error_(src, loc)); + if (d != domain_ && d != kHostDomain) NN_FAIL(domain_error_(src, loc)); return import_impl(src, P); } @@ -535,29 +513,21 @@ class ComputeEngine : public std::enable_shared_from_this // ══════════════════════════════════════════════════════════════════════ // 批量读：张量全量 → dst（行主序），GPU 隐含 flush + 同步。 - template - [[nodiscard]] Result read(const Tensor& t, std::span dst) + template [[nodiscard]] Result read(const Tensor& t, std::span dst) { NN_PROF_OP("read"); - static_assert(!std::is_const_v && - (std::is_same_v || std::is_same_v), + static_assert(!std::is_const_v && (std::is_same_v || std::is_same_v), "read: span 元素类型必须是 nn::Scalar(float) 或 nn::f16（U2 精确匹配）"); NN_TRY(ec, bind_check_({&t})); - constexpr Precision want = - std::is_same_v ? Precision::F16 : Precision::F32; - if (t.precision() != want) - NN_FAIL("read: span 元素类型与张量精度不匹配（U2 精确匹配）"); - if (dst.size() != t.rows() * t.cols()) - NN_FAIL("read: 元素数与张量形状不匹配"); + constexpr Precision want = std::is_same_v ? Precision::F16 : Precision::F32; + if (t.precision() != want) NN_FAIL("read: span 元素类型与张量精度不匹配（U2 精确匹配）"); + if (dst.size() != t.rows() * t.cols()) NN_FAIL("read: 元素数与张量形状不匹配"); // 统一经 to_matrix（GPU 分支自带 flush + wait_in_flight；f16 存储 // 升 f32 精确无损，T=f16 时下面再 RHE 落回——往返位不变）。 auto m = to_matrix(t, Precision::F32); NN_TRY_CHECK(m); const auto src = m->span(); - if constexpr (std::is_same_v) - { - std::copy(src.begin(), src.end(), dst.begin()); - } + if constexpr (std::is_same_v) { std::copy(src.begin(), src.end(), dst.begin()); } else { for (std::size_t i = 0; i < src.size(); ++i) @@ -637,8 +607,7 @@ class ComputeEngine : public std::enable_shared_from_this // （BitLinear 的三值权重）靠写版本判断"要不要重算"，见下方说明。 // fresh_dst 语义见 copy_from（默认 false；仅 detail::upload_span 传 true）。 template - [[nodiscard]] Result write(Tensor& t, std::span src, - bool fresh_dst = false) + [[nodiscard]] Result write(Tensor& t, std::span src, bool fresh_dst = false) { NN_PROF_OP("write"); using Elem = std::remove_cv_t; @@ -646,12 +615,9 @@ class ComputeEngine : public std::enable_shared_from_this "write: span 元素类型必须是 nn::Scalar(float) 或 nn::f16（U2 精确匹配）"); NN_TRY(ec, bind_check_({&t})); WriteEpochGuard wg{t}; // 写版本（见 guard 注释） - constexpr Precision want = - std::is_same_v ? Precision::F16 : Precision::F32; - if (t.precision() != want) - NN_FAIL("write: span 元素类型与张量精度不匹配（U2 精确匹配）"); - if (src.size() != t.rows() * t.cols()) - NN_FAIL("write: 元素数与张量形状不匹配"); + constexpr Precision want = std::is_same_v ? Precision::F16 : Precision::F32; + if (t.precision() != want) NN_FAIL("write: span 元素类型与张量精度不匹配（U2 精确匹配）"); + if (src.size() != t.rows() * t.cols()) NN_FAIL("write: 元素数与张量形状不匹配"); if (t.is_cpu()) { // 宿主张量直写既有存储（保留对象身份）。不走 copy_from：GPU 引擎 @@ -694,12 +660,10 @@ class ComputeEngine : public std::enable_shared_from_this // 索引读（语法糖：宿主直读 / GPU = 一次批量 read）。行主序 (row, col)； // f16 存储提升为 Scalar（精确无损）。GPU 每次调用一整轮 staging 往返， // 不承诺热循环性能（17 §7-1）——热路径请攒批量用 read。 - [[nodiscard]] Result get_index(const Tensor& t, std::size_t row, - std::size_t col) + [[nodiscard]] Result get_index(const Tensor& t, std::size_t row, std::size_t col) { NN_TRY(ec, bind_check_({&t})); - if (row >= t.rows() || col >= t.cols()) - NN_FAIL("get_index: (row, col) 越界"); + if (row >= t.rows() || col >= t.cols()) NN_FAIL("get_index: (row, col) 越界"); const std::size_t idx = row * t.cols() + col; if (t.is_cpu()) { @@ -720,20 +684,17 @@ class ComputeEngine : public std::enable_shared_from_this // 索引写（语法糖：宿主直写 / GPU = read → 改一格 → write）。 // f16 存储：Scalar 值经 RHE 舍入（与 from_matrix/边界 cast 同口径）。 - [[nodiscard]] Result set_index(Tensor& t, std::size_t row, std::size_t col, - Scalar v) + [[nodiscard]] Result set_index(Tensor& t, std::size_t row, std::size_t col, Scalar v) { NN_TRY(ec, bind_check_({&t})); WriteEpochGuard wg{t}; - if (row >= t.rows() || col >= t.cols()) - NN_FAIL("set_index: (row, col) 越界"); + if (row >= t.rows() || col >= t.cols()) NN_FAIL("set_index: (row, col) 越界"); const std::size_t idx = row * t.cols() + col; if (t.is_cpu()) { if (t.precision() == Precision::F16) t.cpu_matrix().span()[idx] = static_cast(v); - else - t.cpu_matrix().span()[idx] = v; + else t.cpu_matrix().span()[idx] = v; return {}; } if (t.precision() == Precision::F16) @@ -754,12 +715,11 @@ class ComputeEngine : public std::enable_shared_from_this // precision_ 复制数据到新形状；设备域随存储派生（v3），空槽保持无效 // 的输入形状由引擎校验。元素数不匹配返回错误（原 NN_ASSERT 升级）。 [[nodiscard]] Result reshape(const Tensor& t, std::size_t new_rows, - std::size_t new_cols) + std::size_t new_cols) { NN_PROF_OP("reshape"); NN_TRY(ec, bind_check_({&t})); - if (t.rows() * t.cols() != new_rows * new_cols) - NN_FAIL("reshape: element count mismatch"); + if (t.rows() * t.cols() != new_rows * new_cols) NN_FAIL("reshape: element count mismatch"); return t.reshape(new_rows, new_cols); } @@ -782,8 +742,7 @@ class ComputeEngine : public std::enable_shared_from_this { if (src.rows() != dst.rows() || src.cols() != dst.cols()) NN_FAIL("cast_into: shape mismatch"); - if (src.precision() == dst.precision()) - return copy_into(dst, src); + if (src.precision() == dst.precision()) return copy_into(dst, src); NN_FAIL("cast_into: 该引擎不支持跨精度转换"); } @@ -791,7 +750,8 @@ class ComputeEngine : public std::enable_shared_from_this // 用于 cast_into 的同精度分支与边界 cast 的"写回原存储"路径。 [[nodiscard]] virtual Result copy_into(Tensor& dst, const Tensor& src) { - (void)dst; (void)src; + (void)dst; + (void)src; NN_FAIL("copy_into: 该引擎未实现"); } @@ -837,21 +797,19 @@ class ComputeEngine : public std::enable_shared_from_this [[nodiscard]] static std::string dump_temp_stats() { auto& v = temp_stats(); - if (v.empty()) - return {}; + if (v.empty()) return {}; std::vector order; order.reserve(v.size()); - for (const auto& e : v) order.push_back(&e); + for (const auto& e : v) + order.push_back(&e); std::sort(order.begin(), order.end(), [](const TempStat* a, const TempStat* b) { return a->bytes > b->bytes; }); std::string s = "── 边界 cast 临时量归因（按字节降序）──\n"; for (const auto* e : order) { char buf[256]; - std::snprintf(buf, sizeof(buf), - " %-11s (%zu,%zu) x%-5zu %8.1f MB L%u\n", - e->to_f32 ? "->f32 副本" : "->按P落回", - e->rows, e->cols, e->count, + std::snprintf(buf, sizeof(buf), " %-11s (%zu,%zu) x%-5zu %8.1f MB L%u\n", + e->to_f32 ? "->f32 副本" : "->按P落回", e->rows, e->cols, e->count, static_cast(e->bytes) / (1024.0 * 1024.0), static_cast(e->line)); s += buf; @@ -894,14 +852,12 @@ class ComputeEngine : public std::enable_shared_from_this // 零回归。 // ⚠ 契约：fresh_dst=true 仅当 dst 未被任何**已录制或在飞**命令引用； // 误传 = 数据竞争。库内唯一承诺方是 detail::upload_span。 - [[nodiscard]] Result copy_from(Tensor& dst, const Matrix& src, - bool fresh_dst = false) + [[nodiscard]] Result copy_from(Tensor& dst, const Matrix& src, bool fresh_dst = false) { NN_PROF_OP("copy_from"); NN_TRY(ec, bind_check_({&dst})); WriteEpochGuard wg{dst}; - if (dst.precision() == Precision::F32) - return copy_from_impl(dst, src, fresh_dst); + if (dst.precision() == Precision::F32) return copy_from_impl(dst, src, fresh_dst); auto t = from_matrix(src, Precision::F32); NN_TRY_CHECK(t); return cast_into(*t, dst); @@ -910,8 +866,8 @@ class ComputeEngine : public std::enable_shared_from_this // ── 行切片原语（op-level 数据操作，不含算法语义） ────────────────── // 返回 src 的行 [start_row, start_row + count) 的连续拷贝。 // 用于多头注意力中 per-head Q/K/V 切片等场景。 - [[nodiscard]] Result slice_rows( - const Tensor& src, std::size_t start_row, std::size_t count) + [[nodiscard]] Result slice_rows(const Tensor& src, std::size_t start_row, + std::size_t count) { NN_PROF_OP("slice_rows"); NN_TRY(ec, bind_check_({&src})); @@ -922,8 +878,8 @@ class ComputeEngine : public std::enable_shared_from_this // 真·就地修改（GPU 用 vkCmdCopyBuffer with dstOffset）。 // 用于多头注意力中 per-head 输出拼接等场景。 // insert_rows：就地写入 dst 的行区间 → dst 精度即目标精度 - [[nodiscard]] Result insert_rows( - Tensor& dst, std::size_t dst_start_row, const Tensor& src) + [[nodiscard]] Result insert_rows(Tensor& dst, std::size_t dst_start_row, + const Tensor& src) { NN_PROF_OP("insert_rows"); NN_TRY(ec, bind_check_({&dst, &src})); @@ -937,8 +893,7 @@ class ComputeEngine : public std::enable_shared_from_this // indices: (num_indices,) — 行索引；越界索引返回零行（防御性，不抛错） // 输出: (num_indices, D)，out[i] = table[indices[i]] // 典型用途：Token embedding 查表（避免 Layer 内手动 to_matrix + at_unchecked） - [[nodiscard]] Result gather_rows( - const Tensor& table, const Tensor& indices) + [[nodiscard]] Result gather_rows(const Tensor& table, const Tensor& indices) { NN_PROF_OP("gather_rows"); NN_TRY(ec, bind_check_({&table, &indices})); @@ -952,8 +907,8 @@ class ComputeEngine : public std::enable_shared_from_this // 语义: dst[indices[i]] += grad[i] (重复 indices 会被多次累加) // 典型用途：Embedding 反向梯度按 token ID 累加（替代 Layer 内手动循环） // scatter_add_rows：就地累加到 dst → 在 f32 空间累加后写回 dst 存储 - [[nodiscard]] Result scatter_add_rows( - Tensor& dst, const Tensor& indices, const Tensor& grad) + [[nodiscard]] Result scatter_add_rows(Tensor& dst, const Tensor& indices, + const Tensor& grad) { NN_PROF_OP("scatter_add_rows"); NN_TRY(ec, bind_check_({&dst, &indices, &grad})); @@ -968,9 +923,8 @@ class ComputeEngine : public std::enable_shared_from_this // out[m, b*N + n] = in[b*M + m, n] // 典型用途：MHA 批量化时把 (H*d_k, batch*seq) 重排为 (batch*H*d_k, seq)， // 使 batched_matmul 能按 batch*H 切分行块。 - [[nodiscard]] Result rearrange_3d( - const Tensor& x, std::size_t M, std::size_t B, std::size_t N, - bool inverse = false) + [[nodiscard]] Result rearrange_3d(const Tensor& x, std::size_t M, std::size_t B, + std::size_t N, bool inverse = false) { NN_PROF_OP("rearrange_3d"); NN_TRY(ec, bind_check_({&x})); @@ -997,11 +951,9 @@ class ComputeEngine : public std::enable_shared_from_this // 各一次即可完成 (C, B*P) → (C*P, B) 的 samples 布局转换。 // 用途：Conv2D 把卷积变成 matmul（W × col）；MaxPool2D 展开窗口做列归约 // （取 k=pool, stride=stride, pad=0）。 - [[nodiscard]] Result im2col( - const Tensor& x, - std::size_t C, std::size_t H, std::size_t W, - std::size_t k, std::size_t stride, std::size_t pad, - std::size_t OH, std::size_t OW) + [[nodiscard]] Result im2col(const Tensor& x, std::size_t C, std::size_t H, + std::size_t W, std::size_t k, std::size_t stride, + std::size_t pad, std::size_t OH, std::size_t OW) { NN_PROF_OP("im2col"); NN_TRY(ec, bind_check_({&x})); @@ -1015,11 +967,9 @@ class ComputeEngine : public std::enable_shared_from_this // 其中 oh = (ih + pad - kh)/stride（须整除且落在 [0, OH)），ow 同理。 // **重叠窗口（stride < k）的贡献在此累加**；每个输出元素由单个线程/循环 // 完整求和，故无需原子操作/预清零。 - [[nodiscard]] Result col2im( - const Tensor& col, - std::size_t C, std::size_t H, std::size_t W, - std::size_t k, std::size_t stride, std::size_t pad, - std::size_t OH, std::size_t OW) + [[nodiscard]] Result col2im(const Tensor& col, std::size_t C, std::size_t H, + std::size_t W, std::size_t k, std::size_t stride, + std::size_t pad, std::size_t OH, std::size_t OW) { NN_PROF_OP("col2im"); NN_TRY(ec, bind_check_({&col})); @@ -1038,19 +988,14 @@ class ComputeEngine : public std::enable_shared_from_this // （f16 直读 + f32 累加 + f16 写出）→ 不为每个操作数物化整份 f32 副本。 // 引擎在"f16 pipeline 不可用 / 小 N GEMV / 操作数精度混合"时自行回退 // f32 空间 + 边界，正确性不变。 - [[nodiscard]] Result matmul( - const Tensor& A, const Tensor& B, - bool transA = false, bool transB = false, - Precision P = Precision::F32) + [[nodiscard]] Result matmul(const Tensor& A, const Tensor& B, bool transA = false, + bool transB = false, Precision P = Precision::F32) { NN_PROF_OP("matmul"); NN_TRY(ec, bind_check_({&A, &B})); - if (P != Precision::F32) - return matmul_impl(A, B, transA, transB, P); + if (P != Precision::F32) return matmul_impl(A, B, transA, transB, P); return binary_(A, B, P, [this, transA, transB](const Tensor& a, const Tensor& b) - { - return matmul_impl(a, b, transA, transB, Precision::F32); - }); + { return matmul_impl(a, b, transA, transB, Precision::F32); }); } // 批量矩阵乘法：对每个 batch b 计算 C_b = alpha * op(A_b, B_b)，结果垂直堆叠 @@ -1063,23 +1008,17 @@ class ComputeEngine : public std::enable_shared_from_this // 供上层折叠 1/sqrt(d_k) 等系数，省去额外全矩阵 scale pass // P: 计算精度（§8.1，同 matmul） // 典型用途：多头注意力的 Q^T×K 和 V×A 批量化（消除 per-head 循环） - [[nodiscard]] Result batched_matmul( - const Tensor& A, const Tensor& B, - std::size_t batch, - bool transA = false, bool transB = false, - Scalar alpha = Scalar{1}, - Precision P = Precision::F32) + [[nodiscard]] Result batched_matmul(const Tensor& A, const Tensor& B, std::size_t batch, + bool transA = false, bool transB = false, + Scalar alpha = Scalar{1}, + Precision P = Precision::F32) { NN_PROF_OP("batched_matmul"); NN_TRY(ec, bind_check_({&A, &B})); - if (P != Precision::F32) - return batched_matmul_impl(A, B, batch, transA, transB, alpha, P); - return binary_(A, B, P, - [this, batch, transA, transB, alpha](const Tensor& a, const Tensor& b) - { - return batched_matmul_impl(a, b, batch, transA, transB, alpha, - Precision::F32); - }); + if (P != Precision::F32) return batched_matmul_impl(A, B, batch, transA, transB, alpha, P); + return binary_( + A, B, P, [this, batch, transA, transB, alpha](const Tensor& a, const Tensor& b) + { return batched_matmul_impl(a, b, batch, transA, transB, alpha, Precision::F32); }); } // ── matmul + broadcast bias（统一精度处理）──────────────────────────── @@ -1087,10 +1026,9 @@ class ComputeEngine : public std::enable_shared_from_this // 入口按 P/操作数精度分派：非 f32 直接下传（经 DSL 融合 matmul 段，带类型 // 变体可用时 f16 直读直写、零边界 cast；无变体时引擎自行回退 f32 空间）； // 全 f32 直通；混合输入抬 f32 计算后按 P 落回。 - [[nodiscard]] Result matmul_with_bias( - const Tensor& A, const Tensor& B, const Tensor& bias, - bool transA = false, bool transB = false, - Precision P = Precision::F32) + [[nodiscard]] Result matmul_with_bias(const Tensor& A, const Tensor& B, + const Tensor& bias, bool transA = false, + bool transB = false, Precision P = Precision::F32) { NN_PROF_OP("matmul_with_bias"); NN_TRY(ec, bind_check_({&A, &B, &bias})); @@ -1101,16 +1039,14 @@ class ComputeEngine : public std::enable_shared_from_this if (P == Precision::F32 && A.precision() == Precision::F32 && B.precision() == Precision::F32 && bias.precision() == Precision::F32) - return matmul_with_bias_impl(A, B, bias, transA, transB, - Precision::F32); + return matmul_with_bias_impl(A, B, bias, transA, transB, Precision::F32); auto a = to_f32(A); NN_TRY_CHECK(a); auto b = to_f32(B); NN_TRY_CHECK(b); auto bi = to_f32(bias); NN_TRY_CHECK(bi); - auto r = matmul_with_bias_impl(*a, *b, *bi, transA, transB, - Precision::F32); + auto r = matmul_with_bias_impl(*a, *b, *bi, transA, transB, Precision::F32); NN_TRY_CHECK(r); return to_prec(std::move(*r), P); } @@ -1203,24 +1139,17 @@ class ComputeEngine : public std::enable_shared_from_this } [[nodiscard]] Result matmul_q(const Tensor& packed, const Tensor& scale, - const Tensor& B, - Precision P = Precision::F32) - { - return matmul_q_dispatch_(packed, scale, B, /*transA=*/false, /*out_rows=*/0, - P); - } + const Tensor& B, Precision P = Precision::F32) + { return matmul_q_dispatch_(packed, scale, B, /*transA=*/false, /*out_rows=*/0, P); } // C = (τ∘W)^T·B：A 存储为 (K, M)（= W 的 (out, in)），out_rows = M。 [[nodiscard]] Result matmul_q_t(const Tensor& packed, const Tensor& scale, const Tensor& B, std::size_t out_rows, Precision P = Precision::F32) - { - return matmul_q_dispatch_(packed, scale, B, /*transA=*/true, out_rows, P); - } + { return matmul_q_dispatch_(packed, scale, B, /*transA=*/true, out_rows, P); } private: - [[nodiscard]] Result matmul_q_dispatch_(const Tensor& packed, - const Tensor& scale, + [[nodiscard]] Result matmul_q_dispatch_(const Tensor& packed, const Tensor& scale, const Tensor& B, bool transA, std::size_t out_rows, Precision P) { @@ -1232,7 +1161,6 @@ class ComputeEngine : public std::enable_shared_from_this } public: - // ══════════════════════════════════════════════════════════════════════ // 扫描级原语（带状态的顺序归约 + matvec 读出；RLA 线性注意力积木） // @@ -1269,18 +1197,17 @@ class ComputeEngine : public std::enable_shared_from_this // 输出 (B·H·5·d_k, seq)，行块（每块 (B·H·d_k, seq)）： // [0) B·P [1) A·P [2) B^T·R [3) s = P·(A·P) [4) r = R·(B·P) // 其中 [3)/[4) 为逐列标量（头内逐行重复）。 - [[nodiscard]] Result scan_prefix_outer( - const Tensor& K, const Tensor& V, const Tensor& P, const Tensor& R, - const Tensor& A0, const Tensor& B0, bool has_state, - std::size_t dk, std::size_t heads, bool causal, - const Tensor& boundary, bool has_bnd, - Precision prec = Precision::F32) + [[nodiscard]] Result + scan_prefix_outer(const Tensor& K, const Tensor& V, const Tensor& P, const Tensor& R, + const Tensor& A0, const Tensor& B0, bool has_state, std::size_t dk, + std::size_t heads, bool causal, const Tensor& boundary, bool has_bnd, + Precision prec = Precision::F32) { NN_PROF_OP("scan_prefix_outer"); NN_TRY(ec, bind_check_({&K, &V, &P, &R, &A0, &B0, &boundary})); // 精度策略（原生 f16 扫描 / 引擎内边界 cast）在引擎的 *_impl 内。 - return scan_prefix_outer_impl(K, V, P, R, A0, B0, has_state, dk, heads, - causal, boundary, has_bnd, prec); + return scan_prefix_outer_impl(K, V, P, R, A0, B0, has_state, dk, heads, causal, boundary, + has_bnd, prec); } // 后缀扫描（RLA 反向 pass 2）： @@ -1289,16 +1216,15 @@ class ComputeEngine : public std::enable_shared_from_this // causal=false: S_i = D_i（Layer 预先把全集梯度沿 seq 广播）。 // D (B·H·d_k², seq)，X/Y (B·H·d_k, seq) // 输出 (B·H·3·d_k, seq)，行块：[0) S·X [1) S·Y [2) S^T·Y - [[nodiscard]] Result scan_suffix_outer( - const Tensor& D, const Tensor& X, const Tensor& Y, - std::size_t dk, std::size_t heads, bool causal, - const Tensor& boundary, bool has_bnd, - Precision prec = Precision::F32) + [[nodiscard]] Result scan_suffix_outer(const Tensor& D, const Tensor& X, + const Tensor& Y, std::size_t dk, + std::size_t heads, bool causal, + const Tensor& boundary, bool has_bnd, + Precision prec = Precision::F32) { NN_PROF_OP("scan_suffix_outer"); NN_TRY(ec, bind_check_({&D, &X, &Y, &boundary})); - return scan_suffix_outer_impl(D, X, Y, dk, heads, causal, boundary, - has_bnd, prec); + return scan_suffix_outer_impl(D, X, Y, dk, heads, causal, boundary, has_bnd, prec); } // 逐列外积（RLA 反向的 dL/dA、dL/dB 物化）： @@ -1306,10 +1232,9 @@ class ComputeEngine : public std::enable_shared_from_this // P/R: (B·H·d_k, seq)；S: (B·H·d_k, seq)（标量头内逐行重复， // 实现读头块首行；has_scale=false → 传 dummy） // 输出: (B·H·d_k², seq) - [[nodiscard]] Result outer_col( - const Tensor& P, const Tensor& R, const Tensor& S, - std::size_t dk, bool has_scale, - Precision prec = Precision::F32) + [[nodiscard]] Result outer_col(const Tensor& P, const Tensor& R, const Tensor& S, + std::size_t dk, bool has_scale, + Precision prec = Precision::F32) { NN_PROF_OP("outer_col"); NN_TRY(ec, bind_check_({&P, &R, &S})); @@ -1327,8 +1252,7 @@ class ComputeEngine : public std::enable_shared_from_this // 按行求和：A (rows, cols) → out (rows, 1) // out[r] = Σ_c A[r][c] // P：输出（与累加）精度（§7.3：归约恒 f32 累加 + 输出舍入到 P） - [[nodiscard]] Result row_reduce_sum( - const Tensor& A, Precision P = Precision::F32) + [[nodiscard]] Result row_reduce_sum(const Tensor& A, Precision P = Precision::F32) { NN_PROF_OP("row_reduce_sum"); NN_TRY(ec, bind_check_({&A})); @@ -1337,8 +1261,7 @@ class ComputeEngine : public std::enable_shared_from_this // 按列求和：A (rows, cols) → out (1, cols) // out[c] = Σ_r A[r][c] - [[nodiscard]] Result col_reduce_sum( - const Tensor& A, Precision P = Precision::F32) + [[nodiscard]] Result col_reduce_sum(const Tensor& A, Precision P = Precision::F32) { NN_PROF_OP("col_reduce_sum"); NN_TRY(ec, bind_check_({&A})); @@ -1347,8 +1270,7 @@ class ComputeEngine : public std::enable_shared_from_this // 按列求最大值：A (rows, cols) → out (1, cols) // out[c] = max_r A[r][c] - [[nodiscard]] Result col_reduce_max( - const Tensor& A, Precision P = Precision::F32) + [[nodiscard]] Result col_reduce_max(const Tensor& A, Precision P = Precision::F32) { NN_PROF_OP("col_reduce_max"); NN_TRY(ec, bind_check_({&A})); @@ -1362,17 +1284,15 @@ class ComputeEngine : public std::enable_shared_from_this // grouped_reduce_max: out[g, n] = max_i x[g*R + i, n] // 用途：把"逐通道/逐头一次 dispatch"的层内循环压成**单次**原语调用 // （池化窗口归约、多头分组统计、分组归一化等）。 - [[nodiscard]] Result grouped_reduce_sum( - const Tensor& x, std::size_t G, std::size_t R, - Precision P = Precision::F32) + [[nodiscard]] Result grouped_reduce_sum(const Tensor& x, std::size_t G, std::size_t R, + Precision P = Precision::F32) { NN_PROF_OP("grouped_reduce_sum"); NN_TRY(ec, bind_check_({&x})); return grouped_reduce_sum_impl(x, G, R, P); } - [[nodiscard]] Result grouped_reduce_max( - const Tensor& x, std::size_t G, std::size_t R, - Precision P = Precision::F32) + [[nodiscard]] Result grouped_reduce_max(const Tensor& x, std::size_t G, std::size_t R, + Precision P = Precision::F32) { NN_PROF_OP("grouped_reduce_max"); NN_TRY(ec, bind_check_({&x})); @@ -1394,11 +1314,9 @@ class ComputeEngine : public std::enable_shared_from_this // in-kernel f16 优先：引擎有该 (结构, 输入精度, 输出精度) 的**带类型** // 融合 shader → 直接把原张量交给它（读 f16 / 写 f16，算术 f32）； // 无变体 → 抬 f32 计算后按 P 落回（trace_miss_ 是唯一可见信号）。 - [[nodiscard]] Result eval_expr( - const ExprSpec& spec, - std::span inputs, - std::size_t rows, std::size_t cols, - Precision P = Precision::F32) + [[nodiscard]] Result eval_expr(const ExprSpec& spec, std::span inputs, + std::size_t rows, std::size_t cols, + Precision P = Precision::F32) { NN_PROF_OP("eval_expr"); NN_TRY(ec, bind_check_(inputs)); @@ -1417,11 +1335,9 @@ class ComputeEngine : public std::enable_shared_from_this // 语义同 eval_expr，但输出为归约向量本身（非广播）： // 行归约轴 → (rows,1)；列归约轴 → (1,cols)。 // 要求表达式归约轴为 0/1（expr_spec_reduce_axis）。 - [[nodiscard]] Result eval_expr_reduce( - const ExprSpec& spec, - std::span inputs, - std::size_t rows, std::size_t cols, - Precision P = Precision::F32) + [[nodiscard]] Result eval_expr_reduce(const ExprSpec& spec, + std::span inputs, std::size_t rows, + std::size_t cols, Precision P = Precision::F32) { NN_PROF_OP("eval_expr_reduce"); NN_TRY(ec, bind_check_(inputs)); @@ -1441,10 +1357,8 @@ class ComputeEngine : public std::enable_shared_from_this // 就地按行广播等）。out 允许与某个输入是同一 buffer：逐元素"先读完全部 // 输入再写 out[i]"，就地安全。 // 仅支持逐元素表达式（无归约）；归约向量输出用 eval_expr_reduce。 - [[nodiscard]] Result eval_expr_into( - const ExprSpec& spec, - std::span inputs, - std::size_t rows, std::size_t cols, Tensor& out) + [[nodiscard]] Result eval_expr_into(const ExprSpec& spec, std::span inputs, + std::size_t rows, std::size_t cols, Tensor& out) { NN_PROF_OP("eval_expr_into"); NN_TRY(ec, bind_check_(inputs)); @@ -1467,9 +1381,10 @@ class ComputeEngine : public std::enable_shared_from_this // 对带 extras 的结构只发 V0 变体（`#x`/typed/native16 跳过），且多输出 // 融合由 fuse_expr_graph 的精度门禁限定 F32。不支持录制段（图执行是 // end_expr 的职责）。 - [[nodiscard]] Result eval_expr_multi_into( - const ExprSpec& spec, std::span inputs, - std::size_t rows, std::size_t cols, std::span outs) + [[nodiscard]] Result eval_expr_multi_into(const ExprSpec& spec, + std::span inputs, + std::size_t rows, std::size_t cols, + std::span outs) { NN_PROF_OP("eval_expr_multi_into"); NN_TRY(ec, bind_check_(inputs)); @@ -1497,7 +1412,8 @@ class ComputeEngine : public std::enable_shared_from_this // ── P1 NVI 的引擎侧实现（原公共虚入口下沉至此，stamp 在公共入口统一做）── // RV-02：实现侧返回 Result —— 分配失败（GPU 显存/设备）的原因必须能被 // 上层拿到，旧形态 `return Tensor();` 把原因整块丢弃。 - [[nodiscard]] virtual Result create_tensor_impl(std::size_t rows, std::size_t cols, Precision P) = 0; + [[nodiscard]] virtual Result create_tensor_impl(std::size_t rows, std::size_t cols, + Precision P) = 0; [[nodiscard]] virtual Result from_matrix_impl(const Matrix& m, Precision P) = 0; // 跨设备拉取（M6；v3 起是唯一 CPU→GPU 跨域入口，公共入口做设备域检查）。 // 默认实现 = 经宿主中转；同设备快路（别名 / 引擎内 cast）也在默认实现内。 @@ -1506,8 +1422,7 @@ class ComputeEngine : public std::enable_shared_from_this { if (src.device() == device()) { - if (P == src.precision()) - return src; // 同设备同精度 = 零拷贝别名 + if (P == src.precision()) return src; // 同设备同精度 = 零拷贝别名 return cast(src, P); // 同设备异精度 = 引擎内 cast } auto m = to_matrix(src, Precision::F32); // 跨设备：宿主中转 @@ -1517,38 +1432,31 @@ class ComputeEngine : public std::enable_shared_from_this // cast 默认实现：同精度 = 返回 src（共享所有权，零拷贝）；跨精度 = 错误（引擎覆盖） [[nodiscard]] virtual Result cast_impl(const Tensor& src, Precision dst) { - if (src.precision() == dst) - return src; // 同精度 = 返回共享所有权（零拷贝） + if (src.precision() == dst) return src; // 同精度 = 返回共享所有权（零拷贝） NN_FAIL("cast: 该引擎不支持跨精度转换"); } // offload 占位默认（激活 offload 是 GPU 特性；CPU 开启只得到 1×1 张量） // D8（15 §4.8）：走 `create_tensor` 而非 `Tensor::cpu` 静态工厂——出生即经 // 引擎入口，形状/语义不变（仍是 1×1 占位）。 [[nodiscard]] virtual Result create_offload_buffer_impl(std::size_t /*bytes*/) - { - return create_tensor(1, 1); - } - [[nodiscard]] virtual Result offload_restore_impl( - const Tensor& /*buffer*/, std::size_t /*offset*/, - std::size_t /*rows*/, std::size_t /*cols*/) - { - return create_tensor(1, 1); - } + { return create_tensor(1, 1); } + [[nodiscard]] virtual Result offload_restore_impl(const Tensor& /*buffer*/, + std::size_t /*offset*/, + std::size_t /*rows*/, + std::size_t /*cols*/) + { return create_tensor(1, 1); } // ── offload / 回读 ──────────────────────────────────────────────────── // CPU no-op 默认（激活 offload 是 GPU 特性；CPU 开启只得到 1×1 张量） - [[nodiscard]] virtual Result offload_save_impl( - const Tensor& /*buffer*/, std::size_t /*offset*/, const Tensor& /*src*/) - { - return {}; - } + [[nodiscard]] virtual Result + offload_save_impl(const Tensor& /*buffer*/, std::size_t /*offset*/, const Tensor& /*src*/) + { return {}; } // 默认实现（CPU / 其他同步引擎）：submit 立刻取标量存值，poll 恒立刻就绪。 // （入口已统一抬到 f32，这里只读 f32；f16 分支保留作健壮性兜底。） - [[nodiscard]] virtual Result submit_scalar_readback_impl( - std::size_t slot, const Tensor& t) + [[nodiscard]] virtual Result submit_scalar_readback_impl(std::size_t slot, + const Tensor& t) { - if (!t.is_cpu()) - NN_FAIL("submit_scalar_readback: 默认实现仅支持 CPU 张量"); + if (!t.is_cpu()) NN_FAIL("submit_scalar_readback: 默认实现仅支持 CPU 张量"); Scalar v = Scalar{0}; if (t.precision() == Precision::F32) { @@ -1564,8 +1472,7 @@ class ComputeEngine : public std::enable_shared_from_this { NN_FAIL("submit_scalar_readback: 不支持的精度"); } - if (sync_readback_slots_.size() <= slot) - sync_readback_slots_.resize(slot + 1, Scalar{0}); + if (sync_readback_slots_.size() <= slot) sync_readback_slots_.resize(slot + 1, Scalar{0}); sync_readback_slots_[slot] = v; return {}; } @@ -1574,136 +1481,139 @@ class ComputeEngine : public std::enable_shared_from_this [[nodiscard]] virtual Result clone_impl(const Tensor& src) = 0; [[nodiscard]] virtual Result copy_from_impl(Tensor& dst, const Matrix& src, bool fresh_dst) = 0; - [[nodiscard]] virtual Result slice_rows_impl( - const Tensor& src, std::size_t start_row, std::size_t count) = 0; - [[nodiscard]] virtual Result insert_rows_impl( - Tensor& dst, std::size_t dst_start_row, const Tensor& src) = 0; - [[nodiscard]] virtual Result gather_rows_impl( - const Tensor& table, const Tensor& indices) = 0; - [[nodiscard]] virtual Result scatter_add_rows_impl( - Tensor& dst, const Tensor& indices, const Tensor& grad) = 0; - [[nodiscard]] virtual Result rearrange_3d_impl( - const Tensor& x, std::size_t M, std::size_t B, std::size_t N, - bool inverse) = 0; + [[nodiscard]] virtual Result slice_rows_impl(const Tensor& src, std::size_t start_row, + std::size_t count) = 0; + [[nodiscard]] virtual Result insert_rows_impl(Tensor& dst, std::size_t dst_start_row, + const Tensor& src) = 0; + [[nodiscard]] virtual Result gather_rows_impl(const Tensor& table, + const Tensor& indices) = 0; + [[nodiscard]] virtual Result scatter_add_rows_impl(Tensor& dst, const Tensor& indices, + const Tensor& grad) = 0; + [[nodiscard]] virtual Result rearrange_3d_impl(const Tensor& x, std::size_t M, + std::size_t B, std::size_t N, + bool inverse) = 0; [[nodiscard]] virtual Result transpose_impl(const Tensor& A) = 0; - [[nodiscard]] virtual Result im2col_impl( - const Tensor& x, - std::size_t C, std::size_t H, std::size_t W, - std::size_t k, std::size_t stride, std::size_t pad, - std::size_t OH, std::size_t OW) = 0; - [[nodiscard]] virtual Result col2im_impl( - const Tensor& col, - std::size_t C, std::size_t H, std::size_t W, - std::size_t k, std::size_t stride, std::size_t pad, - std::size_t OH, std::size_t OW) = 0; + [[nodiscard]] virtual Result im2col_impl(const Tensor& x, std::size_t C, std::size_t H, + std::size_t W, std::size_t k, + std::size_t stride, std::size_t pad, + std::size_t OH, std::size_t OW) = 0; + [[nodiscard]] virtual Result col2im_impl(const Tensor& col, std::size_t C, + std::size_t H, std::size_t W, std::size_t k, + std::size_t stride, std::size_t pad, + std::size_t OH, std::size_t OW) = 0; // ── 矩阵级 ──────────────────────────────────────────────────────────── - [[nodiscard]] virtual Result matmul_impl( - const Tensor& A, const Tensor& B, - bool transA, bool transB, Precision P) = 0; + [[nodiscard]] virtual Result matmul_impl(const Tensor& A, const Tensor& B, bool transA, + bool transB, Precision P) = 0; // 三值打包权重（P2 / F3）：语义见公共入口 trit_pack / matmul_q 的说明。 // 默认 = 引擎未接打包路径（fail-fast）；不支持打包的引擎由这里的默认 // 实现或各自的覆盖 fail-fast（能力不下沉到上层）。 - [[nodiscard]] virtual Result trit_pack_impl(const Tensor& w, - const Tensor& scale) + [[nodiscard]] virtual Result trit_pack_impl(const Tensor& w, const Tensor& scale) { - (void)w; (void)scale; + (void)w; + (void)scale; NN_FAIL("trit_pack: 该引擎未实现三值打包"); } - [[nodiscard]] virtual Result matmul_q_impl(const Tensor& packed, - const Tensor& scale, - const Tensor& B, bool transA, - std::size_t out_rows, - Precision P) + [[nodiscard]] virtual Result matmul_q_impl(const Tensor& packed, const Tensor& scale, + const Tensor& B, bool transA, + std::size_t out_rows, Precision P) { - (void)packed; (void)scale; (void)B; (void)transA; (void)out_rows; (void)P; + (void)packed; + (void)scale; + (void)B; + (void)transA; + (void)out_rows; + (void)P; NN_FAIL("matmul_q: 该引擎未实现三值 GEMM"); } - [[nodiscard]] virtual Result batched_matmul_impl( - const Tensor& A, const Tensor& B, - std::size_t batch, - bool transA, bool transB, - Scalar alpha, - Precision P) = 0; + [[nodiscard]] virtual Result batched_matmul_impl(const Tensor& A, const Tensor& B, + std::size_t batch, bool transA, + bool transB, Scalar alpha, + Precision P) = 0; // 默认实现：matmul + 逐行 add bias（兼容所有引擎） // D6-3（15 §4.6-3，M6 裁定=删除宿主兜底）：原默认实现是 // `to_matrix → 主机逐列加 bias → from_matrix` 的 PCIe 往返兜底，两引擎 // （CpuEngine/GpuEngine）均已 override，属死码且会在第三后端上静默退化。 // 改为纯虚：新引擎必须给出真现值（编译期强制）。 - [[nodiscard]] virtual Result matmul_with_bias_impl( - const Tensor& A, const Tensor& B, const Tensor& bias, - bool transA, bool transB, Precision P) = 0; + [[nodiscard]] virtual Result matmul_with_bias_impl(const Tensor& A, const Tensor& B, + const Tensor& bias, bool transA, + bool transB, Precision P) = 0; [[nodiscard]] virtual Result add_inplace_impl(Tensor& A, const Tensor& B) = 0; [[nodiscard]] virtual Result scale_inplace_impl(Tensor& A, Scalar s) = 0; [[nodiscard]] virtual Result zero_impl(Tensor& A) = 0; // ── 扫描级 ──────────────────────────────────────────────────────────── - [[nodiscard]] virtual Result scan_prefix_outer_impl( - const Tensor& K, const Tensor& V, const Tensor& P, const Tensor& R, - const Tensor& A0, const Tensor& B0, bool has_state, - std::size_t dk, std::size_t heads, bool causal, - const Tensor& boundary, bool has_bnd, - Precision prec) = 0; - [[nodiscard]] virtual Result scan_suffix_outer_impl( - const Tensor& D, const Tensor& X, const Tensor& Y, - std::size_t dk, std::size_t heads, bool causal, - const Tensor& boundary, bool has_bnd, - Precision prec) = 0; - [[nodiscard]] virtual Result outer_col_impl( - const Tensor& P, const Tensor& R, const Tensor& S, - std::size_t dk, bool has_scale, - Precision prec) = 0; + [[nodiscard]] virtual Result + scan_prefix_outer_impl(const Tensor& K, const Tensor& V, const Tensor& P, const Tensor& R, + const Tensor& A0, const Tensor& B0, bool has_state, std::size_t dk, + std::size_t heads, bool causal, const Tensor& boundary, bool has_bnd, + Precision prec) = 0; + [[nodiscard]] virtual Result scan_suffix_outer_impl(const Tensor& D, const Tensor& X, + const Tensor& Y, std::size_t dk, + std::size_t heads, bool causal, + const Tensor& boundary, + bool has_bnd, Precision prec) = 0; + [[nodiscard]] virtual Result outer_col_impl(const Tensor& P, const Tensor& R, + const Tensor& S, std::size_t dk, + bool has_scale, Precision prec) = 0; // ── 归约 ────────────────────────────────────────────────────────────── - [[nodiscard]] virtual Result row_reduce_sum_impl( - const Tensor& A, Precision P) = 0; - [[nodiscard]] virtual Result col_reduce_sum_impl( - const Tensor& A, Precision P) = 0; - [[nodiscard]] virtual Result col_reduce_max_impl( - const Tensor& A, Precision P) = 0; - [[nodiscard]] virtual Result grouped_reduce_sum_impl( - const Tensor& x, std::size_t G, std::size_t R, Precision P) = 0; - [[nodiscard]] virtual Result grouped_reduce_max_impl( - const Tensor& x, std::size_t G, std::size_t R, Precision P) = 0; + [[nodiscard]] virtual Result row_reduce_sum_impl(const Tensor& A, Precision P) = 0; + [[nodiscard]] virtual Result col_reduce_sum_impl(const Tensor& A, Precision P) = 0; + [[nodiscard]] virtual Result col_reduce_max_impl(const Tensor& A, Precision P) = 0; + [[nodiscard]] virtual Result grouped_reduce_sum_impl(const Tensor& x, std::size_t G, + std::size_t R, Precision P) = 0; + [[nodiscard]] virtual Result grouped_reduce_max_impl(const Tensor& x, std::size_t G, + std::size_t R, Precision P) = 0; // ── 表达式求值 ──────────────────────────────────────────────────────── - [[nodiscard]] virtual Result eval_expr_impl( - const ExprSpec& spec, - std::span inputs, - std::size_t rows, std::size_t cols, - Precision P) = 0; + [[nodiscard]] virtual Result eval_expr_impl(const ExprSpec& spec, + std::span inputs, + std::size_t rows, std::size_t cols, + Precision P) = 0; // 默认实现返回错误（未支持的引擎）；CPU/GPU 覆盖。 - [[nodiscard]] virtual Result eval_expr_reduce_impl( - const ExprSpec& spec, - std::span inputs, - std::size_t rows, std::size_t cols, - Precision P) - { - (void)spec; (void)inputs; (void)rows; (void)cols; (void)P; + [[nodiscard]] virtual Result eval_expr_reduce_impl(const ExprSpec& spec, + std::span inputs, + std::size_t rows, std::size_t cols, + Precision P) + { + (void)spec; + (void)inputs; + (void)rows; + (void)cols; + (void)P; NN_FAIL("eval_expr_reduce: 该引擎不支持归约向量输出"); } // 默认实现返回错误（未支持的引擎）；CPU/GPU 覆盖为真原地实现。 - [[nodiscard]] virtual Result eval_expr_into_impl( - const ExprSpec& spec, - std::span inputs, - std::size_t rows, std::size_t cols, Tensor& out) - { - (void)spec; (void)inputs; (void)rows; (void)cols; (void)out; + [[nodiscard]] virtual Result eval_expr_into_impl(const ExprSpec& spec, + std::span inputs, + std::size_t rows, std::size_t cols, + Tensor& out) + { + (void)spec; + (void)inputs; + (void)rows; + (void)cols; + (void)out; NN_FAIL("eval_expr_into: 该引擎不支持原地表达式求值"); } // P2 多输出（IR-C 写穿物化）；默认实现返回错误（未支持的引擎）。 - [[nodiscard]] virtual Result eval_expr_multi_into_impl( - const ExprSpec& spec, - std::span inputs, - std::size_t rows, std::size_t cols, std::span outs) - { - (void)spec; (void)inputs; (void)rows; (void)cols; (void)outs; + [[nodiscard]] virtual Result eval_expr_multi_into_impl(const ExprSpec& spec, + std::span inputs, + std::size_t rows, std::size_t cols, + std::span outs) + { + (void)spec; + (void)inputs; + (void)rows; + (void)cols; + (void)outs; NN_FAIL("eval_expr_multi_into: 该引擎不支持多输出表达式求值"); } @@ -1728,14 +1638,11 @@ class ComputeEngine : public std::enable_shared_from_this // ── IR-C 录制拦截────────────────────────────────── // 录制图归属校验（唯一入口）：同一录制段内的表达式必须经**同一**引擎求值。 // owner 是 shared_ptr：录制图存在期间引擎必然存活。 - [[nodiscard]] Result check_graph_owner_(const ExprGraph& g, - const char* what) const + [[nodiscard]] Result check_graph_owner_(const ExprGraph& g, const char* what) const { - if (!g.owner) - NN_FAIL(std::string(what) + ": 录制图未绑定引擎"); + if (!g.owner) NN_FAIL(std::string(what) + ": 录制图未绑定引擎"); if (g.owner.get() != this) - NN_FAIL(std::string(what) + - ": 录制段内跨引擎求值（段内表达式必须经开启录制的引擎）"); + NN_FAIL(std::string(what) + ": 录制段内跨引擎求值（段内表达式必须经开启录制的引擎）"); return {}; } @@ -1743,9 +1650,10 @@ class ComputeEngine : public std::enable_shared_from_this // 入图（add_node 建依赖边），返回带 virtual_tag 的占位 Tensor（真实存储、 // 出生即绑定——end_expr 的复合 kernel 经 output_override 直接写回该存储， // 调用方持有的句柄即物化结果）。占位精度 = P（输出存储精度）。 - [[nodiscard]] Result record_graph_node_( - ExprGraph& g, const ExprSpec& spec, std::span inputs, - std::size_t rows, std::size_t cols, bool vector_out, Precision P) + [[nodiscard]] Result record_graph_node_(ExprGraph& g, const ExprSpec& spec, + std::span inputs, + std::size_t rows, std::size_t cols, + bool vector_out, Precision P) { NN_TRY_CHECK(check_graph_owner_(g, "eval_expr")); const int node = g.add_node(spec, inputs, rows, cols, vector_out, P); @@ -1765,9 +1673,10 @@ class ComputeEngine : public std::enable_shared_from_this } // 目标传递的录制节点（P2 写穿物化）：输出 = dst 既有张量（无占位）。 - [[nodiscard]] Result record_graph_into_node_( - ExprGraph& g, const ExprSpec& spec, std::span inputs, - std::size_t rows, std::size_t cols, Tensor& out) + [[nodiscard]] Result record_graph_into_node_(ExprGraph& g, const ExprSpec& spec, + std::span inputs, + std::size_t rows, std::size_t cols, + Tensor& out) { NN_TRY_CHECK(check_graph_owner_(g, "eval_expr_into")); const int node = g.add_node(spec, inputs, rows, cols, @@ -1802,13 +1711,15 @@ class ComputeEngine : public std::enable_shared_from_this case InitSpec::Kind::Uniform: { std::uniform_real_distribution dist(spec.lo, spec.hi); - for (std::size_t i = 0; i < sp.size(); ++i) sp[i] = dist(rng); + for (std::size_t i = 0; i < sp.size(); ++i) + sp[i] = dist(rng); break; } case InitSpec::Kind::Normal: { std::normal_distribution dist(spec.mean, spec.stddev); - for (std::size_t i = 0; i < sp.size(); ++i) sp[i] = dist(rng); + for (std::size_t i = 0; i < sp.size(); ++i) + sp[i] = dist(rng); break; } default: @@ -1825,48 +1736,44 @@ class ComputeEngine : public std::enable_shared_from_this // 错误带调用点 source_location（默认参在调用处求值）与张量形状。 [[nodiscard]] Error domain_error_(const Tensor& t, std::source_location loc) const { - std::string s = "设备域不匹配 at " + std::string(loc.file_name()) + ":" - + std::to_string(loc.line()) + " " + t.shape_str() - + " tensor_domain=" + std::to_string(t.device_domain_()) - + " engine_domain=" + std::to_string(domain_); + std::string s = "设备域不匹配 at " + std::string(loc.file_name()) + ":" + + std::to_string(loc.line()) + " " + t.shape_str() + + " tensor_domain=" + std::to_string(t.device_domain_()) + + " engine_domain=" + std::to_string(domain_); const DomainId d = t.device_domain_(); if (d == kHostDomain && domain_ != kHostDomain) s += "\n CPU 张量不能直接进入 GPU 引擎：请先经 engine.import(t) 上载"; else if (d != kHostDomain && domain_ == kHostDomain) s += "\n GPU 张量不能直接进入 CPU 引擎：请先经 engine.to_matrix(t) 落宿主"; - else - s += "\n 跨 GPU backend：请先 to_matrix 落宿主，再 import 到目标 backend"; + else s += "\n 跨 GPU backend：请先 to_matrix 落宿主，再 import 到目标 backend"; return Error{std::move(s)}; } - [[nodiscard]] Result bind_check_one_(const Tensor& t, - std::source_location loc) const + [[nodiscard]] Result bind_check_one_(const Tensor& t, std::source_location loc) const { const DomainId d = t.device_domain_(); - if (d == 0 || d == domain_) - return {}; + if (d == 0 || d == domain_) return {}; NN_FAIL(domain_error_(t, loc)); } - [[nodiscard]] Result bind_check_(std::initializer_list ts, - std::source_location loc = std::source_location::current()) const + [[nodiscard]] Result + bind_check_(std::initializer_list ts, + std::source_location loc = std::source_location::current()) const { for (const Tensor* t : ts) { - if (!t) - continue; - if (auto r = bind_check_one_(*t, loc); !r) - return r; + if (!t) continue; + if (auto r = bind_check_one_(*t, loc); !r) return r; } return {}; } - [[nodiscard]] Result bind_check_(std::span ts, - std::source_location loc = std::source_location::current()) const + [[nodiscard]] Result + bind_check_(std::span ts, + std::source_location loc = std::source_location::current()) const { for (const Tensor& t : ts) - if (auto r = bind_check_one_(t, loc); !r) - return r; + if (auto r = bind_check_one_(t, loc); !r) return r; return {}; } @@ -1878,8 +1785,7 @@ class ComputeEngine : public std::enable_shared_from_this [[nodiscard]] static bool all_f32(std::span ts) noexcept { for (const auto& t : ts) - if (t.precision() != Precision::F32) - return false; + if (t.precision() != Precision::F32) return false; return true; } @@ -1889,10 +1795,8 @@ class ComputeEngine : public std::enable_shared_from_this [[nodiscard]] Result to_f32(const Tensor& t, std::source_location loc = std::source_location::current()) { - if (t.precision() == Precision::F32) - return t; - note_temp_(t.rows(), t.cols(), /*to_f32=*/true, - static_cast(loc.line())); + if (t.precision() == Precision::F32) return t; + note_temp_(t.rows(), t.cols(), /*to_f32=*/true, static_cast(loc.line())); return cast(t, Precision::F32); } @@ -1900,10 +1804,8 @@ class ComputeEngine : public std::enable_shared_from_this [[nodiscard]] Result to_prec(Tensor t, Precision P, std::source_location loc = std::source_location::current()) { - if (P == Precision::F32 || t.precision() == P) - return t; - note_temp_(t.rows(), t.cols(), /*to_f32=*/false, - static_cast(loc.line())); + if (P == Precision::F32 || t.precision() == P) return t; + note_temp_(t.rows(), t.cols(), /*to_f32=*/false, static_cast(loc.line())); // 入参常是 `*_impl` 的产物（unary_/binary_/move_ 等）。v3 起设备域 // 从存储派生，基类内部中转无需补任何身份；cast 的 bind_check_ 直接 // 比较存储域与本引擎域。 @@ -1911,9 +1813,9 @@ class ComputeEngine : public std::enable_shared_from_this } // 一批操作数抬到 f32（顺序与输入一致；ExprSpec 的 views 依赖该顺序） - [[nodiscard]] Result> to_f32_all( - std::span ts, - std::source_location loc = std::source_location::current()) + [[nodiscard]] Result> + to_f32_all(std::span ts, + std::source_location loc = std::source_location::current()) { std::vector out; out.reserve(ts.size()); @@ -1932,14 +1834,12 @@ class ComputeEngine : public std::enable_shared_from_this [[nodiscard]] Result inplace2_(Tensor& A, const Tensor& B, Fn&& fn, std::source_location loc = std::source_location::current()) { - if (A.precision() == Precision::F32 && B.precision() == Precision::F32) - return fn(A, B); + if (A.precision() == Precision::F32 && B.precision() == Precision::F32) return fn(A, B); // 纯边界 cast 助手（不做能力分支）：需要"原生 f16 直通"的引擎在自己 // 的 `*_impl` 里先处理 f16，再落到这里（或根本不进来）。 auto b = to_f32(B, loc); NN_TRY_CHECK(b); - if (A.precision() == Precision::F32) - return fn(A, *b); + if (A.precision() == Precision::F32) return fn(A, *b); auto a = to_f32(A, loc); NN_TRY_CHECK(a); auto r = fn(*a, *b); @@ -1952,8 +1852,7 @@ class ComputeEngine : public std::enable_shared_from_this [[nodiscard]] Result inplace1_(Tensor& A, Fn&& fn, std::source_location loc = std::source_location::current()) { - if (A.precision() == Precision::F32) - return fn(A); + if (A.precision() == Precision::F32) return fn(A); auto a = to_f32(A, loc); NN_TRY_CHECK(a); auto r = fn(*a); @@ -1966,8 +1865,7 @@ class ComputeEngine : public std::enable_shared_from_this [[nodiscard]] Result unary_(const Tensor& A, Precision P, Fn&& fn, std::source_location loc = std::source_location::current()) { - if (P == Precision::F32 && A.precision() == Precision::F32) - return fn(A); + if (P == Precision::F32 && A.precision() == Precision::F32) return fn(A); auto a = to_f32(A, loc); NN_TRY_CHECK(a); auto r = fn(*a); @@ -1977,8 +1875,7 @@ class ComputeEngine : public std::enable_shared_from_this // 运算类通用：抬 f32 → 算 → 按 P 落回（双操作数） template - [[nodiscard]] Result binary_(const Tensor& A, const Tensor& B, - Precision P, Fn&& fn, + [[nodiscard]] Result binary_(const Tensor& A, const Tensor& B, Precision P, Fn&& fn, std::source_location loc = std::source_location::current()) { if (P == Precision::F32 && A.precision() == Precision::F32 && @@ -1999,8 +1896,7 @@ class ComputeEngine : public std::enable_shared_from_this std::source_location loc = std::source_location::current()) { const Precision out_p = A.precision(); - if (out_p == Precision::F32) - return fn(A); + if (out_p == Precision::F32) return fn(A); auto a = to_f32(A, loc); NN_TRY_CHECK(a); auto r = fn(*a); @@ -2010,8 +1906,8 @@ class ComputeEngine : public std::enable_shared_from_this // 不同精度行插入（供引擎 `insert_rows_impl` 在非 f32 时复用）：在 f32 空间 // 插入后写回 dst 的原存储（§8.3 in-place 存储精度不可变）。 - [[nodiscard]] Result insert_rows_cast_( - Tensor& dst, std::size_t dst_start_row, const Tensor& src) + [[nodiscard]] Result insert_rows_cast_(Tensor& dst, std::size_t dst_start_row, + const Tensor& src) { auto d = to_f32(dst); NN_TRY_CHECK(d); @@ -2019,14 +1915,13 @@ class ComputeEngine : public std::enable_shared_from_this NN_TRY_CHECK(s); auto r = insert_rows_impl(*d, dst_start_row, *s); NN_TRY_CHECK(r); - if (dst.precision() == Precision::F32) - return {}; + if (dst.precision() == Precision::F32) return {}; return cast_into(*d, dst); } // 不同精度 scatter-add（供引擎 `scatter_add_rows_impl` 在非 f32 时复用）。 - [[nodiscard]] Result scatter_add_rows_cast_( - Tensor& dst, const Tensor& indices, const Tensor& grad) + [[nodiscard]] Result scatter_add_rows_cast_(Tensor& dst, const Tensor& indices, + const Tensor& grad) { auto d = to_f32(dst); NN_TRY_CHECK(d); @@ -2034,8 +1929,7 @@ class ComputeEngine : public std::enable_shared_from_this NN_TRY_CHECK(g); auto r = scatter_add_rows_impl(*d, indices, *g); NN_TRY_CHECK(r); - if (dst.precision() == Precision::F32) - return {}; + if (dst.precision() == Precision::F32) return {}; return cast_into(*d, dst); } @@ -2049,15 +1943,16 @@ class ComputeEngine : public std::enable_shared_from_this // 时 matmul/unary_/eval_expr 的 cast 无法分辨）。 // 探针门：NN_PROBES_ENABLED=0（默认）时 note_temp_ 是空实现、 // prec_env_flag 恒 false —— 逐算子归因与变体日志**整块不编译进二进制**。 - static void note_temp_(std::size_t rows, std::size_t cols, bool to_f32, - std::uint32_t line = 0) + static void note_temp_(std::size_t rows, std::size_t cols, bool to_f32, std::uint32_t line = 0) { #if !NN_PROBES_ENABLED - (void)rows; (void)cols; (void)to_f32; (void)line; + (void)rows; + (void)cols; + (void)to_f32; + (void)line; return; #else - if (!prec_trace_enabled_()) - return; + if (!prec_trace_enabled_()) return; const std::size_t bytes = rows * cols * (to_f32 ? 4u : 2u); for (auto& e : temp_stats()) if (e.rows == rows && e.cols == cols && e.to_f32 == to_f32 && e.line == line) @@ -2081,47 +1976,36 @@ class ComputeEngine : public std::enable_shared_from_this // 每次读取、零缓存（与 dsl::env_flag 同语义）；探针门：默认构建恒 false // （**不求值 getenv**），"NN_F16_DEBUG" 等诊断开关失效。 [[nodiscard]] static bool prec_env_flag(const char* name) - { - return ::nn::detail::probe_flag_raw(name); - } + { return ::nn::detail::probe_flag_raw(name); } // 诊断（NN_PREC_TRACE=1）：请求了非零精度签名却**没有命中带类型变体** —— 这是 // "静默回退边界 cast"的唯一可见信号，也是下一期该补哪些变体的清单（去重打印）。 - static void trace_miss_(const ExprSpec& spec, - std::span inputs, Precision P) + static void trace_miss_(const ExprSpec& spec, std::span inputs, Precision P) { static std::unordered_set seen; const ExprSpec canon = canonicalize_expr_spec(spec); const std::string k = expr_spec_key(canon); const ExprPrecSig sig = expr_prec_sig_of(inputs, P); const std::string vk = expr_prec_sig_key(k, sig); - if (!seen.insert(vk).second) - return; + if (!seen.insert(vk).second) return; const int raxis = expr_spec_reduce_axis(canon); - std::fprintf(stderr, - "[prec][miss] key=%s sig=0x%04x %s instrs=%zu mm=%d fold=%d raxis=%d\n", - k.c_str(), static_cast(sig), - expr_prec_sig_str(sig, inputs.size()).c_str(), - canon.instrs.size(), canon.matmul ? 1 : 0, - canon.fold ? 1 : 0, raxis); + std::fprintf( + stderr, "[prec][miss] key=%s sig=0x%04x %s instrs=%zu mm=%d fold=%d raxis=%d\n", + k.c_str(), static_cast(sig), expr_prec_sig_str(sig, inputs.size()).c_str(), + canon.instrs.size(), canon.matmul ? 1 : 0, canon.fold ? 1 : 0, raxis); } - static void trace_variant_(const ExprSpec& spec, - std::span inputs, Precision P) + static void trace_variant_(const ExprSpec& spec, std::span inputs, Precision P) { - if (!prec_trace_enabled_()) - return; + if (!prec_trace_enabled_()) return; static std::unordered_set seen; const std::string k = expr_spec_key(canonicalize_expr_spec(spec)); const ExprPrecSig sig = expr_prec_sig_of(inputs, P); const std::string vk = expr_prec_sig_key(k, sig); - if (!seen.insert(vk).second) - return; - std::fprintf(stderr, - "[prec] key=%s sig=0x%04x %s n_in=%zu instrs=%zu mm=%d fold=%d\n", + if (!seen.insert(vk).second) return; + std::fprintf(stderr, "[prec] key=%s sig=0x%04x %s n_in=%zu instrs=%zu mm=%d fold=%d\n", k.c_str(), static_cast(sig), - expr_prec_sig_str(sig, inputs.size()).c_str(), - inputs.size(), spec.instrs.size(), - spec.matmul ? 1 : 0, spec.fold ? 1 : 0); + expr_prec_sig_str(sig, inputs.size()).c_str(), inputs.size(), + spec.instrs.size(), spec.matmul ? 1 : 0, spec.fold ? 1 : 0); } // 默认（CPU / 同步引擎）标量回读槽：submit 立即存值，poll 立即就绪。 @@ -2157,8 +2041,7 @@ class ComputeEngine : public std::enable_shared_from_this std::size_t cols, Precision P, std::span src) { - if (src.size() != rows * cols) - NN_FAIL("upload_span: 元素数与形状不匹配"); + if (src.size() != rows * cols) NN_FAIL("upload_span: 元素数与形状不匹配"); // RV-02：工厂返回 Result —— 分配失败原因原样上抛（含 GPU 侧原因） NN_TRY(t_r, engine.try_create_tensor(rows, cols, P)); Tensor t = std::move(*t_r); @@ -2179,10 +2062,8 @@ class ComputeEngine : public std::enable_shared_from_this [[nodiscard]] inline Result download_span(ComputeEngine& engine, const Tensor& t, std::span dst) { - if (dst.size() != t.rows() * t.cols()) - NN_FAIL("download_span: 元素数与张量形状不匹配"); - if (t.precision() == Precision::F32) - return engine.read(t, dst); + if (dst.size() != t.rows() * t.cols()) NN_FAIL("download_span: 元素数与张量形状不匹配"); + if (t.precision() == Precision::F32) return engine.read(t, dst); std::vector buf(dst.size()); NN_TRY(r, engine.read(t, std::span(buf))); for (std::size_t i = 0; i < dst.size(); ++i) @@ -2191,14 +2072,15 @@ class ComputeEngine : public std::enable_shared_from_this } [[nodiscard]] inline Result> download_vector(ComputeEngine& engine, - const Tensor& t) + const Tensor& t) { std::vector v(t.rows() * t.cols()); NN_TRY(r, download_span(engine, t, std::span(v))); return v; } -namespace detail { +namespace detail +{ // 过渡别名（宿主桥正式归属 nn::）：库内既有调用点不破坏； // 库外请直接用 nn::upload_span / nn::download_span / nn::download_vector。 using nn::download_span; @@ -2214,18 +2096,15 @@ using nn::upload_span; class ExprSegment { public: - explicit ExprSegment(std::shared_ptr e) noexcept - : e_(std::move(e)) + explicit ExprSegment(std::shared_ptr e) noexcept : e_(std::move(e)) { - if (!e_) - NN_FATAL("ExprSegment: 引擎句柄为空（必须由 shared_ptr 管理）"); + if (!e_) NN_FATAL("ExprSegment: 引擎句柄为空（必须由 shared_ptr 管理）"); } ExprSegment(const ExprSegment&) = delete; ExprSegment& operator=(const ExprSegment&) = delete; ~ExprSegment() { - if (!done_) - (void)e_->end_expr(); // 兜底收口（错误已被上层传播或忽略） + if (!done_) (void)e_->end_expr(); // 兜底收口（错误已被上层传播或忽略） } [[nodiscard]] Result end() { diff --git a/include/neuralnet.cpp/compute_gpu_engine.hpp b/include/neuralnet.cpp/compute_gpu_engine.hpp index 963dd04d..67308cf1 100644 --- a/include/neuralnet.cpp/compute_gpu_engine.hpp +++ b/include/neuralnet.cpp/compute_gpu_engine.hpp @@ -63,13 +63,13 @@ class GpuEngine final : public ComputeEngine std::shared_ptr backend_; explicit GpuEngine(std::shared_ptr backend) - : ComputeEngine(backend ? backend->identity() : 0), - backend_(std::move(backend)) {} + : ComputeEngine(backend ? backend->identity() : 0), backend_(std::move(backend)) + { + } public: // 引擎必须由 shared_ptr 管理：Layer / Model / Optimizer 会保存强引用。 - [[nodiscard]] static std::shared_ptr create( - std::shared_ptr backend) + [[nodiscard]] static std::shared_ptr create(std::shared_ptr backend) { NN_ASSERT_ALWAYS(backend != nullptr, "GpuEngine::create: backend 为空"); return std::shared_ptr(new GpuEngine(std::move(backend))); @@ -77,9 +77,7 @@ class GpuEngine final : public ComputeEngine // 兼容仍是 GpuBackend& 的调用点：从 backend 自身取 keepalive。 [[nodiscard]] static std::shared_ptr create(GpuBackend& backend) - { - return create(backend.keepalive()); - } + { return create(backend.keepalive()); } [[nodiscard]] Device device() const noexcept override { return Device::GPU; } @@ -101,10 +99,7 @@ class GpuEngine final : public ComputeEngine // batch 模式查询（基类默认 false；GPU 委托 backend，供 InferenceScope // / Model::forward 自动作用域判定，issue #28） - [[nodiscard]] bool batch_active() const noexcept override - { - return backend_->in_batch(); - } + [[nodiscard]] bool batch_active() const noexcept override { return backend_->in_batch(); } // ── 表达式录制（IR-C）────────────────────────────── // begin_expr 开启录制：期间 eval_expr/eval_expr_reduce 经 NVI 拦截入图 @@ -115,8 +110,7 @@ class GpuEngine final : public ComputeEngine [[nodiscard]] Result begin_expr() override { auto& owner = fused::recording_graph_owner(); - if (owner) - NN_FAIL("begin_expr: 嵌套录制（已有未结束的 begin_expr）"); + if (owner) NN_FAIL("begin_expr: 嵌套录制（已有未结束的 begin_expr）"); owner = std::make_unique(); owner->owner = weak_from_this().lock(); return {}; @@ -167,32 +161,26 @@ class GpuEngine final : public ComputeEngine } // 把 src 复制到 buffer 的 offset（float 单位）处（录制式） - [[nodiscard]] Result offload_save_impl( - const Tensor& buffer, std::size_t offset, const Tensor& src) override + [[nodiscard]] Result offload_save_impl(const Tensor& buffer, std::size_t offset, + const Tensor& src) override { - if (src.is_cpu()) - return {}; - const VkDeviceSize size = - static_cast(src.size() * sizeof(float)); + if (src.is_cpu()) return {}; + const VkDeviceSize size = static_cast(src.size() * sizeof(float)); return backend_->copy_buffer_region_gpu( - src.gpu_tensor().buffer().impl(), 0, - buffer.gpu_tensor().buffer().impl(), + src.gpu_tensor().buffer().impl(), 0, buffer.gpu_tensor().buffer().impl(), static_cast(offset * sizeof(float)), size); } // 从 buffer 的 offset（float 单位）处复制 rows×cols 到新 GPU tensor（录制式） - [[nodiscard]] Result offload_restore_impl( - const Tensor& buffer, std::size_t offset, - std::size_t rows, std::size_t cols) override + [[nodiscard]] Result offload_restore_impl(const Tensor& buffer, std::size_t offset, + std::size_t rows, std::size_t cols) override { auto dst = GpuTensor::create_empty(rows, cols, *backend_); NN_TRY_CHECK(dst); - const VkDeviceSize size = - static_cast(rows * cols * sizeof(float)); - auto r = backend_->copy_buffer_region_gpu( - buffer.gpu_tensor().buffer().impl(), - static_cast(offset * sizeof(float)), - dst->buffer().impl(), 0, size); + const VkDeviceSize size = static_cast(rows * cols * sizeof(float)); + auto r = backend_->copy_buffer_region_gpu(buffer.gpu_tensor().buffer().impl(), + static_cast(offset * sizeof(float)), + dst->buffer().impl(), 0, size); NN_TRY_CHECK(r); return Tensor::from_gpu(std::move(*dst)); } @@ -201,19 +189,16 @@ class GpuEngine final : public ComputeEngine // 把 (1,1) F32 标量排入一次 D2H 拷贝并提交，不等待。调用时机：产出该 // 标量的主帧（flush_batch/end_batch）已提交之后——同队列 FIFO 保证拷贝 // 执行在生产命令之后。poll 非阻塞，就绪即取（见 compute_engine.hpp）。 - [[nodiscard]] Result submit_scalar_readback_impl( - std::size_t slot, const Tensor& t) override + [[nodiscard]] Result submit_scalar_readback_impl(std::size_t slot, + const Tensor& t) override { - if (t.is_cpu()) - NN_FAIL("submit_scalar_readback: 期望 GPU 张量"); + if (t.is_cpu()) NN_FAIL("submit_scalar_readback: 期望 GPU 张量"); if (t.precision() != Precision::F32) NN_FAIL("submit_scalar_readback: 仅支持 F32 标量（调用方先 cast）"); - return backend_->submit_scalar_readback( - slot, t.gpu_tensor().buffer().impl()); + return backend_->submit_scalar_readback(slot, t.gpu_tensor().buffer().impl()); } - [[nodiscard]] Result poll_scalar_readback( - std::size_t slot, Scalar& out) override + [[nodiscard]] Result poll_scalar_readback(std::size_t slot, Scalar& out) override { float v = 0.0f; auto r = backend_->poll_scalar_readback(slot, v); @@ -223,9 +208,7 @@ class GpuEngine final : public ComputeEngine } [[nodiscard]] std::size_t scalar_readback_slots() const override - { - return GpuBackend::scalar_readback_slot_count(); - } + { return GpuBackend::scalar_readback_slot_count(); } // ── 中点刷新：提交当前 command buffer 并开始新的录制 ── // 用于拆分大 batch（如 forward 与 backward 之间），防 TDR。 @@ -245,7 +228,7 @@ class GpuEngine final : public ComputeEngine // RV-02：分配失败原因上抛（旧形态 `if (!r) return Tensor();` 把显存/设备 // 错误整块丢弃，调用方只能看到"空张量"）。 [[nodiscard]] Result create_tensor_impl(std::size_t rows, std::size_t cols, - Precision P) override + Precision P) override { if (P == Precision::F16) { @@ -324,8 +307,7 @@ class GpuEngine final : public ComputeEngine // ── cast 原语（§7.5，统一接口；GPU 原生转换，无 PCIe 往返）───────────── [[nodiscard]] Result cast_impl(const Tensor& src, Precision dst) override { - if (src.precision() == dst) - return src; + if (src.precision() == dst) return src; // 保留精度 Pair（BF16/F64 Phase 1 未实现）清晰报错 NN_TRY(pcheck, check_precision_supported(src.precision())); @@ -338,9 +320,8 @@ class GpuEngine final : public ComputeEngine // f16 GPU → f32 GPU（升 cast，精确无损，GPU kernel） auto dst_gpu = GpuTensor::create_empty(src.rows(), src.cols(), *backend_); NN_TRY_CHECK(dst_gpu); - auto r = backend_->cast_gpu( - src.gpu_tensor().buffer(), dst_gpu->buffer(), - count, /*kind=0*/ 0u); + auto r = backend_->cast_gpu(src.gpu_tensor().buffer(), + dst_gpu->buffer(), count, /*kind=0*/ 0u); NN_TRY_CHECK(r); return Tensor::from_gpu(std::move(*dst_gpu)); } @@ -350,9 +331,8 @@ class GpuEngine final : public ComputeEngine // create_f16_tensor 保证奇数元素 count 时 word 写入不越界。 auto dst_gpu = backend_->create_f16_tensor(src.rows(), src.cols()); NN_TRY_CHECK(dst_gpu); - auto r = backend_->cast_gpu( - src.gpu_tensor().buffer(), dst_gpu->buffer(), - count, /*kind=1*/ 1u); + auto r = backend_->cast_gpu(src.gpu_tensor().buffer(), dst_gpu->buffer(), count, + /*kind=1*/ 1u); NN_TRY_CHECK(r); return Tensor::from_gpu(std::move(*dst_gpu)); } @@ -373,21 +353,18 @@ class GpuEngine final : public ComputeEngine WriteEpochGuard wg{dst}; if (dst.rows() != src.rows() || dst.cols() != src.cols()) NN_FAIL("copy_into: shape mismatch"); - if (dst.precision() != src.precision()) - NN_FAIL("copy_into: precision mismatch"); + if (dst.precision() != src.precision()) NN_FAIL("copy_into: precision mismatch"); // v3：本引擎的 copy_into 只接受本域内的 GPU 张量（CPU 数据先 import） if (dst.device_domain_() != domain() || src.device_domain_() != domain()) NN_FAIL("copy_into: 操作数不在本 GPU 引擎的设备域"); const std::size_t count = dst.rows() * dst.cols(); if (dst.precision() == Precision::F16) - return backend_->copy_buffer_gpu( - src.gpu_tensor().buffer().impl(), - dst.gpu_tensor().buffer().impl(), - static_cast(count * 2u)); - return backend_->copy_buffer_gpu( - src.gpu_tensor().buffer().impl(), - dst.gpu_tensor().buffer().impl(), - static_cast(count * sizeof(float))); + return backend_->copy_buffer_gpu(src.gpu_tensor().buffer().impl(), + dst.gpu_tensor().buffer().impl(), + static_cast(count * 2u)); + return backend_->copy_buffer_gpu(src.gpu_tensor().buffer().impl(), + dst.gpu_tensor().buffer().impl(), + static_cast(count * sizeof(float))); } [[nodiscard]] Result cast_into(const Tensor& src, Tensor& dst) override @@ -396,20 +373,18 @@ class GpuEngine final : public ComputeEngine WriteEpochGuard wg{dst}; if (dst.rows() != src.rows() || dst.cols() != src.cols()) NN_FAIL("cast_into: shape mismatch"); - if (src.precision() == dst.precision()) - return copy_into(dst, src); + if (src.precision() == dst.precision()) return copy_into(dst, src); // v3：本引擎的 cast_into 只接受本域内 GPU 张量 if (dst.device_domain_() != domain() || src.device_domain_() != domain()) NN_FAIL("cast_into: 操作数不在本 GPU 引擎的设备域"); const std::size_t count = dst.rows() * dst.cols(); if (src.precision() == Precision::F16 && dst.precision() == Precision::F32) - return backend_->cast_gpu( - src.gpu_tensor().buffer(), - dst.gpu_tensor().buffer(), count, /*kind=*/0u); + return backend_->cast_gpu(src.gpu_tensor().buffer(), + dst.gpu_tensor().buffer(), count, /*kind=*/0u); if (src.precision() == Precision::F32 && dst.precision() == Precision::F16) - return backend_->cast_gpu( - src.gpu_tensor().buffer(), - dst.gpu_tensor().buffer(), count, /*kind=*/1u); + return backend_->cast_gpu(src.gpu_tensor().buffer(), + dst.gpu_tensor().buffer(), count, + /*kind=*/1u); NN_FAIL("cast_into: unsupported precision conversion"); } @@ -418,8 +393,7 @@ class GpuEngine final : public ComputeEngine { if (dst.rows() != src.rows() || dst.cols() != src.cols()) NN_FAIL("copy_from: shape mismatch"); - if (dst.device_domain_() != domain()) - NN_FAIL("copy_from: dst 不在本 GPU 引擎的设备域"); + if (dst.device_domain_() != domain()) NN_FAIL("copy_from: dst 不在本 GPU 引擎的设备域"); // ── 非 fresh（既有张量）：必须先 drain（提交不等待）────────────── // dst 是当前帧已引用的既有 buffer——当前帧尚未提交，若先把上传提交 // 到队列，GPU 会先执行上传、后执行当前帧命令 → 本帧后续对 dst 的 @@ -468,8 +442,7 @@ class GpuEngine final : public ComputeEngine [[nodiscard]] Result clone_impl(const Tensor& src) override { // v3：只接受本域内 GPU 张量（CPU 数据先 import；原 CPU 回退已删除） - if (src.device_domain_() != domain()) - NN_FAIL("clone: 操作数不在本 GPU 引擎的设备域"); + if (src.device_domain_() != domain()) NN_FAIL("clone: 操作数不在本 GPU 引擎的设备域"); if (src.precision() == Precision::F16) // 原生 f16 字节拷贝（后端模板化实现） { auto r16 = backend_->clone_gpu(src.gpu_tensor()); @@ -482,15 +455,15 @@ class GpuEngine final : public ComputeEngine } // ── 行切片：GPU 内拷贝连续行区间 ── - [[nodiscard]] Result slice_rows_impl( - const Tensor& src, std::size_t start_row, std::size_t count) override + [[nodiscard]] Result slice_rows_impl(const Tensor& src, std::size_t start_row, + std::size_t count) override { auto src_gpu = import(src); NN_TRY_CHECK(src_gpu); if (src_gpu->precision() == Precision::F16) // 原生 f16 行切片（纯字节拷贝） { - auto r16 = backend_->slice_rows_gpu( - src_gpu->gpu_tensor(), start_row, count); + auto r16 = + backend_->slice_rows_gpu(src_gpu->gpu_tensor(), start_row, count); NN_TRY_CHECK(r16); return Tensor::from_gpu(std::move(*r16)); } @@ -501,25 +474,23 @@ class GpuEngine final : public ComputeEngine // ── 行插入：GPU 内就地写入连续行区间 ── // 纯 GPU 架构：dst 必须为 GPU Tensor，传 CPU Tensor 直接报错（无 CPU 回退路径）。 - [[nodiscard]] Result insert_rows_impl( - Tensor& dst, std::size_t dst_start_row, const Tensor& src) override + [[nodiscard]] Result insert_rows_impl(Tensor& dst, std::size_t dst_start_row, + const Tensor& src) override { - if (dst.is_cpu()) - NN_FAIL("insert_rows: dst must be GPU tensor in pure-GPU architecture"); + if (dst.is_cpu()) NN_FAIL("insert_rows: dst must be GPU tensor in pure-GPU architecture"); auto src_gpu = import(src); NN_TRY_CHECK(src_gpu); if (dst.precision() == Precision::F16) // 原生 f16 行插入（纯字节拷贝） - return backend_->insert_rows_gpu(dst.gpu_tensor(), - dst_start_row, - src_gpu->gpu_tensor()); + return backend_->insert_rows_gpu(dst.gpu_tensor(), dst_start_row, + src_gpu->gpu_tensor()); return backend_->insert_rows_gpu(dst.gpu_tensor(), dst_start_row, src_gpu->gpu_tensor()); } // ── gather_rows: 按 indices 从 table 中按行查表 ── // GPU-native 实现：全程在 GPU 执行，无 PCIe 传输。 // indices 支持任意形状，按 flat 遍历所有元素。 - [[nodiscard]] Result gather_rows_impl( - const Tensor& table, const Tensor& indices) override + [[nodiscard]] Result gather_rows_impl(const Tensor& table, + const Tensor& indices) override { auto tbl_gpu = import(table); NN_TRY_CHECK(tbl_gpu); @@ -529,15 +500,12 @@ class GpuEngine final : public ComputeEngine // 精度要求）。无 f16 pipeline → 引擎内边界 cast 回退。 if (tbl_gpu->precision() == Precision::F16) { - if (backend_->has(PipelineId::gather_f16) && - idx_gpu->precision() == Precision::F32) + if (backend_->has(PipelineId::gather_f16) && idx_gpu->precision() == Precision::F32) { - auto r = backend_->gather_gpu(f16_view(*tbl_gpu), - idx_gpu->gpu_tensor(), - /*f16_io=*/true); + auto r = backend_->gather_gpu(f16_view(*tbl_gpu), idx_gpu->gpu_tensor(), + /*f16_io=*/true); NN_TRY_CHECK(r); - return Tensor::from_gpu( - GpuTensorF16(r->shared_buffer(), r->rows(), r->cols())); + return Tensor::from_gpu(GpuTensorF16(r->shared_buffer(), r->rows(), r->cols())); } auto t32 = cast(*tbl_gpu, Precision::F32); NN_TRY_CHECK(t32); @@ -565,11 +533,10 @@ class GpuEngine final : public ComputeEngine // （dst 每 word 2 个 half、grad half 直读；indices 恒 float），免去为 // dst/grad 物化 f32 副本。任一非 f16 或无 pipeline → 引擎内边界 cast // 回退（正确性不变）。 - [[nodiscard]] Result scatter_add_rows_impl( - Tensor& dst, const Tensor& indices, const Tensor& grad) override + [[nodiscard]] Result scatter_add_rows_impl(Tensor& dst, const Tensor& indices, + const Tensor& grad) override { - if (!dst.is_gpu()) - NN_FAIL("scatter_add_rows: dst must be GPU tensor"); + if (!dst.is_gpu()) NN_FAIL("scatter_add_rows: dst must be GPU tensor"); auto idx_gpu = import(indices); NN_TRY_CHECK(idx_gpu); auto grad_gpu = import(grad); @@ -581,29 +548,28 @@ class GpuEngine final : public ComputeEngine { GpuTensor dst_ref = f16_view(dst); // 局部保活（禁 &临时） const GpuTensor grad_ref = f16_view(*grad_gpu); - return backend_->scatter_add_gpu(dst_ref, idx_gpu->gpu_tensor(), - grad_ref, /*f16_io=*/true); + return backend_->scatter_add_gpu(dst_ref, idx_gpu->gpu_tensor(), grad_ref, + /*f16_io=*/true); } if (dst.precision() == Precision::F16 || grad_gpu->precision() == Precision::F16) { // 引擎内 cast 回退：dst 若 f16 → cast_into 写回原存储（§8.3） - auto d32 = (dst.precision() == Precision::F16) - ? cast(dst, Precision::F32) : Result(dst); + auto d32 = (dst.precision() == Precision::F16) ? cast(dst, Precision::F32) + : Result(dst); NN_TRY_CHECK(d32); - auto g32 = (grad_gpu->precision() == Precision::F16) - ? cast(*grad_gpu, Precision::F32) : Result(*grad_gpu); + auto g32 = (grad_gpu->precision() == Precision::F16) ? cast(*grad_gpu, Precision::F32) + : Result(*grad_gpu); NN_TRY_CHECK(g32); auto g32_gpu = import(*g32); NN_TRY_CHECK(g32_gpu); auto r = backend_->scatter_add_gpu(d32->gpu_tensor(), idx_gpu->gpu_tensor(), - g32_gpu->gpu_tensor()); + g32_gpu->gpu_tensor()); NN_TRY_CHECK(r); - if (dst.precision() == Precision::F16) - return cast_into(*d32, dst); + if (dst.precision() == Precision::F16) return cast_into(*d32, dst); return {}; } - return backend_->scatter_add_gpu( - dst.gpu_tensor(), idx_gpu->gpu_tensor(), grad_gpu->gpu_tensor()); + return backend_->scatter_add_gpu(dst.gpu_tensor(), idx_gpu->gpu_tensor(), + grad_gpu->gpu_tensor()); } private: @@ -611,11 +577,9 @@ class GpuEngine final : public ComputeEngine [[nodiscard]] bool in_batch() const noexcept { return backend_->in_batch(); } public: - // ── 3D 维度转置：(M, B, N) ↔ (B, M, N) ── - [[nodiscard]] Result rearrange_3d_impl( - const Tensor& x, std::size_t M, std::size_t B, std::size_t N, - bool inverse) override + [[nodiscard]] Result rearrange_3d_impl(const Tensor& x, std::size_t M, std::size_t B, + std::size_t N, bool inverse) override { auto x_gpu = import(x); NN_TRY_CHECK(x_gpu); @@ -625,29 +589,23 @@ class GpuEngine final : public ComputeEngine if (backend_->has(PipelineId::rearrange_3d_f16)) { auto r = backend_->rearrange_3d_gpu( - f16_view(*x_gpu), static_cast(M), - static_cast(B), static_cast(N), - inverse ? 1u : 0u, /*f16_io=*/true); + f16_view(*x_gpu), static_cast(M), static_cast(B), + static_cast(N), inverse ? 1u : 0u, /*f16_io=*/true); NN_TRY_CHECK(r); - return Tensor::from_gpu( - GpuTensorF16(r->shared_buffer(), r->rows(), r->cols())); + return Tensor::from_gpu(GpuTensorF16(r->shared_buffer(), r->rows(), r->cols())); } auto x32 = cast(*x_gpu, Precision::F32); NN_TRY_CHECK(x32); - auto r = backend_->rearrange_3d_gpu( - x32->gpu_tensor(), static_cast(M), - static_cast(B), static_cast(N), - inverse ? 1u : 0u); + auto r = backend_->rearrange_3d_gpu(x32->gpu_tensor(), static_cast(M), + static_cast(B), static_cast(N), + inverse ? 1u : 0u); NN_TRY_CHECK(r); Tensor out32 = Tensor::from_gpu(std::move(*r)); return cast(std::move(out32), Precision::F16); } - auto r = backend_->rearrange_3d_gpu( - x_gpu->gpu_tensor(), - static_cast(M), - static_cast(B), - static_cast(N), - inverse ? 1u : 0u); + auto r = backend_->rearrange_3d_gpu(x_gpu->gpu_tensor(), static_cast(M), + static_cast(B), static_cast(N), + inverse ? 1u : 0u); NN_TRY_CHECK(r); return Tensor::from_gpu(std::move(*r)); } @@ -666,8 +624,7 @@ class GpuEngine final : public ComputeEngine { auto r = backend_->transpose_gpu(f16_view(*a_gpu), /*f16_io=*/true); NN_TRY_CHECK(r); - return Tensor::from_gpu( - GpuTensorF16(r->shared_buffer(), r->rows(), r->cols())); + return Tensor::from_gpu(GpuTensorF16(r->shared_buffer(), r->rows(), r->cols())); } auto a32 = cast(*a_gpu, Precision::F32); NN_TRY_CHECK(a32); @@ -682,11 +639,10 @@ class GpuEngine final : public ComputeEngine } // ── 卷积/池化窗口展开（纯数据搬运，无算法；契约见 compute_engine.hpp）── - [[nodiscard]] Result im2col_impl( - const Tensor& x, - std::size_t C, std::size_t H, std::size_t W, - std::size_t k, std::size_t stride, std::size_t pad, - std::size_t OH, std::size_t OW) override + [[nodiscard]] Result im2col_impl(const Tensor& x, std::size_t C, std::size_t H, + std::size_t W, std::size_t k, std::size_t stride, + std::size_t pad, std::size_t OH, + std::size_t OW) override { auto x_gpu = import(x); NN_TRY_CHECK(x_gpu); @@ -697,32 +653,27 @@ class GpuEngine final : public ComputeEngine { if (backend_->has(PipelineId::im2col_f16)) { - auto r = backend_->im2col_gpu(f16_view(*x_gpu), - C, H, W, k, stride, pad, OH, OW, - /*f16_io=*/true); + auto r = backend_->im2col_gpu(f16_view(*x_gpu), C, H, W, k, stride, pad, OH, OW, + /*f16_io=*/true); NN_TRY_CHECK(r); - return Tensor::from_gpu( - GpuTensorF16(r->shared_buffer(), r->rows(), r->cols())); + return Tensor::from_gpu(GpuTensorF16(r->shared_buffer(), r->rows(), r->cols())); } auto x32 = cast(*x_gpu, Precision::F32); NN_TRY_CHECK(x32); - auto r = backend_->im2col_gpu(x32->gpu_tensor(), - C, H, W, k, stride, pad, OH, OW); + auto r = backend_->im2col_gpu(x32->gpu_tensor(), C, H, W, k, stride, pad, OH, OW); NN_TRY_CHECK(r); Tensor out32 = Tensor::from_gpu(std::move(*r)); return cast(std::move(out32), Precision::F16); } - auto r = backend_->im2col_gpu(x_gpu->gpu_tensor(), - C, H, W, k, stride, pad, OH, OW); + auto r = backend_->im2col_gpu(x_gpu->gpu_tensor(), C, H, W, k, stride, pad, OH, OW); NN_TRY_CHECK(r); return Tensor::from_gpu(std::move(*r)); } - [[nodiscard]] Result col2im_impl( - const Tensor& col, - std::size_t C, std::size_t H, std::size_t W, - std::size_t k, std::size_t stride, std::size_t pad, - std::size_t OH, std::size_t OW) override + [[nodiscard]] Result col2im_impl(const Tensor& col, std::size_t C, std::size_t H, + std::size_t W, std::size_t k, std::size_t stride, + std::size_t pad, std::size_t OH, + std::size_t OW) override { auto c_gpu = import(col); NN_TRY_CHECK(c_gpu); @@ -731,23 +682,19 @@ class GpuEngine final : public ComputeEngine { if (backend_->has(PipelineId::col2im_f16)) { - auto r = backend_->col2im_gpu(f16_view(*c_gpu), - C, H, W, k, stride, pad, OH, OW, - /*f16_io=*/true); + auto r = backend_->col2im_gpu(f16_view(*c_gpu), C, H, W, k, stride, pad, OH, OW, + /*f16_io=*/true); NN_TRY_CHECK(r); - return Tensor::from_gpu( - GpuTensorF16(r->shared_buffer(), r->rows(), r->cols())); + return Tensor::from_gpu(GpuTensorF16(r->shared_buffer(), r->rows(), r->cols())); } auto c32 = cast(*c_gpu, Precision::F32); NN_TRY_CHECK(c32); - auto r = backend_->col2im_gpu(c32->gpu_tensor(), - C, H, W, k, stride, pad, OH, OW); + auto r = backend_->col2im_gpu(c32->gpu_tensor(), C, H, W, k, stride, pad, OH, OW); NN_TRY_CHECK(r); Tensor out32 = Tensor::from_gpu(std::move(*r)); return cast(std::move(out32), Precision::F16); } - auto r = backend_->col2im_gpu(c_gpu->gpu_tensor(), - C, H, W, k, stride, pad, OH, OW); + auto r = backend_->col2im_gpu(c_gpu->gpu_tensor(), C, H, W, k, stride, pad, OH, OW); NN_TRY_CHECK(r); return Tensor::from_gpu(std::move(*r)); } @@ -755,8 +702,8 @@ class GpuEngine final : public ComputeEngine // ── 分组归约（契约见 compute_engine.hpp）────────────────────────── // Phase C2c：f16 输入 → f16 分组归约 pipeline（half 直读、f32 归约、 // 输出 f32 (G,N)）；无 pipeline → 引擎内边界 cast 回退。 - [[nodiscard]] Result grouped_reduce_f16_or_cast_( - const Tensor& x, std::size_t G, std::size_t R, bool is_max) + [[nodiscard]] Result grouped_reduce_f16_or_cast_(const Tensor& x, std::size_t G, + std::size_t R, bool is_max) { auto x_gpu = import(x); NN_TRY_CHECK(x_gpu); @@ -764,8 +711,8 @@ class GpuEngine final : public ComputeEngine { if (backend_->has(PipelineId::group_reduce_f16)) { - auto r = backend_->grouped_reduce_gpu( - f16_view(*x_gpu), G, R, is_max, /*f16_in=*/true); + auto r = + backend_->grouped_reduce_gpu(f16_view(*x_gpu), G, R, is_max, /*f16_in=*/true); NN_TRY_CHECK(r); return Tensor::from_gpu(std::move(*r)); } @@ -782,19 +729,15 @@ class GpuEngine final : public ComputeEngine return Tensor::from_gpu(std::move(*r)); } - [[nodiscard]] Result grouped_reduce_sum_impl( - const Tensor& x, std::size_t G, std::size_t R, - Precision = Precision::F32) override - { - return grouped_reduce_f16_or_cast_(x, G, R, /*is_max=*/false); - } + [[nodiscard]] Result grouped_reduce_sum_impl(const Tensor& x, std::size_t G, + std::size_t R, + Precision = Precision::F32) override + { return grouped_reduce_f16_or_cast_(x, G, R, /*is_max=*/false); } - [[nodiscard]] Result grouped_reduce_max_impl( - const Tensor& x, std::size_t G, std::size_t R, - Precision = Precision::F32) override - { - return grouped_reduce_f16_or_cast_(x, G, R, /*is_max=*/true); - } + [[nodiscard]] Result grouped_reduce_max_impl(const Tensor& x, std::size_t G, + std::size_t R, + Precision = Precision::F32) override + { return grouped_reduce_f16_or_cast_(x, G, R, /*is_max=*/true); } // ══════════════════════════════════════════════════════════════════════ // 矩阵级原语 @@ -803,13 +746,11 @@ class GpuEngine final : public ComputeEngine // ── 三值打包权重（P2 / F3）：打包字容器 + 量化 GEMM ──────────────────── // 引擎私有能力：打包 pipeline 与 matmul_q pipeline 都在（缺任一 → // `matmul_q_impl` / `trit_pack_impl` 在此 fail-fast；对外事实见 describe()）。 - [[nodiscard]] Result trit_pack_impl(const Tensor& w, - const Tensor& scale) override + [[nodiscard]] Result trit_pack_impl(const Tensor& w, const Tensor& scale) override { if (!(backend_->has(PipelineId::trit_pack) && backend_->has(PipelineId::matmul_q))) NN_FAIL("trit_pack: 该设备未创建三值打包 pipeline（docs 21 §5 P2）"); - if (!w.is_gpu() || !scale.is_gpu()) - NN_FAIL("trit_pack: GPU 张量必需"); + if (!w.is_gpu() || !scale.is_gpu()) NN_FAIL("trit_pack: GPU 张量必需"); if (w.precision() != Precision::F32 || scale.precision() != Precision::F32) NN_FAIL("trit_pack: 需要 f32 权重 + f32 尺度"); if (scale.rows() != w.rows() || scale.cols() != 1) @@ -818,24 +759,22 @@ class GpuEngine final : public ComputeEngine const std::size_t wpr = nn::trit_pack::words_per_row(k); auto out = GpuTensor::create_empty(rows, wpr, *backend_); NN_TRY_CHECK(out); - NN_TRY(r, backend_->trit_pack_gpu( - w.gpu_tensor().buffer(), scale.gpu_tensor().buffer(), out->buffer(), - static_cast(rows), static_cast(k), - static_cast(wpr))); + NN_TRY(r, backend_->trit_pack_gpu(w.gpu_tensor().buffer(), scale.gpu_tensor().buffer(), + out->buffer(), static_cast(rows), + static_cast(k), + static_cast(wpr))); return Tensor::from_gpu(std::move(*out)); } - [[nodiscard]] Result matmul_q_impl(const Tensor& packed, - const Tensor& scale, const Tensor& B, - bool transA, std::size_t out_rows, + [[nodiscard]] Result matmul_q_impl(const Tensor& packed, const Tensor& scale, + const Tensor& B, bool transA, std::size_t out_rows, Precision P) override { if (!(backend_->has(PipelineId::trit_pack) && backend_->has(PipelineId::matmul_q))) NN_FAIL("matmul_q: 该设备未创建三值 GEMM pipeline（docs 21 §5 P2）"); if (!packed.is_gpu() || !scale.is_gpu() || !B.is_gpu()) NN_FAIL("matmul_q: GPU 张量必需（跨设备请先 import）"); - if (packed.precision() != Precision::F32) - NN_FAIL("matmul_q: 打包容器必须是 f32 容器"); + if (packed.precision() != Precision::F32) NN_FAIL("matmul_q: 打包容器必须是 f32 容器"); // 形状推导与校验（与 CPU 实现同一口径） const std::size_t N = B.cols(); const std::size_t M = transA ? out_rows : packed.rows(); @@ -856,27 +795,25 @@ class GpuEngine final : public ComputeEngine auto out16 = GpuTensorF16::create_empty(M, N, *backend_); NN_TRY_CHECK(out16); NN_TRY(r, backend_->matmul_q_gpu( - packed.gpu_tensor().buffer(), scale.gpu_tensor().buffer(), - B.gpu_tensor().buffer(), out16->buffer(), - static_cast(M), static_cast(N), - static_cast(K), transA ? 1u : 0u, - static_cast(wpr), /*f16=*/true)); + packed.gpu_tensor().buffer(), scale.gpu_tensor().buffer(), + B.gpu_tensor().buffer(), out16->buffer(), + static_cast(M), static_cast(N), + static_cast(K), transA ? 1u : 0u, + static_cast(wpr), /*f16=*/true)); return Tensor::from_gpu(std::move(*out16)); } // f32 路径（B 为 f16 时先抬 f32；输出按 P 落回） // 注：to_f32 是基类**私有**助手 → 用公共 cast（同 matmul_impl 的口径） - auto b32 = (B.precision() == Precision::F32) ? Result{B} - : cast(B, Precision::F32); + auto b32 = (B.precision() == Precision::F32) ? Result{B} : cast(B, Precision::F32); NN_TRY_CHECK(b32); auto out = GpuTensor::create_empty(M, N, *backend_); NN_TRY_CHECK(out); NN_TRY(r, backend_->matmul_q_gpu( - packed.gpu_tensor().buffer(), scale.gpu_tensor().buffer(), - b32->gpu_tensor().buffer(), out->buffer(), - static_cast(M), static_cast(N), - static_cast(K), transA ? 1u : 0u, - static_cast(wpr), /*f16=*/false)); + packed.gpu_tensor().buffer(), scale.gpu_tensor().buffer(), + b32->gpu_tensor().buffer(), out->buffer(), static_cast(M), + static_cast(N), static_cast(K), + transA ? 1u : 0u, static_cast(wpr), /*f16=*/false)); Tensor res = Tensor::from_gpu(std::move(*out)); if (P == Precision::F32) return res; // 静态工厂直构的 GPU 张量：设备域从存储派生（v3），直接进公共 cast @@ -884,10 +821,8 @@ class GpuEngine final : public ComputeEngine return cast(std::move(res), P); } - [[nodiscard]] Result matmul_impl( - const Tensor& A, const Tensor& B, - bool transA, bool transB, - Precision P = Precision::F32) override + [[nodiscard]] Result matmul_impl(const Tensor& A, const Tensor& B, bool transA, + bool transB, Precision P = Precision::F32) override { // 确保 A、B 在 GPU 上 auto a_gpu = import(A); @@ -897,9 +832,8 @@ class GpuEngine final : public ComputeEngine // ── F32 路径（现状，零改动）───────────────────────────────────── if (P == Precision::F32) { - auto r = backend_->matmul_gpu( - a_gpu->gpu_tensor(), b_gpu->gpu_tensor(), - transA ? 1u : 0u, transB ? 1u : 0u); + auto r = backend_->matmul_gpu(a_gpu->gpu_tensor(), b_gpu->gpu_tensor(), + transA ? 1u : 0u, transB ? 1u : 0u); NN_TRY_CHECK(r); return Tensor::from_gpu(std::move(*r)); } @@ -915,19 +849,16 @@ class GpuEngine final : public ComputeEngine const std::size_t n_out = transB ? B.rows() : B.cols(); // 大 N：tiled f16；小 N（n_out≤8，GEMV 热路径）：gemv f16 //（C2e：变体 pipeline 缺失时 f16_ok 为假 → 落到下方边界 cast） - const bool f16_ok = - a_gpu->precision() == Precision::F16 && - b_gpu->precision() == Precision::F16 && - ((n_out > 8 && backend_->has(PipelineId::matmul_tiled_f16)) || - (n_out <= 8 && backend_->has(PipelineId::matmul_gemv_f16))); + const bool f16_ok = a_gpu->precision() == Precision::F16 && + b_gpu->precision() == Precision::F16 && + ((n_out > 8 && backend_->has(PipelineId::matmul_tiled_f16)) || + (n_out <= 8 && backend_->has(PipelineId::matmul_gemv_f16))); if (f16_ok) { - auto r = backend_->matmul_gpu( - f16_view(*a_gpu), f16_view(*b_gpu), - transA ? 1u : 0u, transB ? 1u : 0u, /*f16_io=*/true); + auto r = backend_->matmul_gpu(f16_view(*a_gpu), f16_view(*b_gpu), transA ? 1u : 0u, + transB ? 1u : 0u, /*f16_io=*/true); NN_TRY_CHECK(r); - return Tensor::from_gpu( - GpuTensorF16(r->shared_buffer(), r->rows(), r->cols())); + return Tensor::from_gpu(GpuTensorF16(r->shared_buffer(), r->rows(), r->cols())); } } @@ -945,9 +876,8 @@ class GpuEngine final : public ComputeEngine auto b32_gpu = import(*b32); NN_TRY_CHECK(b32_gpu); // f32 matmul - auto r = backend_->matmul_gpu( - a32_gpu->gpu_tensor(), b32_gpu->gpu_tensor(), - transA ? 1u : 0u, transB ? 1u : 0u); + auto r = backend_->matmul_gpu(a32_gpu->gpu_tensor(), b32_gpu->gpu_tensor(), + transA ? 1u : 0u, transB ? 1u : 0u); NN_TRY_CHECK(r); // cast 回 f16 Tensor result_f32 = Tensor::from_gpu(std::move(*r)); @@ -967,21 +897,20 @@ class GpuEngine final : public ComputeEngine // （f16 直读 / f32 累加 / f16 写回，零边界 cast：Linear::forward 是 cast // 临时量最大的单一来源），无变体才在本函数内回退"抬 f32 → f32 融合 → // 落回"（**绝不把缺变体变成硬报错**，闭合世界原则只在 f32 结构缺失时适用）。 - [[nodiscard]] Result matmul_with_bias_impl( - const Tensor& A, const Tensor& B, const Tensor& bias, - bool transA = false, bool transB = false, - Precision P = Precision::F32) override + [[nodiscard]] Result matmul_with_bias_impl(const Tensor& A, const Tensor& B, + const Tensor& bias, bool transA = false, + bool transB = false, + Precision P = Precision::F32) override { const std::size_t rows = transA ? A.cols() : A.rows(); const std::size_t cols = transB ? B.rows() : B.cols(); if (P == Precision::F32) - return nn::dsl::compute(*this, - nn::dsl::matmul(A, B, transA, transB, 1) - + nn::dsl::row_broadcast(bias), + return nn::dsl::compute( + *this, nn::dsl::matmul(A, B, transA, transB, 1) + nn::dsl::row_broadcast(bias), rows, cols); - auto [spec, inputs] = nn::dsl::to_expr_spec( - nn::dsl::matmul(A, B, transA, transB, 1) + nn::dsl::row_broadcast(bias)); + auto [spec, inputs] = nn::dsl::to_expr_spec(nn::dsl::matmul(A, B, transA, transB, 1) + + nn::dsl::row_broadcast(bias)); NN_TRY(v, nn::validate_expr_spec(spec, inputs.size())); // 变体选路与"无变体 → 边界 cast 回退"都在 eval_expr_impl 内部 //（F2 修复后基类入口不再代劳）；这里只需把 P 下传。 @@ -991,12 +920,10 @@ class GpuEngine final : public ComputeEngine // ── 批量矩阵乘法：按 batch 切分行块，单次 dispatch 处理所有 batch ── // alpha 在 shader 写出时一次完成（如注意力 1/sqrt(d_k) 缩放） // P: 计算精度（§8.1，同 matmul） - [[nodiscard]] Result batched_matmul_impl( - const Tensor& A, const Tensor& B, - std::size_t batch, - bool transA, bool transB, - Scalar alpha, - Precision P = Precision::F32) override + [[nodiscard]] Result batched_matmul_impl(const Tensor& A, const Tensor& B, + std::size_t batch, bool transA, bool transB, + Scalar alpha, + Precision P = Precision::F32) override { auto a_gpu = import(A); NN_TRY_CHECK(a_gpu); @@ -1005,11 +932,9 @@ class GpuEngine final : public ComputeEngine // ── F32 路径（现状，零改动）───────────────────────────────────── if (P == Precision::F32) { - auto r = backend_->batched_matmul_gpu( - a_gpu->gpu_tensor(), b_gpu->gpu_tensor(), - static_cast(batch), - transA ? 1u : 0u, transB ? 1u : 0u, - static_cast(alpha)); + auto r = backend_->batched_matmul_gpu(a_gpu->gpu_tensor(), b_gpu->gpu_tensor(), + static_cast(batch), transA ? 1u : 0u, + transB ? 1u : 0u, static_cast(alpha)); NN_TRY_CHECK(r); return Tensor::from_gpu(std::move(*r)); } @@ -1017,16 +942,13 @@ class GpuEngine final : public ComputeEngine // ── F16 路径（Phase 2 in-kernel f16）────────────────────────────── // 同 matmul：两个操作数都 f16 → f16 存储版 batch GEMM（零 f32 副本）。 if (P == Precision::F16 && a_gpu->precision() == Precision::F16 && - b_gpu->precision() == Precision::F16 && - backend_->has(PipelineId::batched_matmul_f16)) + b_gpu->precision() == Precision::F16 && backend_->has(PipelineId::batched_matmul_f16)) { auto r = backend_->batched_matmul_gpu( - f16_view(*a_gpu), f16_view(*b_gpu), - static_cast(batch), transA ? 1u : 0u, transB ? 1u : 0u, - static_cast(alpha), /*f16_io=*/true); + f16_view(*a_gpu), f16_view(*b_gpu), static_cast(batch), transA ? 1u : 0u, + transB ? 1u : 0u, static_cast(alpha), /*f16_io=*/true); NN_TRY_CHECK(r); - return Tensor::from_gpu( - GpuTensorF16(r->shared_buffer(), r->rows(), r->cols())); + return Tensor::from_gpu(GpuTensorF16(r->shared_buffer(), r->rows(), r->cols())); } // ── F16 边界 cast 回退（无原生 f16 GEMM 时）────────────────────── @@ -1040,11 +962,9 @@ class GpuEngine final : public ComputeEngine NN_TRY_CHECK(a32_gpu); auto b32_gpu = import(*b32); NN_TRY_CHECK(b32_gpu); - auto r = backend_->batched_matmul_gpu( - a32_gpu->gpu_tensor(), b32_gpu->gpu_tensor(), - static_cast(batch), - transA ? 1u : 0u, transB ? 1u : 0u, - static_cast(alpha)); + auto r = backend_->batched_matmul_gpu(a32_gpu->gpu_tensor(), b32_gpu->gpu_tensor(), + static_cast(batch), transA ? 1u : 0u, + transB ? 1u : 0u, static_cast(alpha)); NN_TRY_CHECK(r); Tensor result_f32 = Tensor::from_gpu(std::move(*r)); return cast(std::move(result_f32), Precision::F16); @@ -1058,8 +978,7 @@ class GpuEngine final : public ComputeEngine // 免去新 buffer 分配 + 全量写出，消除优化器/梯度累积路径的分配风暴。 [[nodiscard]] Result add_inplace_impl(Tensor& A, const Tensor& B) override { - if (A.rows() != B.rows() || A.cols() != B.cols()) - NN_FAIL("add_inplace: shape mismatch"); + if (A.rows() != B.rows() || A.cols() != B.cols()) NN_FAIL("add_inplace: shape mismatch"); auto a_gpu = import(A); NN_TRY_CHECK(a_gpu); auto b_gpu = import(B); @@ -1069,8 +988,7 @@ class GpuEngine final : public ComputeEngine // · B 同 f16 → 零副本直加； // · B 是 f32（梯度混合精度）→ 只物化 B 的 f16 小副本再加； // · 无 f16 pipeline → 引擎内回退（A 抬 f32 加完写回，正确性不变）。 - if (a_gpu->precision() == Precision::F16 && - backend_->has(PipelineId::elementwise_v2_f16)) + if (a_gpu->precision() == Precision::F16 && backend_->has(PipelineId::elementwise_v2_f16)) { std::optional b16; if (b_gpu->precision() != Precision::F16) @@ -1082,13 +1000,11 @@ class GpuEngine final : public ComputeEngine const GpuTensor& b_ref = b16 ? f16_view(*b16) : f16_view(*b_gpu); const GpuTensor a_ref = f16_view(*a_gpu); // 局部保活（禁 &临时） const uint32_t count16 = static_cast(A.rows() * A.cols()); - auto r = backend_->elementwise_v2_gpu( - a_ref, &b_ref, nullptr, - count16, 1u, 0u, 0u, 0u, 0.0f, 0.0f, 0.0f, - &a_ref, /*f16_io=*/true); // BINARY, Add, 原地写回 A + auto r = backend_->elementwise_v2_gpu(a_ref, &b_ref, nullptr, count16, 1u, 0u, 0u, 0u, + 0.0f, 0.0f, 0.0f, &a_ref, + /*f16_io=*/true); // BINARY, Add, 原地写回 A NN_TRY_CHECK(r); - if (A.is_cpu()) - A = std::move(*a_gpu); + if (A.is_cpu()) A = std::move(*a_gpu); return {}; } if (a_gpu->precision() == Precision::F16) @@ -1105,10 +1021,8 @@ class GpuEngine final : public ComputeEngine } const GpuTensor& b_ref2 = b32 ? b32->gpu_tensor() : b_gpu->gpu_tensor(); const uint32_t cnt = static_cast(A.rows() * A.cols()); - auto r = backend_->elementwise_v2_gpu( - a32->gpu_tensor(), &b_ref2, nullptr, - cnt, 1u, 0u, 0u, 0u, 0.0f, 0.0f, 0.0f, - &a32->gpu_tensor()); + auto r = backend_->elementwise_v2_gpu(a32->gpu_tensor(), &b_ref2, nullptr, cnt, 1u, 0u, + 0u, 0u, 0.0f, 0.0f, 0.0f, &a32->gpu_tensor()); NN_TRY_CHECK(r); return cast_into(*a32, A); // f32 结果写回 A 的 f16 原存储 } @@ -1120,24 +1034,21 @@ class GpuEngine final : public ComputeEngine auto b32 = cast(*b_gpu, Precision::F32); NN_TRY_CHECK(b32); const uint32_t cnt = static_cast(A.rows() * A.cols()); - auto r = backend_->elementwise_v2_gpu( - a_gpu->gpu_tensor(), &b32->gpu_tensor(), nullptr, - cnt, 1u, 0u, 0u, 0u, 0.0f, 0.0f, 0.0f, - &a_gpu->gpu_tensor()); + auto r = backend_->elementwise_v2_gpu(a_gpu->gpu_tensor(), &b32->gpu_tensor(), nullptr, + cnt, 1u, 0u, 0u, 0u, 0.0f, 0.0f, 0.0f, + &a_gpu->gpu_tensor()); NN_TRY_CHECK(r); if (A.is_cpu()) A = std::move(*a_gpu); return {}; } const uint32_t count = static_cast(A.rows() * A.cols()); - auto r = backend_->elementwise_v2_gpu( - a_gpu->gpu_tensor(), &b_gpu->gpu_tensor(), nullptr, - count, 1u, 0u, 0u, 0u, 0.0f, 0.0f, 0.0f, - &a_gpu->gpu_tensor()); // BINARY, Add, 原地写回 A + auto r = backend_->elementwise_v2_gpu(a_gpu->gpu_tensor(), &b_gpu->gpu_tensor(), nullptr, + count, 1u, 0u, 0u, 0u, 0.0f, 0.0f, 0.0f, + &a_gpu->gpu_tensor()); // BINARY, Add, 原地写回 A NN_TRY_CHECK(r); // 原地模式下 A 的 buffer 已被更新；若 import 上传了新 Tensor（防御路径），替换 A - if (A.is_cpu()) - A = std::move(*a_gpu); + if (A.is_cpu()) A = std::move(*a_gpu); return {}; } @@ -1154,33 +1065,29 @@ class GpuEngine final : public ComputeEngine if (backend_->has(PipelineId::elementwise_v2_f16)) { const GpuTensor a_ref = f16_view(*a_gpu); // 局部保活（禁 &临时） - auto r = backend_->elementwise_v2_gpu( - a_ref, nullptr, nullptr, - cnt, 1u, 2u, 0u, 1u, static_cast(s), 0.0f, 0.0f, - &a_ref, /*f16_io=*/true); + auto r = backend_->elementwise_v2_gpu(a_ref, nullptr, nullptr, cnt, 1u, 2u, 0u, 1u, + static_cast(s), 0.0f, 0.0f, &a_ref, + /*f16_io=*/true); NN_TRY_CHECK(r); if (A.is_cpu()) A = std::move(*a_gpu); return {}; } auto a32 = cast(*a_gpu, Precision::F32); NN_TRY_CHECK(a32); - auto r = backend_->elementwise_v2_gpu( - a32->gpu_tensor(), nullptr, nullptr, - cnt, 1u, 2u, 0u, 1u, static_cast(s), 0.0f, 0.0f, - &a32->gpu_tensor()); + auto r = backend_->elementwise_v2_gpu(a32->gpu_tensor(), nullptr, nullptr, cnt, 1u, 2u, + 0u, 1u, static_cast(s), 0.0f, 0.0f, + &a32->gpu_tensor()); NN_TRY_CHECK(r); return cast_into(*a32, A); } const uint32_t count = static_cast(A.rows() * A.cols()); // BINARY, Mul, flags=1 (B is scalar), scalar_b = s - auto r = backend_->elementwise_v2_gpu( - a_gpu->gpu_tensor(), nullptr, nullptr, - count, 1u, 2u, 0u, 1u, static_cast(s), 0.0f, 0.0f, - &a_gpu->gpu_tensor()); // 原地写回 A + auto r = backend_->elementwise_v2_gpu(a_gpu->gpu_tensor(), nullptr, nullptr, count, 1u, 2u, + 0u, 1u, static_cast(s), 0.0f, 0.0f, + &a_gpu->gpu_tensor()); // 原地写回 A NN_TRY_CHECK(r); - if (A.is_cpu()) - A = std::move(*a_gpu); + if (A.is_cpu()) A = std::move(*a_gpu); return {}; } @@ -1200,8 +1107,7 @@ class GpuEngine final : public ComputeEngine NN_TRY_CHECK(r); // fill_zero 修改的是 GPU buffer 本身，Tensor 的 shared_ptr 不变 // 但若 import 上传了新 Tensor，需要替换 A - if (A.is_cpu()) - A = std::move(*a_gpu); + if (A.is_cpu()) A = std::move(*a_gpu); return {}; } @@ -1215,8 +1121,7 @@ class GpuEngine final : public ComputeEngine // 已是 f32 → 共享所有权直通（零拷贝）。 [[nodiscard]] Result cast_if_f16(const Tensor& t) { - if (t.precision() != Precision::F16) - return t; + if (t.precision() != Precision::F16) return t; return cast(t, Precision::F32); } @@ -1224,12 +1129,11 @@ class GpuEngine final : public ComputeEngine // half，消灭适配层全尺寸 f32 副本——输出 (rows*5/3/dk, seq) 是大头）； // 混合精度（如 boundary 仍 f32）→ 引擎内把 f16 侧 cast 到 f32 走 f32 // pipeline（正确性不变）。输出重贴 GpuTensorF16（缓冲按 2B/元素分配）。 - [[nodiscard]] Result scan_prefix_outer_impl( - const Tensor& K, const Tensor& V, const Tensor& P, const Tensor& R, - const Tensor& A0, const Tensor& B0, bool has_state, - std::size_t dk, std::size_t heads, bool causal, - const Tensor& boundary, bool has_bnd, - Precision prec = Precision::F32) override + [[nodiscard]] Result + scan_prefix_outer_impl(const Tensor& K, const Tensor& V, const Tensor& P, const Tensor& R, + const Tensor& A0, const Tensor& B0, bool has_state, std::size_t dk, + std::size_t heads, bool causal, const Tensor& boundary, bool has_bnd, + Precision prec = Precision::F32) override { NN_TRY(k, import(K)); NN_TRY(v, import(V)); @@ -1238,23 +1142,20 @@ class GpuEngine final : public ComputeEngine NN_TRY(a, import(A0)); NN_TRY(b, import(B0)); NN_TRY(bn, import(boundary)); - const bool all_f16 = k->precision() == Precision::F16 && - v->precision() == Precision::F16 && p->precision() == Precision::F16 && - r->precision() == Precision::F16 && a->precision() == Precision::F16 && - b->precision() == Precision::F16 && - bn->precision() == Precision::F16; + const bool all_f16 = k->precision() == Precision::F16 && v->precision() == Precision::F16 && + p->precision() == Precision::F16 && r->precision() == Precision::F16 && + a->precision() == Precision::F16 && b->precision() == Precision::F16 && + bn->precision() == Precision::F16; // 原生 f16 仅在请求输出也是 f16 时启用（输出契约：prec=F32 必须回 // f32——走下方引擎内 cast 路径，输出 f32 恰好免适配层落回） if (prec == Precision::F16 && all_f16 && backend_->has_scan_f16_pipelines()) { auto res = backend_->scan_prefix_outer_gpu( - f16_view(*k), f16_view(*v), f16_view(*p), f16_view(*r), - f16_view(*a), f16_view(*b), has_state, - static_cast(dk), static_cast(heads), causal, + f16_view(*k), f16_view(*v), f16_view(*p), f16_view(*r), f16_view(*a), f16_view(*b), + has_state, static_cast(dk), static_cast(heads), causal, f16_view(*bn), has_bnd, /*f16_io=*/true); NN_TRY_CHECK(res); - return Tensor::from_gpu( - GpuTensorF16(res->shared_buffer(), res->rows(), res->cols())); + return Tensor::from_gpu(GpuTensorF16(res->shared_buffer(), res->rows(), res->cols())); } // 引擎内 cast 回退（混合精度或无 f16 pipeline）：f16 侧抬 f32 NN_TRY(k32, cast_if_f16(*k)); @@ -1266,74 +1167,65 @@ class GpuEngine final : public ComputeEngine NN_TRY(n32, cast_if_f16(*bn)); auto res = backend_->scan_prefix_outer_gpu( k32->gpu_tensor(), v32->gpu_tensor(), p32->gpu_tensor(), r32->gpu_tensor(), - a32->gpu_tensor(), b32->gpu_tensor(), has_state, - static_cast(dk), static_cast(heads), causal, - n32->gpu_tensor(), has_bnd); + a32->gpu_tensor(), b32->gpu_tensor(), has_state, static_cast(dk), + static_cast(heads), causal, n32->gpu_tensor(), has_bnd); NN_TRY_CHECK(res); return to_prec(Tensor::from_gpu(std::move(*res)), prec); } - [[nodiscard]] Result scan_suffix_outer_impl( - const Tensor& D, const Tensor& X, const Tensor& Y, - std::size_t dk, std::size_t heads, bool causal, - const Tensor& boundary, bool has_bnd, - Precision prec = Precision::F32) override + [[nodiscard]] Result scan_suffix_outer_impl(const Tensor& D, const Tensor& X, + const Tensor& Y, std::size_t dk, + std::size_t heads, bool causal, + const Tensor& boundary, bool has_bnd, + Precision prec = Precision::F32) override { NN_TRY(d, import(D)); NN_TRY(x, import(X)); NN_TRY(y, import(Y)); NN_TRY(bn, import(boundary)); - const bool all_f16 = d->precision() == Precision::F16 && - x->precision() == Precision::F16 && y->precision() == Precision::F16 && - bn->precision() == Precision::F16; + const bool all_f16 = d->precision() == Precision::F16 && x->precision() == Precision::F16 && + y->precision() == Precision::F16 && bn->precision() == Precision::F16; if (prec == Precision::F16 && all_f16 && backend_->has_scan_f16_pipelines()) { auto res = backend_->scan_suffix_outer_gpu( - f16_view(*d), f16_view(*x), f16_view(*y), - static_cast(dk), static_cast(heads), causal, - f16_view(*bn), has_bnd, /*f16_io=*/true); + f16_view(*d), f16_view(*x), f16_view(*y), static_cast(dk), + static_cast(heads), causal, f16_view(*bn), has_bnd, /*f16_io=*/true); NN_TRY_CHECK(res); - return Tensor::from_gpu( - GpuTensorF16(res->shared_buffer(), res->rows(), res->cols())); + return Tensor::from_gpu(GpuTensorF16(res->shared_buffer(), res->rows(), res->cols())); } NN_TRY(d32, cast_if_f16(*d)); NN_TRY(x32, cast_if_f16(*x)); NN_TRY(y32, cast_if_f16(*y)); NN_TRY(n32, cast_if_f16(*bn)); auto res = backend_->scan_suffix_outer_gpu( - d32->gpu_tensor(), x32->gpu_tensor(), y32->gpu_tensor(), - static_cast(dk), static_cast(heads), causal, - n32->gpu_tensor(), has_bnd); + d32->gpu_tensor(), x32->gpu_tensor(), y32->gpu_tensor(), static_cast(dk), + static_cast(heads), causal, n32->gpu_tensor(), has_bnd); NN_TRY_CHECK(res); return to_prec(Tensor::from_gpu(std::move(*res)), prec); } - [[nodiscard]] Result outer_col_impl( - const Tensor& P, const Tensor& R, const Tensor& S, - std::size_t dk, bool has_scale, - Precision prec = Precision::F32) override + [[nodiscard]] Result outer_col_impl(const Tensor& P, const Tensor& R, const Tensor& S, + std::size_t dk, bool has_scale, + Precision prec = Precision::F32) override { NN_TRY(p, import(P)); NN_TRY(r, import(R)); NN_TRY(s, import(S)); - const bool all_f16 = p->precision() == Precision::F16 && - r->precision() == Precision::F16 && s->precision() == Precision::F16; - if (prec == Precision::F16 && all_f16 && - backend_->has(PipelineId::outer_col_f16)) - { - auto res = backend_->outer_col_gpu( - f16_view(*p), f16_view(*r), f16_view(*s), - static_cast(dk), has_scale, /*f16_io=*/true); + const bool all_f16 = p->precision() == Precision::F16 && r->precision() == Precision::F16 && + s->precision() == Precision::F16; + if (prec == Precision::F16 && all_f16 && backend_->has(PipelineId::outer_col_f16)) + { + auto res = + backend_->outer_col_gpu(f16_view(*p), f16_view(*r), f16_view(*s), + static_cast(dk), has_scale, /*f16_io=*/true); NN_TRY_CHECK(res); - return Tensor::from_gpu( - GpuTensorF16(res->shared_buffer(), res->rows(), res->cols())); + return Tensor::from_gpu(GpuTensorF16(res->shared_buffer(), res->rows(), res->cols())); } NN_TRY(p32, cast_if_f16(*p)); NN_TRY(r32, cast_if_f16(*r)); NN_TRY(s32, cast_if_f16(*s)); - auto res = backend_->outer_col_gpu( - p32->gpu_tensor(), r32->gpu_tensor(), s32->gpu_tensor(), - static_cast(dk), has_scale); + auto res = backend_->outer_col_gpu(p32->gpu_tensor(), r32->gpu_tensor(), s32->gpu_tensor(), + static_cast(dk), has_scale); NN_TRY_CHECK(res); return to_prec(Tensor::from_gpu(std::move(*res)), prec); } @@ -1345,8 +1237,8 @@ class GpuEngine final : public ComputeEngine // Phase C2：f16 输入 → f16 归约 pipeline（输入 float16_t 直读、f32 归约、 // 输出 f32 向量）——省掉"每个 f16 输入一份整份 f32 副本"的边界 cast。 // 设备无 16bit 存储时 pipeline 未创建 → 引擎内 cast 回退（正确性不变）。 - [[nodiscard]] Result reduce_f16_or_cast_( - const Tensor& A, uint32_t mode, uint32_t reduce_op) + [[nodiscard]] Result reduce_f16_or_cast_(const Tensor& A, uint32_t mode, + uint32_t reduce_op) { auto a_gpu = import(A); NN_TRY_CHECK(a_gpu); @@ -1356,7 +1248,7 @@ class GpuEngine final : public ComputeEngine // 的 f32 view = 断言失败 [gpu_tensor-fail]） if (backend_->has(PipelineId::reduce_f16)) return backend_->reduce_gpu(f16_view(*a_gpu), mode, reduce_op, - /*f16_in=*/true); + /*f16_in=*/true); // 引擎内边界 cast（无 f16 pipeline 的设备） auto a32 = cast(*a_gpu, Precision::F32); NN_TRY_CHECK(a32); @@ -1367,7 +1259,8 @@ class GpuEngine final : public ComputeEngine return backend_->reduce_gpu(a_gpu->gpu_tensor(), mode, reduce_op); } - [[nodiscard]] Result row_reduce_sum_impl(const Tensor& A, Precision P = Precision::F32) override + [[nodiscard]] Result row_reduce_sum_impl(const Tensor& A, + Precision P = Precision::F32) override { // 归约恒 f32 累加；输出按 P 落回（原生 f16 输入由 reduce_f16_or_cast_ 处理） auto r = reduce_f16_or_cast_(A, /*mode=*/0u, /*reduce_op=*/0u); @@ -1375,14 +1268,16 @@ class GpuEngine final : public ComputeEngine return to_prec(Tensor::from_gpu(std::move(*r)), P); } - [[nodiscard]] Result col_reduce_sum_impl(const Tensor& A, Precision P = Precision::F32) override + [[nodiscard]] Result col_reduce_sum_impl(const Tensor& A, + Precision P = Precision::F32) override { auto r = reduce_f16_or_cast_(A, /*mode=*/1u, /*reduce_op=*/0u); NN_TRY_CHECK(r); return to_prec(Tensor::from_gpu(std::move(*r)), P); } - [[nodiscard]] Result col_reduce_max_impl(const Tensor& A, Precision P = Precision::F32) override + [[nodiscard]] Result col_reduce_max_impl(const Tensor& A, + Precision P = Precision::F32) override { auto r = reduce_f16_or_cast_(A, /*mode=*/1u, /*reduce_op=*/1u); NN_TRY_CHECK(r); @@ -1408,7 +1303,7 @@ class GpuEngine final : public ComputeEngine // 录制期生命周期；实测表现为偶发错值，如 A+=B err=0.5）。 struct FusedInputs { - std::vector owners; + std::vector owners; std::vector bufs; }; @@ -1423,8 +1318,7 @@ class GpuEngine final : public ComputeEngine NN_TRY_CHECK(g); if (g->precision() == Precision::F16) out.bufs.push_back(&g->gpu_tensor().buffer()); - else - out.bufs.push_back(&g->gpu_tensor().buffer()); + else out.bufs.push_back(&g->gpu_tensor().buffer()); out.owners.push_back(std::move(*g)); // 保活（见上） } return out; @@ -1451,25 +1345,23 @@ class GpuEngine final : public ComputeEngine // ② 只由真正 dispatch 的调用方传 log_hit=true —— // `supports_expr_precision_variant` 只是查询（NVI 入口每次都调）， // 它打日志会让计数翻倍（实测 12040 行 = 6020 次真实分派 × 2）。 - [[nodiscard]] const nn::fused::FusedShader* find_prec_variant_( - const std::string& key, nn::ExprPrecSig psig, const nn::ExprSpec& spec, - bool log_hit = false) const + [[nodiscard]] const nn::fused::FusedShader* find_prec_variant_(const std::string& key, + nn::ExprPrecSig psig, + const nn::ExprSpec& spec, + bool log_hit = false) const { - if (psig == 0) - return nn::fused::find_fused(key); // V0：全 f32 结构键 + if (psig == 0) return nn::fused::find_fused(key); // V0：全 f32 结构键 const bool trace = log_hit && nn::dsl::env_flag("NN_PREC_TRACE"); - const auto hit = [&](const std::string& k) -> const nn::fused::FusedShader* { + const auto hit = [&](const std::string& k) -> const nn::fused::FusedShader* + { const auto* fs = nn::fused::find_fused(k); - if (!fs || !backend_->has_fused_shader(k)) - return nullptr; - if (trace) - std::fprintf(stderr, "[prec][alu-hit] %s\n", k.c_str()); + if (!fs || !backend_->has_fused_shader(k)) return nullptr; + if (trace) std::fprintf(stderr, "[prec][alu-hit] %s\n", k.c_str()); return fs; }; // V2b：native16 原生 f16 算术（运行期用真实 spec 复核结构谓词） if (nn::expr_prec_sig_native16(spec, psig)) - if (const auto* fs = hit(key + nn::EXPR_PREC_ALU_SUFFIX)) - return fs; + if (const auto* fs = hit(key + nn::EXPR_PREC_ALU_SUFFIX)) return fs; // V1t/V1s：storage-only **编译期特化**（热循环里没有 `prec` 三元）。 // 与 V1 分派变体数值逐字等价；只需要 16 位存储，无 shaderFloat16 的 // 设备（Maxwell）同样命中。优先于 V1。 @@ -1477,35 +1369,31 @@ class GpuEngine final : public ComputeEngine // pass 各跑 f32/f16 两遍，见 scan_exprs / ExprRegistry::add_variant）； // ② `key#s`（V1s）："全输入 f16 + 输出 f16"兜底——始终发射，覆盖 // 扫描没走到的路径（特化集合只是优化，`#x` 仍是语义兜底）。 - if (const auto* fs = hit(nn::expr_prec_sig_key(key, psig))) - return fs; + if (const auto* fs = hit(nn::expr_prec_sig_key(key, psig))) return fs; if (nn::expr_prec_sig_all_f16(spec, psig)) - if (const auto* fs = hit(key + nn::EXPR_PREC_STORAGE_SUFFIX)) - return fs; + if (const auto* fs = hit(key + nn::EXPR_PREC_STORAGE_SUFFIX)) return fs; // V1：运行期精度分派 const std::string xkey = key + nn::EXPR_PREC_DISPATCH_SUFFIX; - if (const auto* fs = hit(xkey)) - return fs; + if (const auto* fs = hit(xkey)) return fs; if (trace) - std::fprintf(stderr, - "[prec][fallback] %s 未预生成 → 回退基类边界 cast\n", + std::fprintf(stderr, "[prec][fallback] %s 未预生成 → 回退基类边界 cast\n", xkey.c_str()); return nullptr; } #endif // ══════════════════════════════════════════════════════════════════════ - [[nodiscard]] Result eval_expr_impl( - const ExprSpec& raw_spec, - std::span inputs, - std::size_t rows, std::size_t cols, - Precision P = Precision::F32) override + [[nodiscard]] Result eval_expr_impl(const ExprSpec& raw_spec, + std::span inputs, std::size_t rows, + std::size_t cols, + Precision P = Precision::F32) override { // 无注册表模式（NN_EXPR_SCAN 收集器 TU / 注册表缺失）：融合分支整体 // 不编译 → rows/cols 仅在 NN_FUSED_REGISTRY_EMBEDDED 内被使用 → // 显式消 unused-parameter（clang -Werror）。 #ifndef NN_FUSED_REGISTRY_EMBEDDED - (void)rows; (void)cols; + (void)rows; + (void)cols; #endif // ── fold 段：canonicalize 对 fold 恒等（expr_opt 入口 // early-return），scan/runtime 两端 key 同源 → 直接用 raw_spec 查表。 @@ -1534,8 +1422,7 @@ class GpuEngine final : public ComputeEngine // fold 的转置仍在 key 里（生成期定死）→ 此处恒 0（无该槽） nn::expr_spec_runtime_matmul_k(raw_spec), nn::expr_spec_runtime_matmul_batch(raw_spec), - /*matmul_trans=*/0u, - nn::expr_spec_runtime_fold_k(raw_spec), fout_f16, fpsig); + /*matmul_trans=*/0u, nn::expr_spec_runtime_fold_k(raw_spec), fout_f16, fpsig); NN_TRY_CHECK(out); // f16 输出：按 2B/元素分配后重贴 GpuTensorF16（同非 fold 路径） if (fout_f16) @@ -1554,7 +1441,9 @@ class GpuEngine final : public ComputeEngine NN_TRY_CHECK(r); return to_prec(std::move(*r), P); } - NN_FAIL("GpuEngine::eval_expr: fold 表达式未命中 AOT 融合 shader" "（闭合世界）；请将该 fold 结构纳入 scan_exprs；key=" + fkey); + NN_FAIL("GpuEngine::eval_expr: fold 表达式未命中 AOT 融合 shader" + "（闭合世界）；请将该 fold 结构纳入 scan_exprs；key=" + + fkey); } // ── canonical IR：canonicalize 为引擎内部优化（IR-A/IR-B）， // key 与 shader 合成两端一致；dispatch 用 canonical 的 consts ── @@ -1573,15 +1462,15 @@ class GpuEngine final : public ComputeEngine // + 目标输出精度取变体（find_prec_variant_）。命中即 shader 直接半精度 // 读/写，无需边界 cast；未命中则退回全 f32 key —— 此时输入必须已全 f32 // （由基类 NVI 入口 cast）。 - const nn::fused::FusedShader* fs = - find_prec_variant_(key, psig, spec, /*log_hit=*/true); + const nn::fused::FusedShader* fs = find_prec_variant_(key, psig, spec, /*log_hit=*/true); if (fs && backend_->has_fused_shader(fs->key)) { // 命中：收集 GPU 输入（同一 buffer 可重复绑定，如 RoPE 的 q×2） // GpuTensor 内部为 shared_ptr，拷贝即共享，零成本 auto fi_r = fused_buffers_(inputs); NN_TRY_CHECK(fi_r); - std::vector& gpu_inputs = fi_r->bufs; // owners 随 fi_r 存活到本作用域末 + std::vector& gpu_inputs = + fi_r->bufs; // owners 随 fi_r 存活到本作用域末 // 运行时视图参数（RowMod 周期 / RotateHalf 块大小）： // 同结构不同形状共享一个融合 shader，按实际 spec 填充 vp 槽 const auto vp = nn::expr_spec_runtime_view_params(spec); @@ -1594,13 +1483,11 @@ class GpuEngine final : public ComputeEngine auto out = backend_->run_fused_gpu( fs->key, gpu_inputs, spec.consts, rows, cols, /*vector_out=*/false, vp, spec.rparams, /*output_override=*/nullptr, - nn::expr_spec_runtime_matmul_k(spec), - nn::expr_spec_runtime_matmul_batch(spec), + nn::expr_spec_runtime_matmul_k(spec), nn::expr_spec_runtime_matmul_batch(spec), nn::expr_spec_runtime_matmul_trans(spec), /*fold_k=*/std::nullopt, out_f16, psig); NN_TRY_CHECK(out); - if (out_f16) - return Tensor::from_gpu(GpuTensorF16(out->shared_buffer(), rows, cols)); + if (out_f16) return Tensor::from_gpu(GpuTensorF16(out->shared_buffer(), rows, cols)); return Tensor::from_gpu(std::move(*out)); } #endif @@ -1611,8 +1498,7 @@ class GpuEngine final : public ComputeEngine // 入口代为处理（能力查询 + 二次归因）；F2 修复后归因/选路只做一次。 if (psig != 0) { - if (prec_trace_enabled_()) - trace_miss_(raw_spec, inputs, P); + if (prec_trace_enabled_()) trace_miss_(raw_spec, inputs, P); auto in32 = to_f32_all(inputs); NN_TRY_CHECK(in32); auto r = eval_expr_impl(raw_spec, *in32, rows, cols, Precision::F32); @@ -1622,7 +1508,10 @@ class GpuEngine final : public ComputeEngine // ── 闭合世界：未命中任何 AOT 融合 shader → 硬报错（绝不静默回退） ── // 带上 key：闭合世界报错必须可定位——key 是登记/查表两侧的唯一标识， // 可用来查 fused_registry.hpp 里有没有该条目、以及它属于哪个 Layer 路径。 - NN_FAIL("GpuEngine::eval_expr: 未找到该内联表达式的 AOT 融合 shader（闭合世界）；" "请将对应表达式纳入构建期扫描（scan_exprs dry-run 需覆盖该 Layer 路径）；" "key=" + key); + NN_FAIL("GpuEngine::eval_expr: 未找到该内联表达式的 AOT 融合 shader（闭合世界）；" + "请将对应表达式纳入构建期扫描（scan_exprs dry-run 需覆盖该 Layer 路径）；" + "key=" + + key); } // ── 能力摘要（唯一对外通道；只披露事实，不用于选路）────────────────── @@ -1635,8 +1524,8 @@ class GpuEngine final : public ComputeEngine i.native_f16_move = true; // 字节级搬运 + 引擎内 cast 回退 i.native_f16_reduce = backend_->has(PipelineId::reduce_f16); i.fused_prec_variant = true; // 有带类型融合 shader 机制（AOT 闭合世界） - i.ternary_gemm = backend_->has(PipelineId::trit_pack) && - backend_->has(PipelineId::matmul_q); + i.ternary_gemm = + backend_->has(PipelineId::trit_pack) && backend_->has(PipelineId::matmul_q); i.int8_dot = backend_->device().has_int8_dot(); i.int8_dot_accelerated = backend_->device().int8_dot_accelerated(); return i; @@ -1645,16 +1534,16 @@ class GpuEngine final : public ComputeEngine // ── 归约向量原生形状输出（LayerNorm/RMSNorm 小向量缓存） ──────── // 与 eval_expr 相同，但以 vector_out=1 调度归约融合 shader（thread 0 写 // (rows,1)/(1,cols) 归约向量，不写全尺寸广播）。 - [[nodiscard]] Result eval_expr_reduce_impl( - const ExprSpec& raw_spec, - std::span inputs, - std::size_t rows, std::size_t cols, - Precision P = Precision::F32) override + [[nodiscard]] Result eval_expr_reduce_impl(const ExprSpec& raw_spec, + std::span inputs, + std::size_t rows, std::size_t cols, + Precision P = Precision::F32) override { // 无注册表模式：融合分支不编译 → rows/cols 仅在 // NN_FUSED_REGISTRY_EMBEDDED 内被使用 → 消 unused-parameter。 #ifndef NN_FUSED_REGISTRY_EMBEDDED - (void)rows; (void)cols; + (void)rows; + (void)cols; #endif // fold 形态只经 eval_expr（PC 需 fold_k，本入口不传）——显式拒绝， // 否则落到 run_fused_gpu 的通用缺参错误，误导排查方向 @@ -1673,10 +1562,11 @@ class GpuEngine final : public ComputeEngine // 小向量用 dsl::compute 施加后处理（LayerNorm/RMSNorm 即此写法）。 // 注意：CPU 端已支持该形态（eval_expr_impl 的按指令下标反向切片）， // 生成器补齐该形态后可在此放开。 - if (!spec.instrs.empty() && - !expr_op_is_reduce(static_cast(spec.instrs.back().op))) + if (!spec.instrs.empty() && !expr_op_is_reduce(static_cast(spec.instrs.back().op))) { - NN_FAIL("GpuEngine::eval_expr_reduce: 归约向量输出的表达式末指令必须是归约指令" "（归约后逐元素后处理在归约融合 shader 中尚未正确实现）；" "请拆成两步：compute_reduce 取归约向量，再用 dsl::compute 做后处理"); + NN_FAIL("GpuEngine::eval_expr_reduce: 归约向量输出的表达式末指令必须是归约指令" + "（归约后逐元素后处理在归约融合 shader 中尚未正确实现）；" + "请拆成两步：compute_reduce 取归约向量，再用 dsl::compute 做后处理"); } const std::string key = nn::expr_spec_key(spec); @@ -1684,22 +1574,21 @@ class GpuEngine final : public ComputeEngine #ifdef NN_FUSED_REGISTRY_EMBEDDED // 变体优先（`#a` native16 → `#x` 运行期分派）：与 eval_expr 同款 // (key,sig) 匹配；归约形态的 native16 谓词恒 false → 自动落 `#x` - const nn::fused::FusedShader* fs = (psig != 0) - ? find_prec_variant_(key, psig, spec, /*log_hit=*/true) - : nn::fused::find_fused(key); + const nn::fused::FusedShader* fs = + (psig != 0) ? find_prec_variant_(key, psig, spec, /*log_hit=*/true) + : nn::fused::find_fused(key); if (fs && backend_->has_fused_shader(fs->key)) { auto fi_r = fused_buffers_(inputs); NN_TRY_CHECK(fi_r); - std::vector& gpu_inputs = fi_r->bufs; // owners 随 fi_r 存活到本作用域末 + std::vector& gpu_inputs = + fi_r->bufs; // owners 随 fi_r 存活到本作用域末 const auto vp = nn::expr_spec_runtime_view_params(spec); const bool out_f16 = nn::expr_variant_out_f16(fs->prec_sig, psig); auto out = backend_->run_fused_gpu( - fs->key, gpu_inputs, spec.consts, rows, cols, /*vector_out=*/true, vp, - spec.rparams, + fs->key, gpu_inputs, spec.consts, rows, cols, /*vector_out=*/true, vp, spec.rparams, /*output_override=*/nullptr, nn::expr_spec_runtime_matmul_k(spec), - nn::expr_spec_runtime_matmul_batch(spec), - nn::expr_spec_runtime_matmul_trans(spec), + nn::expr_spec_runtime_matmul_batch(spec), nn::expr_spec_runtime_matmul_trans(spec), std::nullopt, out_f16, psig); NN_TRY_CHECK(out); if (out_f16) @@ -1716,8 +1605,7 @@ class GpuEngine final : public ComputeEngine // 变体未预生成 → 引擎内边界 cast 回退（抬 f32 重入；F2 修复） if (psig != 0) { - if (prec_trace_enabled_()) - trace_miss_(raw_spec, inputs, P); + if (prec_trace_enabled_()) trace_miss_(raw_spec, inputs, P); auto in32 = to_f32_all(inputs); NN_TRY_CHECK(in32); auto r = eval_expr_reduce_impl(raw_spec, *in32, rows, cols, Precision::F32); @@ -1725,7 +1613,10 @@ class GpuEngine final : public ComputeEngine return to_prec(std::move(*r), P); } // ── 闭合世界：未命中归约融合 shader → 硬报错（绝不静默回退） ── - NN_FAIL("GpuEngine::eval_expr_reduce: 未找到该归约表达式的 AOT 融合 shader（闭合世界）；" "请将对应表达式纳入构建期扫描（scan_exprs dry-run 需覆盖该 Layer 路径）；" "key=" + key); + NN_FAIL("GpuEngine::eval_expr_reduce: 未找到该归约表达式的 AOT 融合 shader（闭合世界）；" + "请将对应表达式纳入构建期扫描（scan_exprs dry-run 需覆盖该 Layer 路径）；" + "key=" + + key); } // ── 目标传递（destination-passing）：结果直接写回已有张量 ───────────── @@ -1733,13 +1624,11 @@ class GpuEngine final : public ComputeEngine // 若表达式引用了 leaf(dst) 则同一 buffer 同时作为 readonly 输入绑定： // 逐元素同索引"先读后写"，无跨调用危害（与 GPU 原地原语的既有做法一致）。 // 语义与限制同 CpuEngine::eval_expr_into（仅逐元素表达式；无分配）。 - [[nodiscard]] Result eval_expr_into_impl( - const ExprSpec& raw_spec, - std::span inputs, - std::size_t rows, std::size_t cols, Tensor& dst) override + [[nodiscard]] Result eval_expr_into_impl(const ExprSpec& raw_spec, + std::span inputs, std::size_t rows, + std::size_t cols, Tensor& dst) override { - if (dst.rows() != rows || dst.cols() != cols) - NN_FAIL("eval_expr_into: dst shape mismatch"); + if (dst.rows() != rows || dst.cols() != cols) NN_FAIL("eval_expr_into: dst shape mismatch"); // fold 形态只经 eval_expr（PC 需 fold_k，本入口不传）——显式拒绝 if (raw_spec.fold) NN_FAIL("GpuEngine::eval_expr_into: fold 表达式请走 eval_expr（需 fold_k 形态参数）"); @@ -1748,17 +1637,17 @@ class GpuEngine final : public ComputeEngine NN_FAIL("eval_expr_into: 仅支持逐元素表达式（无归约）"); const std::string key = nn::expr_spec_key(spec); // 变体优先（`#a` native16 → `#x` 运行期分派）：目标精度 = dst.precision() - const nn::ExprPrecSig psig = - nn::expr_prec_sig_of(inputs, dst.precision()); + const nn::ExprPrecSig psig = nn::expr_prec_sig_of(inputs, dst.precision()); #ifdef NN_FUSED_REGISTRY_EMBEDDED - const nn::fused::FusedShader* fs = (psig != 0) - ? find_prec_variant_(key, psig, spec, /*log_hit=*/true) - : nn::fused::find_fused(key); + const nn::fused::FusedShader* fs = + (psig != 0) ? find_prec_variant_(key, psig, spec, /*log_hit=*/true) + : nn::fused::find_fused(key); if (fs && backend_->has_fused_shader(fs->key)) { auto fi_r = fused_buffers_(inputs); NN_TRY_CHECK(fi_r); - std::vector& gpu_inputs = fi_r->bufs; // owners 随 fi_r 存活到本作用域末 + std::vector& gpu_inputs = + fi_r->bufs; // owners 随 fi_r 存活到本作用域末 auto dst_gpu = import(dst); NN_TRY_CHECK(dst_gpu); // ── output_override 必须按目标存储精度取视图 ────────────────── @@ -1773,29 +1662,28 @@ class GpuEngine final : public ComputeEngine if (dst_gpu->precision() == Precision::F16) { if (!dst_gpu->gpu_shared()) - NN_FAIL("eval_expr_into: f16 目标缺少 f16 GPU 存储" "（precision 标签与 variant 存储不一致）"); + NN_FAIL("eval_expr_into: f16 目标缺少 f16 GPU 存储" + "（precision 标签与 variant 存储不一致）"); dst_view = f16_view(*dst_gpu); dst_override = &*dst_view; } else { if (!dst_gpu->gpu_shared()) - NN_FAIL("eval_expr_into: f32 目标缺少 f32 GPU 存储" "（precision 标签与 variant 存储不一致）"); + NN_FAIL("eval_expr_into: f32 目标缺少 f32 GPU 存储" + "（precision 标签与 variant 存储不一致）"); dst_override = &dst_gpu->gpu_tensor(); } const auto vp = nn::expr_spec_runtime_view_params(spec); auto out = backend_->run_fused_gpu( fs->key, gpu_inputs, spec.consts, rows, cols, /*vector_out=*/false, vp, - spec.rparams, dst_override, - nn::expr_spec_runtime_matmul_k(spec), - nn::expr_spec_runtime_matmul_batch(spec), - nn::expr_spec_runtime_matmul_trans(spec), + spec.rparams, dst_override, nn::expr_spec_runtime_matmul_k(spec), + nn::expr_spec_runtime_matmul_batch(spec), nn::expr_spec_runtime_matmul_trans(spec), /*fold_k=*/std::nullopt, /*out_f16=*/false, psig); NN_TRY_CHECK(out); // dst 原为 CPU staging 时，import 上传了新 buffer（结果在它上面） // → 用 upload 后的张量替换 dst，保证调用方看到更新后的数据 - if (dst.is_cpu()) - dst = std::move(*dst_gpu); + if (dst.is_cpu()) dst = std::move(*dst_gpu); return {}; } #endif @@ -1803,8 +1691,7 @@ class GpuEngine final : public ComputeEngine // cast_into 写回 dst 原存储；F2 修复） if (psig != 0) { - if (prec_trace_enabled_()) - trace_miss_(raw_spec, inputs, dst.precision()); + if (prec_trace_enabled_()) trace_miss_(raw_spec, inputs, dst.precision()); auto in32 = to_f32_all(inputs); NN_TRY_CHECK(in32); auto tmp = create_tensor(rows, cols, Precision::F32); @@ -1813,25 +1700,31 @@ class GpuEngine final : public ComputeEngine return cast_into(tmp, dst); } // ── 闭合世界：未命中 AOT 融合 shader → 硬报错（绝不静默回退） ── - NN_FAIL("GpuEngine::eval_expr_into: 未找到该表达式的 AOT 融合 shader（闭合世界）；" "请将对应表达式纳入构建期扫描（scan_exprs dry-run 需覆盖该 Layer 路径）；" "key=" + key); + NN_FAIL("GpuEngine::eval_expr_into: 未找到该表达式的 AOT 融合 shader（闭合世界）；" + "请将对应表达式纳入构建期扫描（scan_exprs dry-run 需覆盖该 Layer 路径）；" + "key=" + + key); } // ── P2 多输出（IR-C 写穿物化）──────────────────────────────────────── // outs 全部为既有 f32 存储（主输出 + extras）；选路 = V0（NVI 已验全 f32， // extras 结构只生成 V0 变体）。输出视图经 import 取 GpuTensor 句柄， // 以 output_override + extra_outs 交给 run_fused_gpu（无分配路径）。 - [[nodiscard]] Result eval_expr_multi_into_impl( - const ExprSpec& raw_spec, std::span inputs, - std::size_t rows, std::size_t cols, std::span outs) override + [[nodiscard]] Result eval_expr_multi_into_impl(const ExprSpec& raw_spec, + std::span inputs, + std::size_t rows, std::size_t cols, + std::span outs) override { - if (raw_spec.fold || raw_spec.matmul || - expr_spec_reduce_axis(raw_spec) != -1) + if (raw_spec.fold || raw_spec.matmul || expr_spec_reduce_axis(raw_spec) != -1) NN_FAIL("eval_expr_multi_into: 仅支持纯逐元素 spec"); const ExprSpec spec = nn::canonicalize_expr_spec(raw_spec); const std::string key = nn::expr_spec_key(spec); #ifndef NN_FUSED_REGISTRY_EMBEDDED // 无注册表模式（NN_EXPR_SCAN 收集器 TU）：融合分支整体不编译 - (void)inputs; (void)rows; (void)cols; (void)outs; + (void)inputs; + (void)rows; + (void)cols; + (void)outs; #endif #ifdef NN_FUSED_REGISTRY_EMBEDDED const nn::fused::FusedShader* fs = nn::fused::find_fused(key); @@ -1854,19 +1747,18 @@ class GpuEngine final : public ComputeEngine } std::vector extra_ptrs(ptrs.begin() + 1, ptrs.end()); const auto vp = nn::expr_spec_runtime_view_params(spec); - auto r = backend_->run_fused_gpu( - fs->key, fi_r->bufs, spec.consts, rows, cols, - /*vector_out=*/false, vp, spec.rparams, ptrs[0], - /*matmul_k=*/std::nullopt, /*matmul_batch=*/1u, - /*matmul_trans=*/0u, /*fold_k=*/std::nullopt, - /*out_f16=*/false, /*prec_sig=*/0, - extra_ptrs); + auto r = backend_->run_fused_gpu(fs->key, fi_r->bufs, spec.consts, rows, cols, + /*vector_out=*/false, vp, spec.rparams, ptrs[0], + /*matmul_k=*/std::nullopt, /*matmul_batch=*/1u, + /*matmul_trans=*/0u, /*fold_k=*/std::nullopt, + /*out_f16=*/false, /*prec_sig=*/0, extra_ptrs); NN_TRY_CHECK(r); return {}; } #endif NN_FAIL("GpuEngine::eval_expr_multi_into: 未命中 AOT 融合 shader（闭合世界）；" - "请将 begin_expr/end_expr 段纳入构建期扫描（scan_exprs）；key=" + key); + "请将 begin_expr/end_expr 段纳入构建期扫描（scan_exprs）；key=" + + key); } private: @@ -1898,8 +1790,7 @@ class GpuEngine final : public ComputeEngine plans.reserve(kernels.size()); for (const auto& k : kernels) plans.push_back(plan_from_kernel(k)); - if (gcache.size() >= fused::GRAPH_PLAN_CACHE_MAX) - gcache.clear(); + if (gcache.size() >= fused::GRAPH_PLAN_CACHE_MAX) gcache.clear(); gcache.emplace(gkey, std::move(plans)); } @@ -1921,16 +1812,14 @@ class GpuEngine final : public ComputeEngine bool progressed = false; for (std::size_t ki = 0; ki < kernels.size(); ++ki) { - if (done[ki]) - continue; + if (done[ki]) continue; bool ready = true; for (const int m : kernels[ki].members) { const auto& nd = g.nodes[static_cast(m)]; for (const int d : nd.dep_of_input) { - if (d < 0) - continue; + if (d < 0) continue; const int kd = kernel_of[static_cast(d)]; if (kd >= 0 && kd != static_cast(ki) && !done[static_cast(kd)]) @@ -1939,8 +1828,7 @@ class GpuEngine final : public ComputeEngine break; } } - if (!ready) - break; + if (!ready) break; } if (ready) { @@ -1993,14 +1881,14 @@ class GpuEngine final : public ComputeEngine struct DispatchItem { std::vector ins; // 与 spec.views 同序 - Tensor out; // 主输出（tail 占位/目标传递 dst） + Tensor out; // 主输出（tail 占位/目标传递 dst） std::vector extras; // P2 写穿的其余成员输出（成员序） std::vector in_tok; // 输入存储身份（hazard 判定） std::vector out_tok; // 物化输出存储身份 - bool degraded = false; // 复合降级 = 逐成员派发 - bool batchable = false; // P3 批量资格（含 `key#b` 已登记） + bool degraded = false; // 复合降级 = 逐成员派发 + bool batchable = false; // P3 批量资格（含 `key#b` 已登记） std::string bkey; // 组键（结构 key + 参数值序列） - ExprSpec bspec; // canonical spec（批量派发参数来源） + ExprSpec bspec; // canonical spec（批量派发参数来源） ExprPrecSig b_sig = 0; // 组内统一精度签名（0 = 全 f32） std::string b_prefix; // 批量变体前缀（不含 `#b`） }; @@ -2032,16 +1920,14 @@ class GpuEngine final : public ComputeEngine NN_FAIL("GpuEngine::end_expr: tail 占位缺失（图 IR 状态损坏）"); d.out = out_it->second; // 复合降级判定（与原循环同判据：写穿前置条件不满足即逐成员派发） - d.degraded = k.members.size() > 1 && - !composite_kernel_available_(k, d.ins, d.out); + d.degraded = k.members.size() > 1 && !composite_kernel_available_(k, d.ins, d.out); // P2 写穿成员输出（成员序，tail 之后；派发输出槽 [1..]，与 // spec.extras 同序——fuse 构造时即按 members 去尾序登记） if (!k.spec.extras.empty() && !d.degraded) { for (const int m : k.members) { - if (m == k.tail) - continue; + if (m == k.tail) continue; const auto mo = g.node_outputs.find(m); if (mo == g.node_outputs.end()) NN_FAIL("GpuEngine::end_expr: 写穿成员占位缺失（图 IR 状态损坏）"); @@ -2059,8 +1945,7 @@ class GpuEngine final : public ComputeEngine d.out_tok.push_back(tensor_storage_token_(mo->second)); } } - else - d.out_tok.push_back(tensor_storage_token_(d.out)); + else d.out_tok.push_back(tensor_storage_token_(d.out)); d.in_tok.reserve(d.ins.size()); for (const Tensor& t : d.ins) d.in_tok.push_back(tensor_storage_token_(t)); @@ -2068,11 +1953,9 @@ class GpuEngine final : public ComputeEngine // `key#b` 已登记（无 BDA 时注册期即跳过 → 天然不可批量）。 // bspec 与 run_graph_kernel_into_ 的选路同源（fold 原样、其余 // canonicalize）。 - const ExprSpec bspec = k.spec.fold - ? k.spec : nn::canonicalize_expr_spec(k.spec); + const ExprSpec bspec = k.spec.fold ? k.spec : nn::canonicalize_expr_spec(k.spec); d.bspec = bspec; - if (!d.degraded && !k.vector_out && nn::expr_spec_batchable(bspec) && - batch_enabled && + if (!d.degraded && !k.vector_out && nn::expr_spec_batchable(bspec) && batch_enabled && (k.rows == 0 || k.cols <= UINT32_MAX / k.rows)) { // 组内统一的精度签名：输入 + 主输出（extras 不在签名编码里 @@ -2082,17 +1965,14 @@ class GpuEngine final : public ComputeEngine for (const Tensor& t : d.ins) gpu_ok = gpu_ok && t.valid() && t.is_gpu(); for (const Tensor& t : d.extras) - gpu_ok = gpu_ok && t.valid() && t.is_gpu() && - t.precision() == Precision::F32; + gpu_ok = gpu_ok && t.valid() && t.is_gpu() && t.precision() == Precision::F32; gpu_ok = gpu_ok && d.out.valid() && d.out.is_gpu(); - const nn::ExprPrecSig b_sig = - (gpu_ok && d.ins.size() < 16u) - ? nn::expr_prec_sig_of(d.ins, d.out.precision()) - : 0u; - const std::string b_prefix = - (gpu_ok && (b_sig == 0u || d.extras.empty())) - ? batch_variant_prefix_(bspec, b_sig) - : std::string{}; + const nn::ExprPrecSig b_sig = (gpu_ok && d.ins.size() < 16u) + ? nn::expr_prec_sig_of(d.ins, d.out.precision()) + : 0u; + const std::string b_prefix = (gpu_ok && (b_sig == 0u || d.extras.empty())) + ? batch_variant_prefix_(bspec, b_sig) + : std::string{}; if (!b_prefix.empty()) { d.batchable = true; @@ -2114,9 +1994,9 @@ class GpuEngine final : public ComputeEngine constexpr std::size_t kBatchCap = 32; struct BatchGroup { - std::string key; + std::string key; std::vector members; // items 下标（列表序） - std::size_t first = 0; // 首成员下标 = 派发位置 + std::size_t first = 0; // 首成员下标 = 派发位置 }; std::vector groups; std::vector group_of(items.size(), -1); @@ -2128,12 +2008,11 @@ class GpuEngine final : public ComputeEngine }; const auto hazard = [](const DispatchItem& a, const DispatchItem& b) { - const auto hit = [](const std::vector& x, - const std::vector& y) + const auto hit = + [](const std::vector& x, const std::vector& y) { for (const void* p : x) - if (p != nullptr && std::find(y.begin(), y.end(), p) != y.end()) - return true; + if (p != nullptr && std::find(y.begin(), y.end(), p) != y.end()) return true; return false; }; return hit(a.out_tok, b.in_tok) || hit(a.out_tok, b.out_tok) || @@ -2142,28 +2021,27 @@ class GpuEngine final : public ComputeEngine for (std::size_t xi = 0; xi < items.size(); ++xi) { const DispatchItem& x = items[xi]; - if (!x.batchable) - continue; + if (!x.batchable) continue; int chosen = -1; for (std::size_t gi = 0; gi < groups.size() && chosen < 0; ++gi) { const BatchGroup& gr = groups[gi]; - if (gr.key != x.bkey || gr.members.size() >= kBatchCap) - continue; + if (gr.key != x.bkey || gr.members.size() >= kBatchCap) continue; bool ok = true; for (const std::size_t m : gr.members) - if (hazard(x, items[m])) { ok = false; break; } - if (!ok) - continue; + if (hazard(x, items[m])) + { + ok = false; + break; + } + if (!ok) continue; for (std::size_t yi = 0; yi < xi && ok; ++yi) { if (group_of[yi] == static_cast(gi)) continue; // 同组成员与 x 已由上面排除 hazard - if (hazard(x, items[yi]) && action_pos(yi) >= gr.first) - ok = false; + if (hazard(x, items[yi]) && action_pos(yi) >= gr.first) ok = false; } - if (ok) - chosen = static_cast(gi); + if (ok) chosen = static_cast(gi); } if (chosen >= 0) { @@ -2186,13 +2064,11 @@ class GpuEngine final : public ComputeEngine if (gi >= 0) { const BatchGroup& gr = groups[static_cast(gi)]; - if (gr.members.front() != ki) - continue; // 已随本组首成员一次性派发 + if (gr.members.front() != ki) continue; // 已随本组首成员一次性派发 if (gr.members.size() >= 2) { if (nn::dsl::env_flag("NN_IRC_TRACE")) - std::fprintf(stderr, - "[ir-c][batch] %zu 实例 → 1 派发 %s\n", + std::fprintf(stderr, "[ir-c][batch] %zu 实例 → 1 派发 %s\n", gr.members.size(), gr.key.c_str()); std::vector insts; insts.reserve(gr.members.size()); @@ -2203,8 +2079,8 @@ class GpuEngine final : public ComputeEngine const auto gpu_buf = [](const Tensor& t) -> const GpuBuffer* { return t.precision() == Precision::F16 - ? &t.gpu_tensor().buffer() - : &t.gpu_tensor().buffer(); + ? &t.gpu_tensor().buffer() + : &t.gpu_tensor().buffer(); }; for (const std::size_t m : gr.members) { @@ -2217,30 +2093,27 @@ class GpuEngine final : public ComputeEngine bi.extra_outs.reserve(dm.extras.size()); for (const Tensor& t : dm.extras) bi.extra_outs.push_back(gpu_buf(t)); - bi.count = static_cast( - kernels[m].rows * kernels[m].cols); + bi.count = static_cast(kernels[m].rows * kernels[m].cols); bi.cols = static_cast(kernels[m].cols); insts.push_back(std::move(bi)); } const ExprSpec& s = items[gr.members.front()].bspec; auto br = backend_->run_fused_gpu_batch( - items[gr.members.front()].b_prefix + nn::EXPR_BATCH_SUFFIX, - insts, - s.consts, nn::expr_spec_runtime_view_params(s), s.rparams); + items[gr.members.front()].b_prefix + nn::EXPR_BATCH_SUFFIX, insts, s.consts, + nn::expr_spec_runtime_view_params(s), s.rparams); NN_TRY_CHECK(br); continue; } // 单成员组 = 常规逐 kernel（单 kernel 现有语义不变） } // ── 通用降级可观测（NN_IRC_TRACE=1）：逐元素候选但未成组 ────── - if (nn::dsl::env_flag("NN_IRC_TRACE") && !d.degraded && - !k.vector_out && nn::expr_spec_batchable(d.bspec)) - std::fprintf(stderr, - "[ir-c][batch-degrade] %s → 逐 kernel 派发（%s）\n", + if (nn::dsl::env_flag("NN_IRC_TRACE") && !d.degraded && !k.vector_out && + nn::expr_spec_batchable(d.bspec)) + std::fprintf(stderr, "[ir-c][batch-degrade] %s → 逐 kernel 派发（%s）\n", nn::expr_spec_key(d.bspec).c_str(), !batch_enabled ? "NN_IRC_NO_BATCH" : !d.batchable ? "key#b 未登记/无 BDA" - : "同签名未集齐 2 实例（依赖/上限）"); + : "同签名未集齐 2 实例（依赖/上限）"); // ── 复合 kernel 的闭合世界降级（通用机制，非单算子特例）───────── // 融合**分组本身依赖精度**：P2 写穿组要求成员全 F32、P1 链要求成员 @@ -2269,8 +2142,7 @@ class GpuEngine final : public ComputeEngine if (k.spec.extras.empty()) { - auto r = run_graph_kernel_into_(k.spec, d.ins, k.rows, k.cols, - k.vector_out, d.out); + auto r = run_graph_kernel_into_(k.spec, d.ins, k.rows, k.cols, k.vector_out, d.out); NN_TRY_CHECK(r); } else @@ -2293,11 +2165,10 @@ class GpuEngine final : public ComputeEngine for (auto it = g.node_outputs.begin(); it != g.node_outputs.end();) { const int nd = it->first; - if (nd >= 0 && static_cast(nd) < keep.size() - && !keep[static_cast(nd)]) + if (nd >= 0 && static_cast(nd) < keep.size() && + !keep[static_cast(nd)]) it = g.node_outputs.erase(it); - else - ++it; + else ++it; } return {}; } @@ -2307,9 +2178,10 @@ class GpuEngine final : public ComputeEngine // （vector_out）/ 精度变体（`#a`/`#x`/V0）——与 eval_expr_impl / // eval_expr_into_impl 的选路同源，output_override 精度视图同 // eval_expr_into_impl（f16 目标必须经 f16_view 包 f16 buffer）。 - [[nodiscard]] Result run_graph_kernel_into_( - const ExprSpec& raw_spec, std::span inputs, - std::size_t rows, std::size_t cols, bool vector_out, Tensor& out) + [[nodiscard]] Result run_graph_kernel_into_(const ExprSpec& raw_spec, + std::span inputs, + std::size_t rows, std::size_t cols, + bool vector_out, Tensor& out) { const bool is_fold = raw_spec.fold.has_value(); const ExprSpec spec = is_fold ? raw_spec : nn::canonicalize_expr_spec(raw_spec); @@ -2317,12 +2189,15 @@ class GpuEngine final : public ComputeEngine const nn::ExprPrecSig psig = nn::expr_prec_sig_of(inputs, out.precision()); #ifndef NN_FUSED_REGISTRY_EMBEDDED // 无注册表模式（NN_EXPR_SCAN 收集器 TU）：融合分支整体不编译 - (void)rows; (void)cols; (void)vector_out; (void)psig; + (void)rows; + (void)cols; + (void)vector_out; + (void)psig; #endif #ifdef NN_FUSED_REGISTRY_EMBEDDED - const nn::fused::FusedShader* fs = (psig != 0) - ? find_prec_variant_(key, psig, spec, /*log_hit=*/true) - : nn::fused::find_fused(key); + const nn::fused::FusedShader* fs = + (psig != 0) ? find_prec_variant_(key, psig, spec, /*log_hit=*/true) + : nn::fused::find_fused(key); if (fs && backend_->has_fused_shader(fs->key)) { auto fi_r = fused_buffers_(inputs); @@ -2349,11 +2224,9 @@ class GpuEngine final : public ComputeEngine } const auto vp = nn::expr_spec_runtime_view_params(spec); auto r = backend_->run_fused_gpu( - fs->key, gpu_inputs, spec.consts, rows, cols, vector_out, vp, - spec.rparams, out_override, - nn::expr_spec_runtime_matmul_k(spec), - nn::expr_spec_runtime_matmul_batch(spec), - nn::expr_spec_runtime_matmul_trans(spec), + fs->key, gpu_inputs, spec.consts, rows, cols, vector_out, vp, spec.rparams, + out_override, nn::expr_spec_runtime_matmul_k(spec), + nn::expr_spec_runtime_matmul_batch(spec), nn::expr_spec_runtime_matmul_trans(spec), is_fold ? nn::expr_spec_runtime_fold_k(spec) : std::nullopt, /*out_f16=*/false, psig); NN_TRY_CHECK(r); @@ -2361,7 +2234,8 @@ class GpuEngine final : public ComputeEngine } #endif NN_FAIL("GpuEngine::end_expr: 融合 kernel 未命中 AOT 融合 shader（闭合世界）；" - "请将 begin_expr/end_expr 段纳入构建期扫描（scan_exprs）；key=" + key); + "请将 begin_expr/end_expr 段纳入构建期扫描（scan_exprs）；key=" + + key); } // ── 复合 kernel 可用性查询（闭合世界降级的判据）──────────────────────── @@ -2370,11 +2244,14 @@ class GpuEngine final : public ComputeEngine // execute_fused_graph）。与 run_graph_kernel_into_ / eval_expr_multi_into // 的选路同源——这里若"过宽"（说可用但实际查不到）不会静默错值，只是让 // 下游把原来的硬报错照常抛出。 - [[nodiscard]] bool composite_kernel_available_( - const FusedKernel& k, std::span ins, const Tensor& out) const + [[nodiscard]] bool composite_kernel_available_(const FusedKernel& k, + std::span ins, + const Tensor& out) const { #ifndef NN_FUSED_REGISTRY_EMBEDDED - (void)k; (void)ins; (void)out; + (void)k; + (void)ins; + (void)out; return true; // 收集器 TU：融合执行分支整体不编译 #else if (k.spec.extras.empty()) @@ -2384,20 +2261,17 @@ class GpuEngine final : public ComputeEngine const ExprSpec spec = is_fold ? k.spec : nn::canonicalize_expr_spec(k.spec); const std::string key = nn::expr_spec_key(spec); const nn::ExprPrecSig psig = nn::expr_prec_sig_of(ins, out.precision()); - const nn::fused::FusedShader* fs = (psig != 0) - ? find_prec_variant_(key, psig, spec, /*log_hit=*/false) - : nn::fused::find_fused(key); + const nn::fused::FusedShader* fs = + (psig != 0) ? find_prec_variant_(key, psig, spec, /*log_hit=*/false) + : nn::fused::find_fused(key); return fs != nullptr && backend_->has_fused_shader(fs->key); } // 写穿物化（多输出）：V0 专用——eval_expr_multi_into 的 NVI 前置校验 // 要求输入/输出全 f32，这里同条件（不满足 = 不可用 = 逐成员降级） for (const auto& t : ins) - if (t.precision() != Precision::F32) - return false; - if (out.precision() != Precision::F32) - return false; - const std::string key = - nn::expr_spec_key(nn::canonicalize_expr_spec(k.spec)); + if (t.precision() != Precision::F32) return false; + if (out.precision() != Precision::F32) return false; + const std::string key = nn::expr_spec_key(nn::canonicalize_expr_spec(k.spec)); const nn::fused::FusedShader* fs = nn::fused::find_fused(key); return fs != nullptr && backend_->has_fused_shader(fs->key); #endif @@ -2409,19 +2283,16 @@ class GpuEngine final : public ComputeEngine // 无依赖边形态，如 AdamW 的 decay 与 update 都写 p）。 [[nodiscard]] static const void* tensor_storage_token_(const Tensor& t) { - if (!t.valid()) - return nullptr; + if (!t.valid()) return nullptr; if (t.is_gpu()) { if (t.precision() == Precision::F16) { auto s = t.gpu_shared(); - return s ? static_cast(s->shared_buffer().get()) - : nullptr; + return s ? static_cast(s->shared_buffer().get()) : nullptr; } auto s = t.gpu_shared(); - return s ? static_cast(s->shared_buffer().get()) - : nullptr; + return s ? static_cast(s->shared_buffer().get()) : nullptr; } if (t.precision() == Precision::F16) { @@ -2440,20 +2311,19 @@ class GpuEngine final : public ComputeEngine // **返回命中的批量变体"前缀"**（`key` / `key#xxxx` / `key#s`，不含 `#b`） // ——组键与派发键都必须用它，否则"精确签名未发射、退回 `#s`"的一类会被 // 派发到不存在的键上。未命中返回空串（→ 调用方逐 kernel 降级）。 - [[nodiscard]] std::string batch_variant_prefix_(const ExprSpec& spec, - nn::ExprPrecSig sig) const + [[nodiscard]] std::string batch_variant_prefix_(const ExprSpec& spec, nn::ExprPrecSig sig) const { #ifdef NN_FUSED_REGISTRY_EMBEDDED const std::string key = nn::expr_spec_key(spec); - const auto hit = [this](const std::string& k) { + const auto hit = [this](const std::string& k) + { const nn::fused::FusedShader* fs = nn::fused::find_fused(k); return fs != nullptr && backend_->has_fused_shader(fs->key); }; // 与 find_prec_variant_ 同序：精确签名键 → 全 f16 兜底（`key#s`， // 全 f16 签名不重复发射 `key#xxxx` 变体）。 const std::string exact = nn::expr_prec_sig_key(key, sig); - if (hit(exact + nn::EXPR_BATCH_SUFFIX)) - return exact; + if (hit(exact + nn::EXPR_BATCH_SUFFIX)) return exact; if (nn::expr_prec_sig_all_f16(spec, sig) && hit(key + nn::EXPR_PREC_STORAGE_SUFFIX + nn::EXPR_BATCH_SUFFIX)) return key + nn::EXPR_PREC_STORAGE_SUFFIX; @@ -2524,15 +2394,13 @@ class GpuEngine final : public ComputeEngine } else { - mins.push_back(k < nd.input_tensors.size() ? nd.input_tensors[k] - : Tensor{}); + mins.push_back(k < nd.input_tensors.size() ? nd.input_tensors[k] : Tensor{}); } } const auto mo = g.node_outputs.find(m); if (mo == g.node_outputs.end()) NN_FAIL("GpuEngine::end_expr: 降级派发占位缺失（图 IR 状态损坏）"); - return run_graph_kernel_into_(nd.spec, mins, nd.rows, nd.cols, nd.vector_out, - mo->second); + return run_graph_kernel_into_(nd.spec, mins, nd.rows, nd.cols, nd.vector_out, mo->second); } // ── 辅助：f16 Tensor → GpuTensor **绑定视图** ───────────────────────── @@ -2560,12 +2428,10 @@ class GpuEngine final : public ComputeEngine { if (src.is_gpu()) { - if (P == src.precision()) - return src; // 同设备同精度 = 零拷贝别名（共享拷贝） - return cast(src, P); // 同设备异精度 = 引擎内 cast + if (P == src.precision()) return src; // 同设备同精度 = 零拷贝别名（共享拷贝） + return cast(src, P); // 同设备异精度 = 引擎内 cast } - if (!src.valid()) - NN_FAIL("import: invalid tensor"); + if (!src.valid()) NN_FAIL("import: invalid tensor"); if (src.precision() == Precision::F32 && P == Precision::F32) { auto r = GpuTensor::from_matrix(src.cpu_matrix(), *backend_); diff --git a/include/neuralnet.cpp/compute_layer_attention.hpp b/include/neuralnet.cpp/compute_layer_attention.hpp index 81bb36c1..3ad48d8c 100644 --- a/include/neuralnet.cpp/compute_layer_attention.hpp +++ b/include/neuralnet.cpp/compute_layer_attention.hpp @@ -70,35 +70,34 @@ enum class AttnMaskKind // fold 构造：掩码 × 位置偏置两个正交维度；生成的 IR 指令顺序固定为 // 掩码（causal → doc）→ ALiBi 斜率项，与拆分前的 5 个融合变体逐指令一致。 [[nodiscard]] inline ExprSpec make_fold_attn_o(std::uint32_t seq, std::uint32_t dk, - std::uint32_t bh, AttnMaskKind mask, - bool score_bias, + std::uint32_t bh, AttnMaskKind mask, bool score_bias, std::uint32_t kv_ratio = 1) { ExprSpec s; FoldSpec f; - f.k = seq; - f.num_state = 2; // 0 = m, 1 = l - f.inits = { -std::numeric_limits::infinity(), Scalar{0} }; + f.k = seq; + f.num_state = 2; // 0 = m, 1 = l + f.inits = {-std::numeric_limits::infinity(), Scalar{0}}; f.vec_state_len = dk; - f.matmul = MatmulSpec{ 0, 1, 1, 0, dk, bh }; // A=Q(trans), B=K, k=dk, batch=bh + f.matmul = MatmulSpec{0, 1, 1, 0, dk, bh}; // A=Q(trans), B=K, k=dk, batch=bh // 常量：c0 = 0（拷贝/中性元）、c1 = -inf（掩码屏蔽） - s.consts = { Scalar{0}, -std::numeric_limits::infinity() }; + s.consts = {Scalar{0}, -std::numeric_limits::infinity()}; // 输入/视图槽位（与掩码策略 append_fold_inputs 的追加顺序约定一致）： // 0=Q 1=K 2=V_t；[slopes (1,bh)→BatchMod]；[doc_col (BH*seq,1)→RowBroadcast]； // [doc_ids (1,bh*seq)→BatchCol(seq)] const bool has_slope = score_bias; - const bool has_doc = (mask == AttnMaskKind::CausalDoc); + const bool has_doc = (mask == AttnMaskKind::CausalDoc); const bool has_causal = (mask != AttnMaskKind::Plain); f.tri_skip = has_causal; // 生成器据此钳 valid（整块/边界跳过 -inf 区） - s.views = { linear(), linear(), linear() }; + s.views = {linear(), linear(), linear()}; // GQA（docs 22 §3.2 P2）：K/V_t 只存 bh/kv_ratio 份行块（块映射） // —— 挂成 HeadGroup 视图，比率走运行期 vp 槽（不进 key → 任意 n_rep // 共用一份 shader）。kv_ratio==1 时保持 linear()，既有结构逐字节不变。 if (kv_ratio > 1) { - s.views[1].kind = static_cast(ExprViewKind::HeadGroup); + s.views[1].kind = static_cast(ExprViewKind::HeadGroup); s.views[1].param = kv_ratio; - s.views[2].kind = static_cast(ExprViewKind::HeadGroup); + s.views[2].kind = static_cast(ExprViewKind::HeadGroup); s.views[2].param = kv_ratio; } std::uint8_t slot_slope = 0, slot_dc = 0, slot_ids = 0; @@ -115,13 +114,13 @@ enum class AttnMaskKind s.views.push_back(batch_col(seq)); // doc_ids：b[batch*seq + col] } - const auto ins = [](ExprOp op, std::uint8_t dst, ExprOperand a, - ExprOperand b = {}) { + const auto ins = [](ExprOp op, std::uint8_t dst, ExprOperand a, ExprOperand b = {}) + { ExprInstr i; - i.op = static_cast(op); + i.op = static_cast(op); i.dst = dst; - i.a = a; - i.b = b; + i.a = a; + i.b = b; return i; }; std::uint8_t nreg = 2; // 状态占 0..1，临时从 2 递增 @@ -138,11 +137,11 @@ enum class AttnMaskKind const uint8_t sel = nreg++; // Select(gt, c1=-inf, c0=0) { ExprInstr isel; // 三操作数指令：手填 c 字段 - isel.op = static_cast(ExprOp::Select); + isel.op = static_cast(ExprOp::Select); isel.dst = sel; - isel.a = reg(gt); - isel.b = cst(1); - isel.c = cst(0); + isel.a = reg(gt); + isel.b = cst(1); + isel.c = cst(0); body.push_back(isel); } const uint8_t sm = nreg++; // S + 屏蔽项 @@ -158,11 +157,11 @@ enum class AttnMaskKind const uint8_t ds = nreg++; { ExprInstr isel; - isel.op = static_cast(ExprOp::Select); + isel.op = static_cast(ExprOp::Select); isel.dst = ds; - isel.a = reg(ne); - isel.b = cst(1); - isel.c = cst(0); + isel.a = reg(ne); + isel.b = cst(1); + isel.c = cst(0); body.push_back(isel); } const uint8_t sm = nreg++; @@ -201,17 +200,17 @@ enum class AttnMaskKind body.push_back(ins(ExprOp::Add, 1, reg(la), reduce(blk_l))); // l ← 状态更新 // ── vecacc：O *= α；O += Σ_块 e·V_t ── - f.vecacc = VecAccSpec{ 0, weight, 2, alpha, 1 }; - f.body = std::move(body); + f.vecacc = VecAccSpec{0, weight, 2, alpha, 1}; + f.body = std::move(body); // ── 向量域 finalize：out[d] = O[d] / l ── // dst 必须用**临时寄存器**（新号）：向量域逐列循环执行，若 dst 复用被读 // 状态（如 l），第 1 列的输出会覆盖除数 → 第 2 列起读到被污染的值 // （实测 0.4 = 10/25：l=4 被首列输出 25 覆盖）。 { const uint8_t fout = nreg++; - f.finalize = { ins(ExprOp::Div, fout, nn::expr::vec_state(0), reg(1)) }; + f.finalize = {ins(ExprOp::Div, fout, nn::expr::vec_state(0), reg(1))}; } - s.fold = std::move(f); + s.fold = std::move(f); s.num_regs = nreg; return s; } @@ -246,13 +245,12 @@ inline constexpr AttnMaskKind kMaskKinds[] = { constexpr std::uint32_t seq = 64, dk = 2, bh = 2; for (const AttnMaskKind mask : kMaskKinds) { - for (const bool score_bias : { false, true }) + for (const bool score_bias : {false, true}) { if (mask == AttnMaskKind::Plain && score_bias) continue; // 无掩码 + 位置偏置：本仓不构造该组合 - for (const std::uint32_t kv_ratio : { 1u, 2u }) - (void)dsl::FoldAnchor{ - make_fold_attn_o(seq, dk, bh, mask, score_bias, kv_ratio) }; + for (const std::uint32_t kv_ratio : {1u, 2u}) + (void)dsl::FoldAnchor{make_fold_attn_o(seq, dk, bh, mask, score_bias, kv_ratio)}; } } return true; @@ -305,8 +303,8 @@ class AttnScoreMask [[nodiscard]] virtual expr::AttnMaskKind mask_kind() const noexcept = 0; // 每步准备掩码输入张量（无掩码分量 = no-op） - [[nodiscard]] virtual Result prepare( - ComputeEngine& /*engine*/, std::size_t /*batch*/, std::size_t /*seq*/) + [[nodiscard]] virtual Result prepare(ComputeEngine& /*engine*/, std::size_t /*batch*/, + std::size_t /*seq*/) { return {}; } // fold 输入追加：顺序必须与 make_fold_attn_o 的 views 顺序一致 @@ -315,10 +313,10 @@ class AttnScoreMask // 反向：S = mask(Q·Kᵀ)（DSL 文本 —— 与 forward fold 掩码语义同构； // 位置偏置由位置编码策略另行叠加，见 PositionEncoder::apply_score_bias） - [[nodiscard]] virtual Result masked_scores( - ComputeEngine& engine, const Tensor& Q, const Tensor& K, - std::size_t BH, std::size_t seq, const PrecisionProfile& p, - std::uint32_t kv_ratio) const = 0; + [[nodiscard]] virtual Result masked_scores(ComputeEngine& engine, const Tensor& Q, + const Tensor& K, std::size_t BH, + std::size_t seq, const PrecisionProfile& p, + std::uint32_t kv_ratio) const = 0; // 文档 id 注入（仅文档感知策略覆写；其余 no-op = 无文档掩码） virtual void set_doc_ids(std::span /*ids*/) {} @@ -331,10 +329,10 @@ class PlainScoreMask final : public AttnScoreMask [[nodiscard]] expr::AttnMaskKind mask_kind() const noexcept override { return expr::AttnMaskKind::Plain; } - [[nodiscard]] Result masked_scores( - ComputeEngine& engine, const Tensor& Q, const Tensor& K, - std::size_t BH, std::size_t /*seq*/, const PrecisionProfile& p, - std::uint32_t kv_ratio) const override + [[nodiscard]] Result masked_scores(ComputeEngine& engine, const Tensor& Q, + const Tensor& K, std::size_t BH, std::size_t /*seq*/, + const PrecisionProfile& p, + std::uint32_t kv_ratio) const override { return dsl::compute(engine, dsl::matmul_kv(Q, K, true, false, BH, 1, kv_ratio), p.compute); // 输出 (BH·seq, seq) 由 matmul 段推导 @@ -351,17 +349,15 @@ class CausalScoreMask : public AttnScoreMask [[nodiscard]] expr::AttnMaskKind mask_kind() const noexcept override { return expr::AttnMaskKind::Causal; } - [[nodiscard]] Result masked_scores( - ComputeEngine& engine, const Tensor& Q, const Tensor& K, - std::size_t BH, std::size_t seq, const PrecisionProfile& p, - std::uint32_t kv_ratio) const override + [[nodiscard]] Result masked_scores(ComputeEngine& engine, const Tensor& Q, + const Tensor& K, std::size_t BH, std::size_t seq, + const PrecisionProfile& p, + std::uint32_t kv_ratio) const override { const auto scores = dsl::matmul_kv(Q, K, true, false, BH, 1, kv_ratio); - const auto causal = dsl::select(dsl::col() > dsl::row(), - Scalar{1}, Scalar{0}); - return dsl::compute(engine, - scores + dsl::select(causal != Scalar{0}, kNegInf, Scalar{0}), - BH * seq, seq, p.compute); + const auto causal = dsl::select(dsl::col() > dsl::row(), Scalar{1}, Scalar{0}); + return dsl::compute(engine, scores + dsl::select(causal != Scalar{0}, kNegInf, Scalar{0}), + BH * seq, seq, p.compute); } }; @@ -380,7 +376,12 @@ class CausalDocScoreMask final : public CausalScoreMask void set_doc_ids(std::span ids) override { - if (ids.empty()) { doc_ids_.clear(); has_doc_ids_ = false; return; } + if (ids.empty()) + { + doc_ids_.clear(); + has_doc_ids_ = false; + return; + } doc_ids_.assign(ids.begin(), ids.end()); has_doc_ids_ = true; } @@ -390,8 +391,8 @@ class CausalDocScoreMask final : public CausalScoreMask [[nodiscard]] expr::AttnMaskKind mask_kind() const noexcept override { return expr::AttnMaskKind::CausalDoc; } - [[nodiscard]] Result prepare( - ComputeEngine& engine, std::size_t batch, std::size_t seq) override + [[nodiscard]] Result prepare(ComputeEngine& engine, std::size_t batch, + std::size_t seq) override { // 文档感知：doc_ids_ 每 step 变化，每步重建（小张量 O(BH*seq)）。 // doc_ids_cache_ 为 (1, BH*seq) 布局——每 (b,h) 块重复 doc_ids[b*seq..] @@ -403,8 +404,7 @@ class CausalDocScoreMask final : public CausalScoreMask for (std::size_t b = 0; b < batch; ++b) for (std::size_t h = 0; h < num_heads_; ++h) for (std::size_t i = 0; i < seq; ++i) - d[(b * num_heads_ + h) * seq + i] = - static_cast(doc_ids_[b * seq + i]); + d[(b * num_heads_ + h) * seq + i] = static_cast(doc_ids_[b * seq + i]); auto t = detail::upload_span(engine, 1, BH * seq, Precision::F32, std::span(d)); NN_TRY_CHECK(t); doc_ids_cache_ = std::move(*t); @@ -419,8 +419,7 @@ class CausalDocScoreMask final : public CausalScoreMask else { // CPU：直接构造目标形状，避免为一个小表引入额外转置遍历。 - auto tc = detail::upload_span(engine, BH * seq, 1, Precision::F32, - std::span(d)); + auto tc = detail::upload_span(engine, BH * seq, 1, Precision::F32, std::span(d)); NN_TRY_CHECK(tc); doc_col_ = std::move(*tc); } @@ -433,20 +432,18 @@ class CausalDocScoreMask final : public CausalScoreMask io.push_back(doc_ids_cache_); } - [[nodiscard]] Result masked_scores( - ComputeEngine& engine, const Tensor& Q, const Tensor& K, - std::size_t BH, std::size_t seq, const PrecisionProfile& p, - std::uint32_t kv_ratio) const override + [[nodiscard]] Result masked_scores(ComputeEngine& engine, const Tensor& Q, + const Tensor& K, std::size_t BH, std::size_t seq, + const PrecisionProfile& p, + std::uint32_t kv_ratio) const override { const auto scores = dsl::matmul_kv(Q, K, true, false, BH, 1, kv_ratio); - const auto causal = dsl::select(dsl::col() > dsl::row(), - Scalar{1}, Scalar{0}); - const auto blocked = causal + dsl::select( - dsl::row_broadcast(doc_col_) != dsl::batch_col(doc_ids_cache_, seq), - Scalar{1}, Scalar{0}); - return dsl::compute(engine, - scores + dsl::select(blocked != Scalar{0}, kNegInf, Scalar{0}), - BH * seq, seq, p.compute); + const auto causal = dsl::select(dsl::col() > dsl::row(), Scalar{1}, Scalar{0}); + const auto blocked = causal + dsl::select(dsl::row_broadcast(doc_col_) != + dsl::batch_col(doc_ids_cache_, seq), + Scalar{1}, Scalar{0}); + return dsl::compute(engine, scores + dsl::select(blocked != Scalar{0}, kNegInf, Scalar{0}), + BH * seq, seq, p.compute); } }; @@ -571,37 +568,31 @@ class AttentionBase : public Layer } public: - AttentionBase(std::size_t d_model, std::size_t num_heads, - std::size_t seq_len = 0, - bool subln = false, - NormType subln_norm_type = NormType::LayerNorm, - std::size_t n_head_kv = 0, - WeightQuant weight_quant = WeightQuant::None) + AttentionBase(std::size_t d_model, std::size_t num_heads, std::size_t seq_len = 0, + bool subln = false, NormType subln_norm_type = NormType::LayerNorm, + std::size_t n_head_kv = 0, WeightQuant weight_quant = WeightQuant::None) : d_model_(d_model), num_heads_(num_heads), n_head_kv_(n_head_kv == 0 ? num_heads : n_head_kv), kv_ratio_(num_heads / (n_head_kv == 0 ? num_heads : n_head_kv)), - d_k_(d_model / num_heads), - seq_len_(seq_len), + d_k_(d_model / num_heads), seq_len_(seq_len), scale_(Scalar{1} / std::sqrt(static_cast(d_model / num_heads))), w_q_(make_linear_layer(d_model, d_model, weight_quant)), // Linear(in_features, out_features)：GQA 只收窄**输出**宽度 // （K/V 的 out = n_head_kv·d_k），输入侧仍是 d_model。 - w_k_(make_linear_layer( - d_model, (n_head_kv == 0 ? num_heads : n_head_kv) * (d_model / num_heads), - weight_quant)), - w_v_(make_linear_layer( - d_model, (n_head_kv == 0 ? num_heads : n_head_kv) * (d_model / num_heads), - weight_quant)), - w_o_(make_linear_layer(d_model, d_model, weight_quant)), - weight_quant_(weight_quant) + w_k_(make_linear_layer(d_model, + (n_head_kv == 0 ? num_heads : n_head_kv) * (d_model / num_heads), + weight_quant)), + w_v_(make_linear_layer(d_model, + (n_head_kv == 0 ? num_heads : n_head_kv) * (d_model / num_heads), + weight_quant)), + w_o_(make_linear_layer(d_model, d_model, weight_quant)), weight_quant_(weight_quant) { NN_ASSERT_ALWAYS(d_model % num_heads == 0, "AttentionBase: d_model must be divisible by num_heads"); NN_ASSERT_ALWAYS(n_head_kv_ != 0 && num_heads % n_head_kv_ == 0, "AttentionBase: num_heads must be divisible by n_head_kv"); mask_ = std::make_unique(); - if (subln) - attn_sub_norm_ = make_norm_layer(d_model, subln_norm_type); + if (subln) attn_sub_norm_ = make_norm_layer(d_model, subln_norm_type); } [[nodiscard]] PositionEncoder& position_encoder() noexcept { return *pos_; } @@ -618,8 +609,8 @@ class AttentionBase : public Layer // 的那一份。非三值模式 split 退化为恒等 → 逐位不变。 void set_precision_profile(const PrecisionProfile& profile) override { - const TernaryProfileSplit sp = split_ternary_profile( - layer_name(), profile, weight_quant_ == WeightQuant::T1_58); + const TernaryProfileSplit sp = + split_ternary_profile(layer_name(), profile, weight_quant_ == WeightQuant::T1_58); Layer::set_precision_profile(sp.self); w_q_->set_precision_profile(sp.linear); w_k_->set_precision_profile(sp.linear); @@ -648,12 +639,13 @@ class AttentionBase : public Layer for (std::size_t r = 0; r < d_k_; ++r) idx[h * d_k_ + r] = static_cast(src + r); } - NN_TRY(gi, detail::upload_span(engine, num_heads_ * d_k_, 1, - Precision::F32, idx)); + NN_TRY(gi, detail::upload_span(engine, num_heads_ * d_k_, 1, Precision::F32, idx)); gqa_expand_idx_ = std::move(*gi); } // M6 段 C：softmax_ 是子 Layer（engine 由 init 绑定），否则 forward 内 fail-fast - { NN_TRY(r6, softmax_.init(engine)); } + { + NN_TRY(r6, softmax_.init(engine)); + } // 注意力侧位置编码（RoPE/ALiBi 的 init 是 no-op；恒等策略同样 no-op）—— // 按 M6 段 C 的不变量"复合层 init_impl 必须 init 全部子对象"统一调用。 if (attn_sub_norm_) { NN_TRY(r7, attn_sub_norm_->init(engine)); } @@ -662,8 +654,8 @@ class AttentionBase : public Layer std::vector parameters() override { - auto r = collect_refs(w_q_->parameters(), w_k_->parameters(), - w_v_->parameters(), w_o_->parameters()); + auto r = collect_refs(w_q_->parameters(), w_k_->parameters(), w_v_->parameters(), + w_o_->parameters()); if (attn_sub_norm_) { auto s = attn_sub_norm_->parameters(); @@ -705,10 +697,7 @@ class AttentionBase : public Layer } // 文档感知：转发给掩码策略（非文档策略 = no-op，与旧"无 doc 钩子"一致） - void set_doc_ids(std::span ids) override - { - score_mask().set_doc_ids(ids); - } + void set_doc_ids(std::span ids) override { score_mask().set_doc_ids(ids); } void clear_cache() override { @@ -731,11 +720,16 @@ class AttentionBase : public Layer if (Q_cache_.valid()) r.emplace_back(Q_cache_); if (K_cache_.valid()) r.emplace_back(K_cache_); if (V_cache_.valid()) r.emplace_back(V_cache_); - auto wq = w_q_->activation_cache(); r.insert(r.end(), wq.begin(), wq.end()); - auto wk = w_k_->activation_cache(); r.insert(r.end(), wk.begin(), wk.end()); - auto wv = w_v_->activation_cache(); r.insert(r.end(), wv.begin(), wv.end()); - auto wo = w_o_->activation_cache(); r.insert(r.end(), wo.begin(), wo.end()); - auto sm = softmax_.activation_cache(); r.insert(r.end(), sm.begin(), sm.end()); + auto wq = w_q_->activation_cache(); + r.insert(r.end(), wq.begin(), wq.end()); + auto wk = w_k_->activation_cache(); + r.insert(r.end(), wk.begin(), wk.end()); + auto wv = w_v_->activation_cache(); + r.insert(r.end(), wv.begin(), wv.end()); + auto wo = w_o_->activation_cache(); + r.insert(r.end(), wo.begin(), wo.end()); + auto sm = softmax_.activation_cache(); + r.insert(r.end(), sm.begin(), sm.end()); if (attn_sub_norm_) { auto sn = attn_sub_norm_->activation_cache(); @@ -744,15 +738,13 @@ class AttentionBase : public Layer return r; } - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); - if (input.rows() != d_model_) - NN_FAIL("AttentionBase forward: input shape mismatch"); + if (input.rows() != d_model_) NN_FAIL("AttentionBase forward: input shape mismatch"); const std::size_t total_seq = input.cols(); - const std::size_t seq = (seq_len_ > 0) ? seq_len_ : total_seq; - const std::size_t batch = (seq_len_ > 0) ? (total_seq / seq_len_) : 1; + const std::size_t seq = (seq_len_ > 0) ? seq_len_ : total_seq; + const std::size_t batch = (seq_len_ > 0) ? (total_seq / seq_len_) : 1; if (total_seq != batch * seq) NN_FAIL("AttentionBase forward: cols not divisible by seq_len"); // 1. 线性投影 → Q/K/V: (H*d_k, batch*seq) @@ -810,8 +802,7 @@ class AttentionBase : public Layer // （另分配一整块缓冲），此处不适用。 // scale_ 用 rparam 承载 → 值不进 expr_spec_key（融合二期教训 1）。 { - auto qs = dsl::compute_into(engine, - dsl::leaf(Q) * dsl::rparam(scale_), Q); + auto qs = dsl::compute_into(engine, dsl::leaf(Q) * dsl::rparam(scale_), Q); NN_TRY_CHECK(qs); } @@ -835,9 +826,9 @@ class AttentionBase : public Layer std::vector fold_in{Q, K, *V_t}; pos_->append_score_bias_inputs(fold_in); mask_->append_fold_inputs(fold_in); - const nn::ExprSpec fold_spec = nn::expr::make_fold_attn_o( - seq, d_k_, BH, mask_->mask_kind(), pos_->has_score_bias(), - static_cast(kv_ratio_)); + const nn::ExprSpec fold_spec = + nn::expr::make_fold_attn_o(seq, d_k_, BH, mask_->mask_kind(), pos_->has_score_bias(), + static_cast(kv_ratio_)); NN_TRY(fv, nn::validate_expr_spec(fold_spec, fold_in.size())); // scale 已折进 Q（见上 2.6；运行时值不进 expr_spec_key）——fold 的 mm 段直接消费 NN_TRY(O_t_r, engine.eval_expr(fold_spec, fold_in, BH * seq, d_k_, p_.compute)); @@ -879,14 +870,13 @@ class AttentionBase : public Layer return w_o_->forward(attn_out); } - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { ComputeEngine& engine = engine_ref(); // 推断 batch/seq（与 forward 一致） const std::size_t total_seq = grad_output.cols(); - const std::size_t seq = (seq_len_ > 0) ? seq_len_ : total_seq; - const std::size_t batch = (seq_len_ > 0) ? (total_seq / seq_len_) : 1; + const std::size_t seq = (seq_len_ > 0) ? seq_len_ : total_seq; + const std::size_t batch = (seq_len_ > 0) ? (total_seq / seq_len_) : 1; const std::size_t H_dk = num_heads_ * d_k_; const std::size_t BH = batch * num_heads_; @@ -923,8 +913,8 @@ class AttentionBase : public Layer // 输出形状：transA=true → rows=BH*A.cols()；transB=false → cols=B.cols() // GQA：V_cache_ 存的是 n_head_kv 个头 → B 侧按头分组比取行块 const std::uint32_t kvr = static_cast(kv_ratio_); - auto grad_A = dsl::compute(engine, - dsl::matmul_kv(grad_concat_re, V_cache_, true, false, BH, 1, kvr), + auto grad_A = dsl::compute( + engine, dsl::matmul_kv(grad_concat_re, V_cache_, true, false, BH, 1, kvr), BH * grad_concat_re.cols(), V_cache_.cols(), p_.compute); NN_TRY_CHECK(grad_A); nn_dbg_scan("attn.grad_A", engine, *grad_A); @@ -941,44 +931,38 @@ class AttentionBase : public Layer // R = row_sum(W·P) → (BH*seq, 1) // X = scale·W·(P − R) → (BH*seq, seq)（物化） // grad_Q = K × X^T；grad_K = Q × X；grad_V = W^T × G - auto R = dsl::compute_reduce(engine, - dsl::row_reduce_sum(dsl::leaf(*W_re) * dsl::leaf(*grad_A)), - BH * seq, seq); + auto R = dsl::compute_reduce( + engine, dsl::row_reduce_sum(dsl::leaf(*W_re) * dsl::leaf(*grad_A)), BH * seq, seq); NN_TRY_CHECK(R); nn_dbg_scan("attn.R", engine, *R); auto X = dsl::compute(engine, - dsl::leaf(*W_re) - * (dsl::leaf(*grad_A) - dsl::row_broadcast(*R)), - BH * seq, seq, p_.compute); + dsl::leaf(*W_re) * (dsl::leaf(*grad_A) - dsl::row_broadcast(*R)), + BH * seq, seq, p_.compute); NN_TRY_CHECK(X); nn_dbg_scan("attn.X", engine, *X); // grad_Q = K × X^T（K_b (d_k,seq)，X_b (seq,seq) 按 X^T 使用） // dsl::matmul(batch)（transA=F,transB=T → rows=K.rows(), cols=X.rows()/BH） // GQA：A 侧 K 只存 n_head_kv 个头 → 按头分组比取行块；输出仍是 // **全 H 头**（rows = BH·d_k），即"紧输入 → 满输出"的映射。 - auto gq = dsl::compute(engine, - dsl::matmul_kv(K_cache_, *X, false, true, BH, kvr, 1), - BH * d_k_, X->rows() / BH, p_.compute); + auto gq = dsl::compute(engine, dsl::matmul_kv(K_cache_, *X, false, true, BH, kvr, 1), + BH * d_k_, X->rows() / BH, p_.compute); NN_TRY_CHECK(gq); nn_dbg_scan("attn.gq(pre-scale)", engine, *gq); // grad_Q 补乘 scale（forward 把 scale 折进了 Q）：compute_into 原地 // （同 forward 2.6 的写法；rparam 承载 scale 值不进 key） { - auto gqs = dsl::compute_into(engine, - dsl::leaf(*gq) * dsl::rparam(scale_), *gq); + auto gqs = dsl::compute_into(engine, dsl::leaf(*gq) * dsl::rparam(scale_), *gq); NN_TRY_CHECK(gqs); } // grad_K = Q × X（dsl::matmul(batch)：false,false → rows=Q.rows(), cols=X.cols()） - auto gk = dsl::compute(engine, - dsl::matmul(Q_cache_, *X, false, false, BH), - Q_cache_.rows(), X->cols(), p_.compute); + auto gk = dsl::compute(engine, dsl::matmul(Q_cache_, *X, false, false, BH), + Q_cache_.rows(), X->cols(), p_.compute); NN_TRY_CHECK(gk); nn_dbg_scan("attn.gk", engine, *gk); // grad_V = W^T × G（W_b (seq,seq) 按 W^T 使用，G_b (seq,d_k)）→ (BH*seq, d_k) // dsl::matmul(batch)：transA=T,transB=F → rows=BH*W.cols(), cols=G.cols() - auto gv_t = dsl::compute(engine, - dsl::matmul(*W_re, *G, true, false, BH), - BH * W_re->cols(), G->cols(), p_.compute); + auto gv_t = dsl::compute(engine, dsl::matmul(*W_re, *G, true, false, BH), + BH * W_re->cols(), G->cols(), p_.compute); NN_TRY_CHECK(gv_t); // grad_V 转置回 (BH*d_k, seq)（与 forward 的 V_t→V 逆变换一致） NN_TRY(gv_T, engine.transpose(*gv_t)); @@ -1040,8 +1024,8 @@ class AttentionBase : public Layer // grad_input = grad_Q + grad_K + grad_V：三路累加**原地**融合为单趟 // （dsl::compute_into 目标传递，不额外分配）。结合顺序固定为 // (giq + gik) + giv —— 求和顺序确定，结果可复现。 - auto acc = dsl::compute_into(engine, - dsl::leaf(*giq) + dsl::leaf(*gik) + dsl::leaf(*giv), *giq); + auto acc = + dsl::compute_into(engine, dsl::leaf(*giq) + dsl::leaf(*gik) + dsl::leaf(*giv), *giq); NN_TRY_CHECK(acc); return giq; } @@ -1059,12 +1043,8 @@ class AttentionBase : public Layer // // 输入: x_new (d_model, 1) 单个新 token 的嵌入 // 输出: (d_model, 1) attention 输出 - [[nodiscard]] Result forward_step( - ComputeEngine& engine, - const Tensor& x_new, - Tensor& k_cache, - Tensor& v_cache, - std::size_t cur_len) + [[nodiscard]] Result forward_step(ComputeEngine& engine, const Tensor& x_new, + Tensor& k_cache, Tensor& v_cache, std::size_t cur_len) { if (x_new.rows() != d_model_ || x_new.cols() != 1) NN_FAIL("AttentionBase forward_step: x_new must be (d_model, 1)"); @@ -1122,8 +1102,8 @@ class AttentionBase : public Layer // 每头: (1, d_k) × (d_k, new_len) = (1, new_len) // 堆叠: (H, new_len) // scale (1/sqrt(d_k)) 经 rparam 尾链乘在 matmul 结果上（不进 expr_spec_key） - auto scores = dsl::compute(engine, - dsl::matmul(*q_res, *K_T, true, false, num_heads_) * dsl::rparam(scale_), + auto scores = dsl::compute( + engine, dsl::matmul(*q_res, *K_T, true, false, num_heads_) * dsl::rparam(scale_), num_heads_ * q_res->cols(), K_T->cols(), p_.compute); NN_TRY_CHECK(scores); @@ -1140,9 +1120,8 @@ class AttentionBase : public Layer // 每头: (d_k, new_len) × (new_len, 1) = (d_k, 1) // 堆叠: (H*d_k, 1) // dsl::matmul(batch)：纯 {0,1} 结构（attention backward 已登记同 key） - auto attn_out = dsl::compute(engine, - dsl::matmul(*V_T, *attn, false, true, num_heads_), - V_T->rows(), attn->rows() / num_heads_, p_.compute); + auto attn_out = dsl::compute(engine, dsl::matmul(*V_T, *attn, false, true, num_heads_), + V_T->rows(), attn->rows() / num_heads_, p_.compute); NN_TRY_CHECK(attn_out); // 9. SubLN（子层内部归一化；关 = 直通）→ 输出投影 → (d_model, 1) @@ -1164,10 +1143,9 @@ class AttentionBase : public Layer // 3) W = softmax(S)（归约快路径） // 返回 (BH*seq, seq) 的 softmax 归一化权重；attention FLOPs ×1.5–2 // （QK^T 在 backward 重算），训练可接受。 - [[nodiscard]] Result recompute_W_( - ComputeEngine& engine, - const Tensor& Q, const Tensor& K, - std::size_t BH, std::size_t seq) const + [[nodiscard]] Result recompute_W_(ComputeEngine& engine, const Tensor& Q, + const Tensor& K, std::size_t BH, + std::size_t seq) const { NN_TRY(sres, mask_->masked_scores(engine, Q, K, BH, seq, p_, kv_ratio_)); Tensor S = std::move(*sres); @@ -1176,10 +1154,10 @@ class AttentionBase : public Layer // W = softmax(S)：单表达式归约快路径——与 Softmax::forward // 同构：ReduceRef 直接参与算术按行广播；m/l 在 kernel 内部归一化， // 不作为输入/缓存存在。 - return dsl::compute(engine, - dsl::exp(dsl::leaf(S) - dsl::row_reduce_max(S)) - / dsl::row_reduce_sum( - dsl::exp(dsl::leaf(S) - dsl::row_reduce_max(S))), + return dsl::compute( + engine, + dsl::exp(dsl::leaf(S) - dsl::row_reduce_max(S)) / + dsl::row_reduce_sum(dsl::exp(dsl::leaf(S) - dsl::row_reduce_max(S))), BH * seq, seq, p_.compute); } }; @@ -1195,11 +1173,11 @@ class MultiHeadAttention final : public AttentionBase public: [[nodiscard]] const char* layer_name() const noexcept override { return "MultiHeadAttention"; } - MultiHeadAttention(std::size_t d_model, std::size_t num_heads, - std::size_t seq_len = 0, + MultiHeadAttention(std::size_t d_model, std::size_t num_heads, std::size_t seq_len = 0, WeightQuant weight_quant = WeightQuant::None) - : AttentionBase(d_model, num_heads, seq_len, false, NormType::LayerNorm, 0, - weight_quant) {} + : AttentionBase(d_model, num_heads, seq_len, false, NormType::LayerNorm, 0, weight_quant) + { + } }; // ══════════════════════════════════════════════════════════════════════════ @@ -1237,27 +1215,21 @@ class CausalSelfAttention final : public AttentionBase { // 配置期工厂：这是**唯一**决定掩码种类的地方（只依赖文档掩码的开关， // 与位置编码无关）。forward/backward 里没有任何变体判断。 - if (doc_enabled_) - return std::make_unique(num_heads_); + if (doc_enabled_) return std::make_unique(num_heads_); return std::make_unique(); } public: - 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, - bool subln = false, - NormType subln_norm_type = NormType::LayerNorm, - std::size_t n_head_kv = 0, - WeightQuant weight_quant = WeightQuant::None) + 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, + bool subln = false, NormType subln_norm_type = NormType::LayerNorm, + std::size_t n_head_kv = 0, WeightQuant weight_quant = WeightQuant::None) : AttentionBase(d_model, num_heads, seq_len, subln, subln_norm_type, n_head_kv, weight_quant) { // ① 位置编码：注意力侧**由本层自持**（RoPE 的 cos/sin 表随层构建； // Learned/Sinusoidal 在注意力侧是恒等 → 交回模型侧施加）。 - install_position_encoder( - make_attention_position_encoder(pos_enc, d_k_, num_heads_)); + install_position_encoder(make_attention_position_encoder(pos_enc, d_k_, num_heads_)); // ② 掩码：本类的因果语义（与位置编码正交，故独立安装） install_score_mask_(); } @@ -1321,19 +1293,18 @@ class CrossAttention final : public AttentionBase std::size_t seq_kv_len_; // memory 每样本序列长（0 = 单样本，cols 即 seq） public: - [[nodiscard]] const char* layer_name() const noexcept override - { return "CrossAttention"; } + [[nodiscard]] const char* layer_name() const noexcept override { return "CrossAttention"; } // seq_len = query 每样本序列长（0 = 单样本，cols 即 seq） // seq_kv_len = memory 每样本序列长（0 = 单样本；一般 ≠ seq_len） - CrossAttention(std::size_t d_model, std::size_t num_heads, - std::size_t seq_len = 0, std::size_t seq_kv_len = 0, - std::size_t n_head_kv = 0, + CrossAttention(std::size_t d_model, std::size_t num_heads, std::size_t seq_len = 0, + std::size_t seq_kv_len = 0, std::size_t n_head_kv = 0, WeightQuant weight_quant = WeightQuant::None) - : AttentionBase(d_model, num_heads, seq_len, false, NormType::LayerNorm, - n_head_kv, weight_quant), + : AttentionBase(d_model, num_heads, seq_len, false, NormType::LayerNorm, n_head_kv, + weight_quant), seq_kv_len_(seq_kv_len) - {} + { + } // 注意：掩码基类默认 PlainScoreMask、位置编码默认恒等 —— 正是本层语义 // （见类注释），无需 install。 @@ -1353,22 +1324,21 @@ class CrossAttention final : public AttentionBase // query (d_model, batch*seq_q) —— 输出侧序列（decoder 位置） // memory (d_model, batch*seq_kv) —— 被注意序列（encoder 输出等） // 输出 (d_model, batch*seq_q)（列布局与自注意力一致）。 - [[nodiscard]] Result forward_pair(const Tensor& query, - const Tensor& memory) + [[nodiscard]] Result forward_pair(const Tensor& query, const Tensor& memory) { ComputeEngine& engine = engine_ref(); if (query.rows() != d_model_ || memory.rows() != d_model_) NN_FAIL("CrossAttention forward_pair: 输入形状不匹配（rows != d_model）"); // batch 一致性：两侧按各自 seq 长推 batch，必须相等 - const std::size_t seq_q = (seq_len_ > 0) ? seq_len_ : query.cols(); + const std::size_t seq_q = (seq_len_ > 0) ? seq_len_ : query.cols(); const std::size_t seq_kv = (seq_kv_len_ > 0) ? seq_kv_len_ : memory.cols(); if (query.cols() % seq_q != 0 || memory.cols() % seq_kv != 0) NN_FAIL("CrossAttention forward_pair: cols 不能被各自的 seq 长整除"); - const std::size_t batch_q = query.cols() / seq_q; + const std::size_t batch_q = query.cols() / seq_q; const std::size_t batch_kv = memory.cols() / seq_kv; if (batch_q != batch_kv) - NN_FAIL("CrossAttention forward_pair: query 与 memory 的 batch 不一致（" - + std::to_string(batch_q) + " vs " + std::to_string(batch_kv) + "）"); + NN_FAIL("CrossAttention forward_pair: query 与 memory 的 batch 不一致（" + + std::to_string(batch_q) + " vs " + std::to_string(batch_kv) + "）"); // 1. 投影：Q 吃 query，K/V 吃 memory NN_TRY(q_res, w_q_->forward(query)); @@ -1376,7 +1346,7 @@ class CrossAttention final : public AttentionBase NN_TRY(v_proj, w_v_->forward(memory)); // 2. rearrange：Q → (batch*H*d_k, seq_q)；K/V → (batch*KV_dk, seq_kv) - const std::size_t H_dk = num_heads_ * d_k_; + const std::size_t H_dk = num_heads_ * d_k_; const std::size_t KV_dk = n_head_kv_ * d_k_; Tensor Q, K, V; if (batch_q > 1) @@ -1397,8 +1367,7 @@ class CrossAttention final : public AttentionBase // 3. scale 折进 Q（与自注意力 2.6 同口径；rparam 承载不进 key） { - auto qs = dsl::compute_into(engine, - dsl::leaf(Q) * dsl::rparam(scale_), Q); + auto qs = dsl::compute_into(engine, dsl::leaf(Q) * dsl::rparam(scale_), Q); NN_TRY_CHECK(qs); } @@ -1411,12 +1380,11 @@ class CrossAttention final : public AttentionBase std::vector fold_in{Q, K, *V_t}; mask_->append_fold_inputs(fold_in); // Plain/no-bias 构造：f.k = seq_kv；eval 输出行 = BH·seq_q（矩形） - const nn::ExprSpec fold_spec = nn::expr::make_fold_attn_o( - seq_kv, d_k_, BH, mask_->mask_kind(), /*score_bias=*/false, - static_cast(kv_ratio_)); + const nn::ExprSpec fold_spec = + nn::expr::make_fold_attn_o(seq_kv, d_k_, BH, mask_->mask_kind(), /*score_bias=*/false, + static_cast(kv_ratio_)); NN_TRY(fv, nn::validate_expr_spec(fold_spec, fold_in.size())); - NN_TRY(O_t_r, engine.eval_expr(fold_spec, fold_in, BH * seq_q, d_k_, - p_.compute)); + NN_TRY(O_t_r, engine.eval_expr(fold_spec, fold_in, BH * seq_q, d_k_, p_.compute)); // O_t (BH*seq_q, dk) → (BH*dk, seq_q) → (H*dk, batch*seq_q) NN_TRY(O_T_full, engine.transpose(*O_t_r)); NN_TRY(co, engine.rearrange_3d(*O_T_full, d_k_, BH, seq_q, false)); @@ -1442,15 +1410,14 @@ class CrossAttention final : public AttentionBase // ── backward_pair：双路梯度 ──────────────────────────────────────── // grad_output (d_model, batch*seq_q) → {grad_query, grad_memory} - [[nodiscard]] Result backward_pair( - const Tensor& grad_output) + [[nodiscard]] Result backward_pair(const Tensor& grad_output) { ComputeEngine& engine = engine_ref(); if (!Q_cache_.valid() || !K_cache_.valid() || !V_cache_.valid()) NN_FAIL("CrossAttention backward_pair: 无缓存（backward_pair 前必须" "先 forward_pair；checkpoint 模式下缓存不保留）"); // batch/seq 从缓存推（Q: (BH*dk, seq_q)，K: (BH_kv*dk, seq_kv)） - const std::size_t seq_q = Q_cache_.cols(); + const std::size_t seq_q = Q_cache_.cols(); const std::size_t seq_kv = K_cache_.cols(); if (grad_output.cols() % seq_q != 0) NN_FAIL("CrossAttention backward_pair: grad_output.cols() 不能被 seq_q 整除"); @@ -1477,8 +1444,8 @@ class CrossAttention final : public AttentionBase { const std::uint32_t kvr = static_cast(kv_ratio_); // P = grad_A = grad_O^T × V → (BH*seq_q, seq_kv) - auto grad_A = dsl::compute(engine, - dsl::matmul_kv(grad_concat_re, V_cache_, true, false, BH, 1, kvr), + auto grad_A = dsl::compute( + engine, dsl::matmul_kv(grad_concat_re, V_cache_, true, false, BH, 1, kvr), BH * grad_concat_re.cols(), V_cache_.cols(), p_.compute); NN_TRY_CHECK(grad_A); nn_dbg_scan("xattn.grad_A", engine, *grad_A); @@ -1486,40 +1453,34 @@ class CrossAttention final : public AttentionBase NN_TRY(G_T_full, engine.transpose(grad_concat_re)); NN_TRY(G, engine.rearrange_3d(*G_T_full, seq_q, BH, d_k_, false)); // W = softmax(mask(Q·Kᵀ))（矩形；Plain 掩码 = 无屏蔽） - NN_TRY(W_re, recompute_W_pair_(engine, Q_cache_, K_cache_, BH, - seq_q, seq_kv)); + NN_TRY(W_re, recompute_W_pair_(engine, Q_cache_, K_cache_, BH, seq_q, seq_kv)); nn_dbg_scan("xattn.W_re", engine, *W_re); // R = row_sum(W·P) → (BH*seq_q, 1) auto R = dsl::compute_reduce(engine, - dsl::row_reduce_sum(dsl::leaf(*W_re) * dsl::leaf(*grad_A)), - BH * seq_q, seq_kv); + dsl::row_reduce_sum(dsl::leaf(*W_re) * dsl::leaf(*grad_A)), + BH * seq_q, seq_kv); NN_TRY_CHECK(R); // X = W·(P − R) → (BH*seq_q, seq_kv) auto X = dsl::compute(engine, - dsl::leaf(*W_re) - * (dsl::leaf(*grad_A) - dsl::row_broadcast(*R)), - BH * seq_q, seq_kv, p_.compute); + dsl::leaf(*W_re) * (dsl::leaf(*grad_A) - dsl::row_broadcast(*R)), + BH * seq_q, seq_kv, p_.compute); NN_TRY_CHECK(X); nn_dbg_scan("xattn.X", engine, *X); // grad_Q = K × X^T → (BH*d_k, seq_q)（scale 补乘同自反向） - auto gq = dsl::compute(engine, - dsl::matmul_kv(K_cache_, *X, false, true, BH, kvr, 1), - BH * d_k_, X->rows() / BH, p_.compute); + auto gq = dsl::compute(engine, dsl::matmul_kv(K_cache_, *X, false, true, BH, kvr, 1), + BH * d_k_, X->rows() / BH, p_.compute); NN_TRY_CHECK(gq); { - auto gqs = dsl::compute_into(engine, - dsl::leaf(*gq) * dsl::rparam(scale_), *gq); + auto gqs = dsl::compute_into(engine, dsl::leaf(*gq) * dsl::rparam(scale_), *gq); NN_TRY_CHECK(gqs); } // grad_K = Q × X → (BH*d_k, seq_kv) - auto gk = dsl::compute(engine, - dsl::matmul(Q_cache_, *X, false, false, BH), - Q_cache_.rows(), X->cols(), p_.compute); + auto gk = dsl::compute(engine, dsl::matmul(Q_cache_, *X, false, false, BH), + Q_cache_.rows(), X->cols(), p_.compute); NN_TRY_CHECK(gk); // grad_V = W^T × G → (BH*seq_kv, d_k) → 转置回 (BH*d_k, seq_kv) - auto gv_t = dsl::compute(engine, - dsl::matmul(*W_re, *G, true, false, BH), - BH * W_re->cols(), G->cols(), p_.compute); + auto gv_t = dsl::compute(engine, dsl::matmul(*W_re, *G, true, false, BH), + BH * W_re->cols(), G->cols(), p_.compute); NN_TRY_CHECK(gv_t); NN_TRY(gv_T, engine.transpose(*gv_t)); NN_TRY(gv_re, engine.rearrange_3d(*gv_T, d_k_, BH, seq_kv, false)); @@ -1551,10 +1512,8 @@ class CrossAttention final : public AttentionBase auto fold_kv = [&](Tensor&& grad_full) -> Result { if (!gqa_expand_idx_.valid()) return std::move(grad_full); - Tensor collapsed = engine.create_tensor(n_head_kv_ * d_k_, - grad_full.cols(), - grad_full.precision(), - InitSpec::zero()); + Tensor collapsed = engine.create_tensor(n_head_kv_ * d_k_, grad_full.cols(), + grad_full.precision(), InitSpec::zero()); if (!collapsed.valid()) NN_FAIL("CrossAttention: GQA 折叠缓冲分配失败"); NN_TRY(sr, engine.scatter_add_rows(collapsed, gqa_expand_idx_, grad_full)); return collapsed; @@ -1567,9 +1526,8 @@ class CrossAttention final : public AttentionBase CrossAttentionGrads out; out.grad_query = std::move(*giq); // memory 侧两路累加（结合顺序固定 = 可复现） - auto acc = dsl::compute(engine, - dsl::leaf(*gik) + dsl::leaf(*giv), - gik->rows(), gik->cols(), p_.compute); + auto acc = dsl::compute(engine, dsl::leaf(*gik) + dsl::leaf(*giv), gik->rows(), gik->cols(), + p_.compute); NN_TRY_CHECK(acc); out.grad_memory = std::move(*acc); return out; @@ -1578,15 +1536,15 @@ class CrossAttention final : public AttentionBase protected: // W 重算（矩形版；与自反向 recompute_W_ 同构，仅形状参数分离为 // seq_q/seq_kv）。掩码恒 Plain → 屏蔽链为空；位置编码恒等 → 无偏置叠加。 - [[nodiscard]] Result recompute_W_pair_( - ComputeEngine& engine, const Tensor& Q, const Tensor& K, - std::size_t BH, std::size_t seq_q, std::size_t seq_kv) const + [[nodiscard]] Result recompute_W_pair_(ComputeEngine& engine, const Tensor& Q, + const Tensor& K, std::size_t BH, + std::size_t seq_q, std::size_t seq_kv) const { NN_TRY(sres, mask_->masked_scores(engine, Q, K, BH, seq_kv, p_, kv_ratio_)); - return dsl::compute(engine, - dsl::exp(dsl::leaf(*sres) - dsl::row_reduce_max(*sres)) - / dsl::row_reduce_sum( - dsl::exp(dsl::leaf(*sres) - dsl::row_reduce_max(*sres))), + return dsl::compute( + engine, + dsl::exp(dsl::leaf(*sres) - dsl::row_reduce_max(*sres)) / + dsl::row_reduce_sum(dsl::exp(dsl::leaf(*sres) - dsl::row_reduce_max(*sres))), BH * seq_q, seq_kv, p_.compute); } }; diff --git a/include/neuralnet.cpp/compute_layer_base.hpp b/include/neuralnet.cpp/compute_layer_base.hpp index f28be006..ee97e2c6 100644 --- a/include/neuralnet.cpp/compute_layer_base.hpp +++ b/include/neuralnet.cpp/compute_layer_base.hpp @@ -39,8 +39,7 @@ inline void nn_dbg_scan(const char* tag, ComputeEngine& eng, const Tensor& t) std::vector buf(t.rows() * t.cols()); // 宿主桥（17 §3 D11）：不经 Matrix if (auto er = detail::download_span(eng, t, std::span(buf)); !er) { - std::fprintf(stderr, "[bwd][%s] download fail: %s\n", tag, - er.error().message.c_str()); + std::fprintf(stderr, "[bwd][%s] download fail: %s\n", tag, er.error().message.c_str()); return; } double mx = 0.0; @@ -50,8 +49,8 @@ inline void nn_dbg_scan(const char* tag, ComputeEngine& eng, const Tensor& t) if (!std::isfinite(v)) bad = true; else mx = std::max(mx, std::fabs(static_cast(v))); } - std::fprintf(stderr, "[bwd][%s] %s max=%.6g%s\n", tag, t.shape_str().c_str(), - mx, bad ? " <<< NONFINITE" : ""); + std::fprintf(stderr, "[bwd][%s] %s max=%.6g%s\n", tag, t.shape_str().c_str(), mx, + bad ? " <<< NONFINITE" : ""); std::fflush(stderr); } #else @@ -83,17 +82,13 @@ struct PrecisionSet std::uint32_t mask = 0; [[nodiscard]] constexpr bool has(Precision p) const noexcept - { - return (mask & (1u << static_cast(p))) != 0; - } - constexpr void add(Precision p) noexcept - { - mask |= (1u << static_cast(p)); - } + { return (mask & (1u << static_cast(p))) != 0; } + constexpr void add(Precision p) noexcept { mask |= (1u << static_cast(p)); } [[nodiscard]] static constexpr PrecisionSet of(std::initializer_list list) noexcept { PrecisionSet s; - for (Precision p : list) s.add(p); + for (Precision p : list) + s.add(p); return s; } @@ -101,8 +96,8 @@ struct PrecisionSet [[nodiscard]] std::string to_string() const { std::string s; - for (const Precision p : {Precision::F16, Precision::F32, Precision::BF16, - Precision::F64, Precision::T1_58}) + for (const Precision p : + {Precision::F16, Precision::F32, Precision::BF16, Precision::F64, Precision::T1_58}) { if (!has(p)) continue; if (!s.empty()) s += '|'; @@ -115,9 +110,7 @@ struct PrecisionSet // 舍入精度集合 {f16, f32}：**一切既有层的默认能力**（§4.2 —— 这是"不影响其他层" // 的落实方式：默认声明 = 现状，既有层一行不改）。 [[nodiscard]] constexpr PrecisionSet rounding_precision_set() noexcept -{ - return PrecisionSet::of({Precision::F16, Precision::F32}); -} +{ return PrecisionSet::of({Precision::F16, Precision::F32}); } // 按槽的精度能力（四槽与 PrecisionProfile 一一对应；语义见 §4.3.2： // param = "该层 forward 的**有效权重精度**"，不是"参数张量的存储布局"）。 @@ -154,16 +147,16 @@ struct TernaryProfileSplit PrecisionProfile linear; // 直接持有的线性子层（BitLinear 或 Linear） }; -[[nodiscard]] inline TernaryProfileSplit split_ternary_profile( - const char* layer_name, const PrecisionProfile& profile, bool ternary) +[[nodiscard]] inline TernaryProfileSplit +split_ternary_profile(const char* layer_name, const PrecisionProfile& profile, bool ternary) { if (!ternary && profile.param == Precision::T1_58) { - NN_UNWRAP(false, - std::string("层 ") + layer_name + "：收到 param=t1_58，但该层不是" - "三值模式（weight_quant=None）—— 三值权重只对 BitLinear 有意义；" - "请以 weight_quant=t1_58 构造该层，或把 param 改回 f16/f32" - "（docs/internal/21-quantized-weights.md §4.8.3）"); + NN_UNWRAP(false, std::string("层 ") + layer_name + + "：收到 param=t1_58，但该层不是" + "三值模式（weight_quant=None）—— 三值权重只对 BitLinear 有意义；" + "请以 weight_quant=t1_58 构造该层，或把 param 改回 f16/f32" + "（docs/internal/21-quantized-weights.md §4.8.3）"); } TernaryProfileSplit out{profile, profile}; if (out.self.param == Precision::T1_58) out.self.param = Precision::F32; @@ -172,8 +165,7 @@ struct TernaryProfileSplit } // 「param 归一化」单用版（给自建张量的复合层用：`GPTModel` 的 token 嵌入）。 -[[nodiscard]] inline PrecisionProfile without_ternary_param( - PrecisionProfile profile) noexcept +[[nodiscard]] inline PrecisionProfile without_ternary_param(PrecisionProfile profile) noexcept { if (profile.param == Precision::T1_58) profile.param = Precision::F32; return profile; @@ -184,11 +176,9 @@ struct TernaryProfileSplit // param 槽声明了 T1_58，就等价于 weight_quant=T1_58（显式开关优先保留）。 // 这样 CLI `--precision-param t1_58`、`ModelSpec.weight_quant` 与库 API 三条 // 路径收敛到同一个判据，不存在"声明了三值却构造出普通 Linear"的组合。 -[[nodiscard]] inline WeightQuant effective_weight_quant( - const PrecisionProfile& profile, WeightQuant explicit_quant) noexcept -{ - return (profile.param == Precision::T1_58) ? WeightQuant::T1_58 : explicit_quant; -} +[[nodiscard]] inline WeightQuant effective_weight_quant(const PrecisionProfile& profile, + WeightQuant explicit_quant) noexcept +{ return (profile.param == Precision::T1_58) ? WeightQuant::T1_58 : explicit_quant; } // ══════════════════════════════════════════════════════════════════════════ // 参数形态协议（v7 文件角色；docs/internal/23 §2.2） @@ -221,19 +211,22 @@ enum class ModelRole : std::uint8_t { switch (k) { - case ParamKind::Dense: return "Dense"; - case ParamKind::TernaryLatent: return "TernaryLatent"; - case ParamKind::TernaryPacked2: return "TernaryPacked2"; - case ParamKind::TernaryPacked3: return "TernaryPacked3"; - case ParamKind::I8Packed4: return "I8Packed4"; + case ParamKind::Dense: + return "Dense"; + case ParamKind::TernaryLatent: + return "TernaryLatent"; + case ParamKind::TernaryPacked2: + return "TernaryPacked2"; + case ParamKind::TernaryPacked3: + return "TernaryPacked3"; + case ParamKind::I8Packed4: + return "I8Packed4"; } return "?"; } [[nodiscard]] inline const char* model_role_name(ModelRole r) noexcept -{ - return (r == ModelRole::Checkpoint) ? "checkpoint(.nnckpt)" : "model(.nnmodel)"; -} +{ return (r == ModelRole::Checkpoint) ? "checkpoint(.nnckpt)" : "model(.nnmodel)"; } // ── 参数载入意图（P1-7）──────────────────────────────────────────────── // 同一个文件（三值层为打包形态）有两种载入方式——**角色**回答"这份文件是什么"， @@ -297,22 +290,23 @@ class Layer // 绑定读取：未 init / init 失败 / 引擎已销毁都是调用方错误 → fail-fast， // 先打**调用点**（source_location 默认参在调用处求值 → 定位到层的哪个方法）。 - [[nodiscard]] ComputeEngine& engine_ref( - std::source_location loc = std::source_location::current()) const + [[nodiscard]] ComputeEngine& + engine_ref(std::source_location loc = std::source_location::current()) const { if (binding_state_ == BindingState::failed) { - const std::string msg = std::string("[layer-lifetime] ") + layer_name() + + const std::string msg = + std::string("[layer-lifetime] ") + layer_name() + ": 上一次 init 失败（层处于半初始化状态）—— 修复原因后重新 init(engine)"; - detail::fatal_at("NN_FATAL", loc.file_name(), - static_cast(loc.line()), msg.c_str()); + detail::fatal_at("NN_FATAL", loc.file_name(), static_cast(loc.line()), + msg.c_str()); } if (!engine_) { const std::string msg = std::string("[layer-lifetime] ") + layer_name() + - ": 未绑定引擎（先调 init(engine)）"; - detail::fatal_at("NN_FATAL", loc.file_name(), - static_cast(loc.line()), msg.c_str()); + ": 未绑定引擎（先调 init(engine)）"; + detail::fatal_at("NN_FATAL", loc.file_name(), static_cast(loc.line()), + msg.c_str()); } return *engine_; } @@ -346,10 +340,11 @@ class Layer { if (binding_state_ == BindingState::ready) { - NN_UNWRAP(false, - std::string("set_precision_profile: 层已 init，精度 profile 之后不可变（层=") - + layer_name() + "）—— 请在 init(engine) 之前注入" - "（docs/internal/21-quantized-weights.md §4.3.2(b)）"); + NN_UNWRAP(false, std::string( + "set_precision_profile: 层已 init，精度 profile 之后不可变（层=") + + layer_name() + + "）—— 请在 init(engine) 之前注入" + "（docs/internal/21-quantized-weights.md §4.3.2(b)）"); } p_ = profile; for (Layer* c : param_children()) @@ -360,9 +355,7 @@ class Layer // ── 层能力声明（R1，§4.2）────────────────────────────────────────── // 默认 = 现状（f16/f32 × 四槽）→ 既有层零改动。声明支持的层 override。 [[nodiscard]] virtual PrecisionSupport precision_support() const - { - return PrecisionSupport::rounding_all_slots(); - } + { return PrecisionSupport::rounding_all_slots(); } // forward/backward 只写一次，CPU/GPU 由引擎实现自动分发（引擎来自 init 绑定） [[nodiscard]] virtual Result forward(const Tensor& input) = 0; @@ -375,24 +368,23 @@ class Layer // ModelSpec 中登记的 schema 显式传入。 [[nodiscard]] virtual Result output_schema() const { - return TensorSchema::create( - std::string(layer_name()) + ".output", - {tensor_field("output")}); + return TensorSchema::create(std::string(layer_name()) + ".output", + {tensor_field("output")}); } - [[nodiscard]] virtual Result forward_named( - const Tensor& input) + [[nodiscard]] virtual Result forward_named(const Tensor& input) { NN_TRY(schema, output_schema()); return forward_named(input, *schema); } - [[nodiscard]] virtual Result forward_named( - const Tensor& input, const TensorSchema& schema) + [[nodiscard]] virtual Result forward_named(const Tensor& input, + const TensorSchema& schema) { if (!schema.valid() || schema.size() != 1) NN_FAIL(std::string("Layer::forward_named: 默认单输出层要求 schema 恰有" - "一个字段（层=") + layer_name() + "）"); + "一个字段（层=") + + layer_name() + "）"); NN_TRY(out, forward(input)); std::vector values; values.push_back(NamedTensor{ @@ -402,22 +394,23 @@ class Layer return NamedTensorSet::create(schema, std::move(values)); } - [[nodiscard]] virtual Result backward_named( - const GradientBundle& gradients) + [[nodiscard]] virtual Result backward_named(const GradientBundle& gradients) { NN_TRY(schema, output_schema()); return backward_named(gradients, *schema); } - [[nodiscard]] virtual Result backward_named( - const GradientBundle& gradients, const TensorSchema& schema) + [[nodiscard]] virtual Result backward_named(const GradientBundle& gradients, + const TensorSchema& schema) { if (!schema.valid() || schema.size() != 1) NN_FAIL(std::string("Layer::backward_named: 默认单输出层要求 schema " - "恰有一个字段（层=") + layer_name() + "）"); + "恰有一个字段（层=") + + layer_name() + "）"); if (gradients.schema() != schema) NN_FAIL(std::string("Layer::backward_named: gradient schema 不匹配" - "（层=") + layer_name() + "）"); + "（层=") + + layer_name() + "）"); const FieldHandle handle{gradients.schema().fingerprint(), 0}; NN_TRY(g, gradients.require(handle)); return backward(**g); @@ -508,7 +501,8 @@ class Layer auto children = param_children(); if (children.empty()) return param_count(); std::size_t child_total = 0; - for (Layer* c : children) child_total += c->param_count(); + for (Layer* c : children) + child_total += c->param_count(); const std::size_t total = param_count(); return total >= child_total ? total - child_total : 0; } @@ -533,8 +527,7 @@ class Layer NN_FAIL(std::string("import_param: 形状不匹配（层=") + layer_name() + "，期望 " + std::to_string(dst.rows()) + "x" + std::to_string(dst.cols()) + "，实际 " + std::to_string(e.rows) + "x" + std::to_string(e.cols) + "）"); - if (e.payload.precision() == dst.precision()) - return engine_ref().copy_into(dst, e.payload); + if (e.payload.precision() == dst.precision()) return engine_ref().copy_into(dst, e.payload); return engine_ref().cast_into(e.payload, dst); } @@ -549,8 +542,7 @@ class Layer [[nodiscard]] static std::size_t ref_count_impl_(const Tensor&) noexcept { return 1; } [[nodiscard]] static std::size_t ref_count_impl_(const std::vector& v) noexcept { return v.size(); } - static void append_refs_impl_(std::vector& out, Tensor& t) - { out.push_back(t); } + static void append_refs_impl_(std::vector& out, Tensor& t) { out.push_back(t); } static void append_refs_impl_(std::vector& out, const std::vector& v) { out.insert(out.end(), v.begin(), v.end()); } @@ -566,8 +558,8 @@ class Layer // 把「对块容器逐块调用某成员函数」的结果拼起来（三个模型的 blocks_ 级联共用）。 // RefGetter = 指向 Layer 子类 parameters / param_gradients 的成员函数指针。 template - [[nodiscard]] static std::vector collect_block_refs_( - Blocks& blocks, RefGetter getter) + [[nodiscard]] static std::vector collect_block_refs_(Blocks& blocks, + RefGetter getter) { std::vector out; for (auto& b : blocks) @@ -588,40 +580,34 @@ class Layer // add_prec = 两处残差相加的**表达式输出精度**。各块历史行为不同（GPT/ // Transformer 传 p.compute；RAPT 的残差相加历史上未传 profile、恒 F32）， // 故显式传入以保持逐位不变——不要"顺手统一"成 p.compute。 - [[nodiscard]] static Result prenorm_residual_backward_( - ComputeEngine& engine, const Tensor& grad_output, - Layer& norm1, Layer& sub1, Layer& norm2, Layer& sub2, - Precision add_prec) + [[nodiscard]] static Result + prenorm_residual_backward_(ComputeEngine& engine, const Tensor& grad_output, Layer& norm1, + Layer& sub1, Layer& norm2, Layer& sub2, Precision add_prec) { NN_TRY(grad_s2, sub2.backward(grad_output)); NN_TRY(b_n2, norm2.backward(*grad_s2)); - auto grad_x2 = dsl::compute(engine, - dsl::leaf(grad_output) + dsl::leaf(*b_n2), - grad_output.rows(), grad_output.cols(), add_prec); + auto grad_x2 = dsl::compute(engine, dsl::leaf(grad_output) + dsl::leaf(*b_n2), + grad_output.rows(), grad_output.cols(), add_prec); NN_TRY_CHECK(grad_x2); NN_TRY(grad_s1, sub1.backward(*grad_x2)); NN_TRY(b_n1, norm1.backward(*grad_s1)); - return dsl::compute(engine, - dsl::leaf(*grad_x2) + dsl::leaf(*b_n1), - grad_x2->rows(), grad_x2->cols(), add_prec); + return dsl::compute(engine, dsl::leaf(*grad_x2) + dsl::leaf(*b_n1), grad_x2->rows(), + grad_x2->cols(), add_prec); } // ── Pre-Norm 残差块前向的一个残差分支 ───────────────────────────────── // 返回 x + sub(norm(x))（GPT/Transformer/RAPT 三块的两个残差分支共用）。 // add_prec 同上：显式传以保持各块历史精度行为逐位不变。 - [[nodiscard]] static Result prenorm_residual_forward_( - ComputeEngine& engine, const Tensor& x, Layer& norm, Layer& sub, - Precision add_prec) + [[nodiscard]] static Result prenorm_residual_forward_(ComputeEngine& engine, + const Tensor& x, Layer& norm, + Layer& sub, Precision add_prec) { NN_TRY(n, norm.forward(x)); NN_TRY(s, sub.forward(*n)); - return dsl::compute(engine, - dsl::leaf(x) + dsl::leaf(*s), - x.rows(), x.cols(), add_prec); + return dsl::compute(engine, dsl::leaf(x) + dsl::leaf(*s), x.rows(), x.cols(), add_prec); } public: - // 梯度清零（每个训练 step 开始前调用） [[nodiscard]] virtual Result zero_grad() { @@ -673,8 +659,8 @@ class Layer auto keepalive = engine.weak_from_this().lock(); if (!keepalive) NN_FAIL(std::string("Layer::init: 引擎必须由 shared_ptr 管理" - "（使用 CpuEngine::create() / GpuEngine::create()）；层=") - + layer_name()); + "（使用 CpuEngine::create() / GpuEngine::create()）；层=") + + layer_name()); if (binding_state_ == BindingState::ready) { @@ -703,8 +689,7 @@ class Layer // 便捷重载：引擎工厂返回 shared_ptr，调用方不必手写 *engine。 [[nodiscard]] Result init(const std::shared_ptr& engine) { - if (!engine) - NN_FAIL("Layer::init: engine 为空"); + if (!engine) NN_FAIL("Layer::init: engine 为空"); return init(*engine); } @@ -721,17 +706,15 @@ class Layer // 的 latent 恒 f32）。 [[nodiscard]] Result check_param_precision_() { - if (precision_bytes(p_.param) == 0) - return {}; + if (precision_bytes(p_.param) == 0) return {}; for (const TensorRef& ref : parameters()) { const Tensor& t = ref.get(); if (t.precision() != p_.param) - NN_FAIL(std::string("参数存储精度与 param 槽不一致: 层 ") - + layer_name() + " 的 " + t.shape_str() + " 实际 " - + precision_name(t.precision()) + "，但 param=" - + precision_name(p_.param) - + "（硬编码 F32 会让 --f16 在该层静默失效；见 docs/internal/18 P1-6 R2）"); + NN_FAIL(std::string("参数存储精度与 param 槽不一致: 层 ") + layer_name() + " 的 " + + t.shape_str() + " 实际 " + precision_name(t.precision()) + + "，但 param=" + precision_name(p_.param) + + "（硬编码 F32 会让 --f16 在该层静默失效；见 docs/internal/18 P1-6 R2）"); } return {}; } @@ -742,12 +725,10 @@ class Layer { const PrecisionSupport sup = precision_support(); const PrecisionProfile& p = p_; - if (!sup.param.has(p.param)) - return precision_slot_error_("param", p.param, sup.param); + if (!sup.param.has(p.param)) return precision_slot_error_("param", p.param, sup.param); if (!sup.compute.has(p.compute)) return precision_slot_error_("compute", p.compute, sup.compute); - if (!sup.stable.has(p.stable)) - return precision_slot_error_("stable", p.stable, sup.stable); + if (!sup.stable.has(p.stable)) return precision_slot_error_("stable", p.stable, sup.stable); if (!sup.optimizer.has(p.optimizer)) return precision_slot_error_("optimizer", p.optimizer, sup.optimizer); return {}; @@ -757,14 +738,12 @@ class Layer [[nodiscard]] Result precision_slot_error_(const char* slot, Precision got, const PrecisionSet& allowed) const { - NN_FAIL(std::string("精度能力校验失败: 层 ") + layer_name() - + " 不支持 " + slot + "=" + precision_name(got) - + "（该槽允许: " + allowed.to_string() - + "）—— 见 docs/internal/21-quantized-weights.md §4.2"); + NN_FAIL(std::string("精度能力校验失败: 层 ") + layer_name() + " 不支持 " + slot + "=" + + precision_name(got) + "（该槽允许: " + allowed.to_string() + + "）—— 见 docs/internal/21-quantized-weights.md §4.2"); } public: - // 梯度检查点（激活重计算）契约 ────────────────────────────────── // checkpoint_mode_ = true 时，forward 不保留中间激活（供 L1 激活重计算）； // forward_recompute 重算 forward 并重建缓存（供 backward 使用）。 @@ -807,7 +786,6 @@ class Layer [[nodiscard]] virtual Result init_impl(ComputeEngine& /*engine*/) { return {}; } public: - // 梯度检查点粒度（默认 0 = 不启用；由 GPTModel 等 override） virtual void set_checkpoint_every(std::size_t /*stride*/) {} @@ -877,8 +855,8 @@ class ActivationOffloader } // 导出：逐个写入 slab，写入后置空成员张量（数据已在 host slab） - [[nodiscard]] Result export_activations( - ComputeEngine& engine, std::vector refs) + [[nodiscard]] Result export_activations(ComputeEngine& engine, + std::vector refs) { if (!enabled_) return {}; // 状态机：Offloaded 状态只能经 import（消费）或 discard（作废）离开。 @@ -891,8 +869,7 @@ class ActivationOffloader // 写回 refs_），因此在这里**压缩**，而不是在写入循环里 skip 无效项。 refs_.clear(); for (auto& ref : refs) - if (ref.get().valid()) - refs_.push_back(ref); + if (ref.get().valid()) refs_.push_back(ref); offsets_.clear(); shapes_.clear(); // slab 容量校验：slab 必须容纳**本次**导出的激活总量（批大小/ @@ -908,13 +885,17 @@ class ActivationOffloader // 不一致——该组合无测试、无精度协议，静默漂移不可接受，一律 // fail-fast（CLI 侧另有更早的配置期拒绝）。 if (ref.get().precision() != Precision::F32) - NN_FAIL("ActivationOffloader::export_activations: 仅支持 F32 激活（当前 " - + std::string(precision_name(ref.get().precision())) - + "）；offload slab 恒 f32，混合精度/三值精度请改用" - " --checkpoint-every"); + NN_FAIL("ActivationOffloader::export_activations: 仅支持 F32 激活（当前 " + + std::string(precision_name(ref.get().precision())) + + "）；offload slab 恒 f32，混合精度/三值精度请改用" + " --checkpoint-every"); needed += ref.get().size(); } - if (needed == 0) { offloaded_ = false; return {}; } + if (needed == 0) + { + offloaded_ = false; + return {}; + } // CPU 引擎的 offload 原语是 no-op（restore 返回 1×1）——export 会把成员 // 缓存置空、import 再把 1×1 写回，等于毒化 backward。CLI 早已忽略该组合， // 这里对**所有**非 CLI 调用方统一 fail-fast（唯一的正确做法是别开）。 @@ -946,8 +927,8 @@ class ActivationOffloader if (!offloaded_) return {}; for (std::size_t i = 0; i < offsets_.size(); ++i) { - auto t = engine.offload_restore(slab_, offsets_[i], - shapes_[i].first, shapes_[i].second); + auto t = + engine.offload_restore(slab_, offsets_[i], shapes_[i].first, shapes_[i].second); NN_TRY_CHECK(t); refs_[i].get() = std::move(*t); } @@ -957,16 +938,11 @@ class ActivationOffloader // 实际 slab 字节数（诊断用；未创建时 0） [[nodiscard]] std::size_t slab_bytes() const noexcept - { - return slab_.valid() ? slab_.size() * sizeof(float) : 0; - } + { return slab_.valid() ? slab_.size() * sizeof(float) : 0; } }; -[[nodiscard]] inline Result clone_tensor( - ComputeEngine& engine, const Tensor& src) -{ - return engine.clone(src); -} +[[nodiscard]] inline Result clone_tensor(ComputeEngine& engine, const Tensor& src) +{ return engine.clone(src); } // ══════════════════════════════════════════════════════════════════════════ // 采样：重复惩罚 → temperature 缩放 → top-k / top-p 截断 → @@ -1032,8 +1008,7 @@ struct SamplingConfig // ── ⓪ 外部 logits 变换钩子（issue #34）：全部库内策略之前 ──────── // 未设置（默认）时零开销，RNG 消耗序与旧路径逐位一致。 - if (cfg.logits_transform) - cfg.logits_transform(std::span(logits.data(), logits.size())); + if (cfg.logits_transform) cfg.logits_transform(std::span(logits.data(), logits.size())); // ── ① 重复惩罚（去重；判定只依赖 history 的内容，不依赖扫描顺序）── if (cfg.repeat_penalty > 0.0 && cfg.repeat_penalty != 1.0 && !history.empty()) @@ -1143,7 +1118,11 @@ struct SamplingConfig for (const std::size_t v : cand) { cumulative += logits[v]; - if (r <= cumulative) { next = v; break; } + if (r <= cumulative) + { + next = v; + break; + } } return next; } diff --git a/include/neuralnet.cpp/compute_layer_bitlinear.hpp b/include/neuralnet.cpp/compute_layer_bitlinear.hpp index 78d2c7b7..4eb9f9df 100644 --- a/include/neuralnet.cpp/compute_layer_bitlinear.hpp +++ b/include/neuralnet.cpp/compute_layer_bitlinear.hpp @@ -97,7 +97,9 @@ class BitLinear final : public Layer public: BitLinear(std::size_t in_features, std::size_t out_features) - : in_features_(in_features), out_features_(out_features) {} + : in_features_(in_features), out_features_(out_features) + { + } [[nodiscard]] const char* layer_name() const noexcept override { return "BitLinear"; } @@ -157,10 +159,7 @@ class BitLinear final : public Layer } // ── 参数视图：latent（f32）+ 偏置（f32）────────────────────────────── - [[nodiscard]] std::vector parameters() override - { - return {latent_w_, b_}; - } + [[nodiscard]] std::vector parameters() override { return {latent_w_, b_}; } [[nodiscard]] std::vector param_gradients() override { ensure_grad_(); @@ -226,85 +225,79 @@ class BitLinear final : public Layer std::to_string(index) + "）"); switch (entry.kind) { - case ParamKind::TernaryLatent: + case ParamKind::TernaryLatent: + { + if (!latent_w_.valid()) { - if (!latent_w_.valid()) - { - // 从打包只读态恢复可训练：重建 f32 latent（零初始化）并回到 - // 默认打包缓冲形态；量化缓存作废（首次 forward 会重算）。 - latent_w_ = engine_ref().create_tensor( - out_features_, in_features_, Precision::F32, InitSpec::zero()); - NN_UNWRAP(latent_w_.valid(), "BitLinear: latent 恢复失败"); - packed_only_ = false; - use_packed_ = true; - ensure_quant_buffer_(); - } - quant_valid_ = false; - NN_TRY_CHECK(import_into_(latent_w_, entry)); - // ── P1-7：反量化 warm start（aux = 文件里的 τ）──────────────── - // 反量化本身在序列化层（宿主侧逐位解码；L2 不碰宿主字节/Matrix）， - // 层这边只做**播种**：直接用文件 τ 重建打包派生物。这样首次 - // forward 的 wq/τ 与文件里的模型**逐位相同**，不经过"用自己的 - // absmean 重算 τ"那一步——重算只会得到同一个值的 ±1 ulp，却会让 - // "首次 forward 逐位复现"这条验收失去确定性。 - // τ 是权威 ⇒ 重算出的 trit 必然与文件一致（见 ternary_pack.hpp - // 的 latent_row_scale：非零位置 |L| = τ·K/n ≥ τ > 0.5τ）。 - if (!entry.aux.valid()) return {}; // .nnckpt 语义：无 τ，首 forward 自行量化 - if (entry.aux.precision() != Precision::F32) - NN_FAIL("BitLinear::import_param: 反量化 warm start 的 τ 必须是 f32"); - if (entry.aux.rows() != out_features_ || entry.aux.cols() != 1) - NN_FAIL("BitLinear::import_param: τ 形状与层不匹配（期望 " + - std::to_string(out_features_) + "x1）"); - tau_ = entry.aux; - // 播种的是**打包**缓存 ⇒ 把缓冲形态钉成打包（覆盖此前的 A/B - // `set_ternary_pack_enabled(false)`；要改形态在载入之后再调它）。 - force_pack_ = 1; + // 从打包只读态恢复可训练：重建 f32 latent（零初始化）并回到 + // 默认打包缓冲形态；量化缓存作废（首次 forward 会重算）。 + latent_w_ = engine_ref().create_tensor(out_features_, in_features_, Precision::F32, + InitSpec::zero()); + NN_UNWRAP(latent_w_.valid(), "BitLinear: latent 恢复失败"); + packed_only_ = false; use_packed_ = true; ensure_quant_buffer_(); - NN_TRY(packed, engine_ref().trit_pack(latent_w_, tau_)); - wq_packed_ = std::move(*packed); - quant_epoch_ = latent_w_.write_epoch(); - quant_valid_ = true; - ++quantize_calls_; - return {}; } - case ParamKind::TernaryPacked2: - { - if (!entry.aux.valid()) - NN_FAIL("BitLinear::import_param: 打包形态缺少 τ（aux）"); - const std::size_t wpr = nn::trit_pack::words_per_row(in_features_); - if (entry.rows != out_features_ || - entry.payload.rows() != out_features_ || - entry.payload.cols() != wpr) - NN_FAIL("BitLinear::import_param: 打包字形状与层不匹配（期望 " + - std::to_string(out_features_) + "x" + std::to_string(wpr) + "）"); - if (entry.aux.rows() != out_features_ || entry.aux.cols() != 1) - NN_FAIL("BitLinear::import_param: τ 形状与层不匹配（期望 " + - std::to_string(out_features_) + "x1）"); - // 打包只读态：只保留打包字 + τ（推理足迹 0.25 B/权重 + τ） - force_pack_ = 1; - use_packed_ = true; - wq_packed_ = entry.payload; - tau_ = entry.aux; - quant_valid_ = true; - packed_only_ = true; - latent_w_ = Tensor{}; - wq_ = Tensor{}; - grad_w_ = Tensor{}; - grad_b_ = Tensor{}; - return {}; - } - default: - NN_FAIL(std::string("BitLinear::import_param: 不接受该参数形态（kind=") + - param_kind_name(entry.kind) + - "，可接受 {TernaryLatent, TernaryPacked2}）"); + quant_valid_ = false; + NN_TRY_CHECK(import_into_(latent_w_, entry)); + // ── P1-7：反量化 warm start（aux = 文件里的 τ）──────────────── + // 反量化本身在序列化层（宿主侧逐位解码；L2 不碰宿主字节/Matrix）， + // 层这边只做**播种**：直接用文件 τ 重建打包派生物。这样首次 + // forward 的 wq/τ 与文件里的模型**逐位相同**，不经过"用自己的 + // absmean 重算 τ"那一步——重算只会得到同一个值的 ±1 ulp，却会让 + // "首次 forward 逐位复现"这条验收失去确定性。 + // τ 是权威 ⇒ 重算出的 trit 必然与文件一致（见 ternary_pack.hpp + // 的 latent_row_scale：非零位置 |L| = τ·K/n ≥ τ > 0.5τ）。 + if (!entry.aux.valid()) return {}; // .nnckpt 语义：无 τ，首 forward 自行量化 + if (entry.aux.precision() != Precision::F32) + NN_FAIL("BitLinear::import_param: 反量化 warm start 的 τ 必须是 f32"); + if (entry.aux.rows() != out_features_ || entry.aux.cols() != 1) + NN_FAIL("BitLinear::import_param: τ 形状与层不匹配（期望 " + + std::to_string(out_features_) + "x1）"); + tau_ = entry.aux; + // 播种的是**打包**缓存 ⇒ 把缓冲形态钉成打包（覆盖此前的 A/B + // `set_ternary_pack_enabled(false)`；要改形态在载入之后再调它）。 + force_pack_ = 1; + use_packed_ = true; + ensure_quant_buffer_(); + NN_TRY(packed, engine_ref().trit_pack(latent_w_, tau_)); + wq_packed_ = std::move(*packed); + quant_epoch_ = latent_w_.write_epoch(); + quant_valid_ = true; + ++quantize_calls_; + return {}; + } + case ParamKind::TernaryPacked2: + { + if (!entry.aux.valid()) NN_FAIL("BitLinear::import_param: 打包形态缺少 τ（aux）"); + const std::size_t wpr = nn::trit_pack::words_per_row(in_features_); + if (entry.rows != out_features_ || entry.payload.rows() != out_features_ || + entry.payload.cols() != wpr) + NN_FAIL("BitLinear::import_param: 打包字形状与层不匹配（期望 " + + std::to_string(out_features_) + "x" + std::to_string(wpr) + "）"); + if (entry.aux.rows() != out_features_ || entry.aux.cols() != 1) + NN_FAIL("BitLinear::import_param: τ 形状与层不匹配（期望 " + + std::to_string(out_features_) + "x1）"); + // 打包只读态：只保留打包字 + τ（推理足迹 0.25 B/权重 + τ） + force_pack_ = 1; + use_packed_ = true; + wq_packed_ = entry.payload; + tau_ = entry.aux; + quant_valid_ = true; + packed_only_ = true; + latent_w_ = Tensor{}; + wq_ = Tensor{}; + grad_w_ = Tensor{}; + grad_b_ = Tensor{}; + return {}; + } + default: + NN_FAIL(std::string("BitLinear::import_param: 不接受该参数形态（kind=") + + param_kind_name(entry.kind) + "，可接受 {TernaryLatent, TernaryPacked2}）"); } } - [[nodiscard]] bool grad_buffers_allocated() const noexcept override - { - return grad_w_.valid(); - } + [[nodiscard]] bool grad_buffers_allocated() const noexcept override { return grad_w_.valid(); } // ── 量化派生物只读视图（诊断 / 测试 / 序列化对拍用）──────────────────── [[nodiscard]] const Tensor& latent_weights() const noexcept { return latent_w_; } @@ -313,10 +306,7 @@ class BitLinear final : public Layer // 打包路径（P2/F3）：本层是否用 2 bit/权重的打包缓冲（`quantized_weights()` // 在打包路径下为空 —— 用 `quantized_weights_packed()` + ternary_pack.hpp 解码） [[nodiscard]] bool uses_packed_weights() const noexcept { return use_packed_; } - [[nodiscard]] const Tensor& quantized_weights_packed() const noexcept - { - return wq_packed_; - } + [[nodiscard]] const Tensor& quantized_weights_packed() const noexcept { return wq_packed_; } // 量化实际执行次数（诊断/测试：缓存是否生效看它，不看 forward 次数） [[nodiscard]] std::uint64_t quantize_calls() const noexcept { return quantize_calls_; } @@ -326,16 +316,14 @@ class BitLinear final : public Layer // **逐位复现**训练时的有效权重 τ∘W（重新量化会得到 mean|trit|·τ 的比例差异）。 [[nodiscard]] Result load_packed_weights(Tensor packed_words, Tensor tau) { - if (!packed_words.valid() || !tau.valid()) - NN_FAIL("BitLinear: 部署权重为空张量"); + if (!packed_words.valid() || !tau.valid()) NN_FAIL("BitLinear: 部署权重为空张量"); if (packed_words.rows() != out_features_ || packed_words.cols() != nn::trit_pack::words_per_row(in_features_)) - NN_FAIL("BitLinear: 部署打包字形状与层不匹配（期望 " + - std::to_string(out_features_) + "x" + - std::to_string(nn::trit_pack::words_per_row(in_features_)) + "）"); + NN_FAIL("BitLinear: 部署打包字形状与层不匹配（期望 " + std::to_string(out_features_) + + "x" + std::to_string(nn::trit_pack::words_per_row(in_features_)) + "）"); if (tau.rows() != out_features_ || tau.cols() != 1) - NN_FAIL("BitLinear: 部署 τ 形状与层不匹配（期望 " + - std::to_string(out_features_) + "x1）"); + NN_FAIL("BitLinear: 部署 τ 形状与层不匹配（期望 " + std::to_string(out_features_) + + "x1）"); force_pack_ = 1; use_packed_ = true; wq_packed_ = std::move(packed_words); @@ -384,35 +372,29 @@ class BitLinear final : public Layer [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); - if (input.rows() != in_features_) - NN_FAIL("bitlinear forward: input shape mismatch"); + if (input.rows() != in_features_) NN_FAIL("bitlinear forward: input shape mismatch"); // 部署态：只允许"打包字 + τ"这条路（引擎不支持三值 GEMM 就没法跑） if (packed_only_ && !(use_packed_ && wq_packed_.valid() && tau_.valid())) NN_FAIL("BitLinear(packed-only)：部署态缺少打包权重/τ（载入失败或缓冲被清）"); - if (!checkpoint_mode_) - input_cache_ = input; + if (!checkpoint_mode_) input_cache_ = input; // 缓冲形态：默认跟随引擎能力（分块解码版下训练/解码都不亏，见成员注释）； // `set_ternary_pack_enabled` 之后走强制模式。换形态即作废量化缓存。 sync_buffer_mode_(); // τ / wq 只在 latent 权重变过时重算（同权重下的 N 次 forward 共享）。 // 重算失败：直接上抛（不静默沿用上一步的旧缓冲）。 - if (needs_quantize_()) - { - NN_TRY(q, quantize_weights_()); - } + if (needs_quantize_()) { NN_TRY(q, quantize_weights_()); } // 打包路径：Y = (τ∘W)·X 是一趟量化 GEMM（权重读 0.25 B/元素），偏置用 // 一次行广播加（原语不带 bias，见 compute_engine.hpp 的说明）。 if (use_packed_) { NN_TRY(y, engine.matmul_q(wq_packed_, tau_, input, p_.compute)); - return dsl::compute(engine, - dsl::leaf(*y) + dsl::row_broadcast(b_), - out_features_, input.cols(), p_.compute); + return dsl::compute(engine, dsl::leaf(*y) + dsl::row_broadcast(b_), out_features_, + input.cols(), p_.compute); } return dsl::compute(engine, - dsl::matmul(wq_, input, false, false) * dsl::row_broadcast(tau_) - + dsl::row_broadcast(b_), - out_features_, input.cols(), p_.compute); + dsl::matmul(wq_, input, false, false) * dsl::row_broadcast(tau_) + + dsl::row_broadcast(b_), + out_features_, input.cols(), p_.compute); } // ── backward：STE（把 wq 当常数；梯度回 f32 latent）────────────────── @@ -433,39 +415,35 @@ class BitLinear final : public Layer Tensor grad_input; if (use_packed_) { - NN_TRY(gi, engine.matmul_q_t(wq_packed_, tau_, grad_output, - in_features_, p_.compute)); + NN_TRY(gi, engine.matmul_q_t(wq_packed_, tau_, grad_output, in_features_, p_.compute)); grad_input = std::move(*gi); } else { // f16 稠密缓冲路径：τ 先落到 dY 的行上再左乘（等价形式） - NN_TRY(dy_scaled, dsl::compute(engine, - dsl::leaf(grad_output) * dsl::row_broadcast(tau_), - out_features_, grad_output.cols(), p_.compute)); - auto gi = dsl::compute(engine, - dsl::matmul(wq_, *dy_scaled, /*transA=*/true, false), - in_features_, grad_output.cols(), p_.compute); + NN_TRY(dy_scaled, + dsl::compute(engine, dsl::leaf(grad_output) * dsl::row_broadcast(tau_), + out_features_, grad_output.cols(), p_.compute)); + auto gi = dsl::compute(engine, dsl::matmul(wq_, *dy_scaled, /*transA=*/true, false), + in_features_, grad_output.cols(), p_.compute); NN_TRY_CHECK(gi); grad_input = std::move(*gi); } // dW += (dY·Xᵀ)∘τ（STE：∂wq/∂W ≡ 1，再乘去量化尺度 τ） // matmul 段 + 逐行缩放 + 原地累加融合为单次 dispatch（与 Linear 同构）。 - auto grad_w_acc = dsl::compute_into(engine, - dsl::leaf(grad_w_) - + dsl::matmul(grad_output, input_cache_, false, true) - * dsl::row_broadcast(tau_), + auto grad_w_acc = dsl::compute_into( + engine, + dsl::leaf(grad_w_) + + dsl::matmul(grad_output, input_cache_, false, true) * dsl::row_broadcast(tau_), grad_w_); NN_TRY_CHECK(grad_w_acc); // db += Σ_batch dY（归约向量 + 原地累加两步；与 Linear 同款理由： // 归约向量输出契约要求输出链只经归约/广播视图访问输入） - NN_TRY(gb, dsl::compute_reduce(engine, - dsl::row_reduce_sum(dsl::leaf(grad_output)), - grad_output.rows(), grad_output.cols(), p_.compute)); - auto r2 = dsl::compute_into(engine, - dsl::leaf(grad_b_) + dsl::leaf(*gb), grad_b_); + NN_TRY(gb, dsl::compute_reduce(engine, dsl::row_reduce_sum(dsl::leaf(grad_output)), + grad_output.rows(), grad_output.cols(), p_.compute)); + auto r2 = dsl::compute_into(engine, dsl::leaf(grad_b_) + dsl::leaf(*gb), grad_b_); NN_TRY_CHECK(r2); return grad_input; @@ -493,9 +471,9 @@ class BitLinear final : public Layer if (use_packed_) { if (!wq_packed_.valid()) - wq_packed_ = engine.create_tensor( - out_features_, nn::trit_pack::words_per_row(in_features_), - Precision::F32, InitSpec::zero()); + wq_packed_ = + engine.create_tensor(out_features_, nn::trit_pack::words_per_row(in_features_), + Precision::F32, InitSpec::zero()); // 终止式检查（本函数返回 void） NN_UNWRAP(wq_packed_.valid(), "BitLinear: 打包量化缓冲初始化失败"); wq_ = Tensor{}; @@ -520,12 +498,11 @@ class BitLinear final : public Layer { if (grad_w_.valid()) return; ComputeEngine& engine = engine_ref(); - grad_w_ = engine.create_tensor(out_features_, in_features_, Precision::F32, - InitSpec::zero()); + grad_w_ = + engine.create_tensor(out_features_, in_features_, Precision::F32, InitSpec::zero()); grad_b_ = engine.create_tensor(out_features_, 1, Precision::F32, InitSpec::zero()); // 终止式检查（本函数返回 void，不能用 NN_FAIL 的 return 语义） - NN_UNWRAP(grad_w_.valid() && grad_b_.valid(), - "BitLinear: 梯度缓冲（惰性）初始化失败"); + NN_UNWRAP(grad_w_.valid() && grad_b_.valid(), "BitLinear: 梯度缓冲（惰性）初始化失败"); } // 量化派生物是否还有效：写版本相等 ⇒ latent 权重未被写过。 @@ -556,9 +533,10 @@ class BitLinear final : public Layer ComputeEngine& engine = engine_ref(); const Scalar inv_in = Scalar{1} / static_cast(in_features_); - NN_TRY(tau, dsl::compute_reduce(engine, - dsl::row_reduce_sum(dsl::abs(dsl::leaf(latent_w_)) * dsl::rparam(inv_in)), - out_features_, in_features_, Precision::F32)); + NN_TRY(tau, dsl::compute_reduce( + engine, + dsl::row_reduce_sum(dsl::abs(dsl::leaf(latent_w_)) * dsl::rparam(inv_in)), + out_features_, in_features_, Precision::F32)); tau_ = std::move(*tau); // 打包路径（F3）：一趟 trit_pack 直接出 2 bit/权重容器（8.25 B/元素 @@ -579,9 +557,9 @@ class BitLinear final : public Layer // 唯一差异正在这个测度零的等值点（RoundClip 按 round-half-* 取值）。 auto lo = dsl::row_broadcast(tau_) * dsl::rparam(Scalar{-0.5}); auto hi = dsl::row_broadcast(tau_) * dsl::rparam(Scalar{0.5}); - auto r = dsl::compute_into(engine, - dsl::select(dsl::leaf(latent_w_) > hi, - Scalar{1}, + auto r = dsl::compute_into( + engine, + dsl::select(dsl::leaf(latent_w_) > hi, Scalar{1}, dsl::select(dsl::leaf(latent_w_) < lo, Scalar{-1}, Scalar{0})), wq_); NN_TRY_CHECK(r); diff --git a/include/neuralnet.cpp/compute_layer_bundle.hpp b/include/neuralnet.cpp/compute_layer_bundle.hpp index 8b47ed54..29196ab0 100644 --- a/include/neuralnet.cpp/compute_layer_bundle.hpp +++ b/include/neuralnet.cpp/compute_layer_bundle.hpp @@ -36,11 +36,8 @@ class RowBandPacking std::vector bands_; std::size_t total_rows_ = 0; - RowBandPacking(TensorSchema schema, std::vector bands, - std::size_t total_rows) - : schema_(std::move(schema)), - bands_(std::move(bands)), - total_rows_(total_rows) + RowBandPacking(TensorSchema schema, std::vector bands, std::size_t total_rows) + : schema_(std::move(schema)), bands_(std::move(bands)), total_rows_(total_rows) { } @@ -52,11 +49,10 @@ class RowBandPacking RowBandPacking& operator=(RowBandPacking&&) noexcept = default; ~RowBandPacking() = default; - [[nodiscard]] static Result create( - const TensorSchema& schema, std::vector order = {}) + [[nodiscard]] static Result create(const TensorSchema& schema, + std::vector order = {}) { - if (!schema.valid()) - NN_FAIL("RowBandPacking::create: schema 未初始化"); + if (!schema.valid()) NN_FAIL("RowBandPacking::create: schema 未初始化"); if (order.empty()) { order.reserve(schema.size()); @@ -74,8 +70,7 @@ class RowBandPacking { NN_TRY(handle, schema.field(name)); if (seen[handle->index]) - NN_FAIL("RowBandPacking::create: 打包字段 '" + name + - "' 重复"); + NN_FAIL("RowBandPacking::create: 打包字段 '" + name + "' 重复"); seen[handle->index] = true; const TensorFieldSpec& field = schema.fields()[handle->index]; if (!field.shape.rows) @@ -92,17 +87,12 @@ class RowBandPacking [[nodiscard]] bool valid() const noexcept { return schema_.valid(); } [[nodiscard]] const TensorSchema& schema() const noexcept { return schema_; } - [[nodiscard]] const std::vector& bands() const noexcept - { - return bands_; - } + [[nodiscard]] const std::vector& bands() const noexcept { return bands_; } [[nodiscard]] std::size_t total_rows() const noexcept { return total_rows_; } - [[nodiscard]] Result pack( - ComputeEngine& engine, const NamedTensorSet& outputs) const + [[nodiscard]] Result pack(ComputeEngine& engine, const NamedTensorSet& outputs) const { - if (!valid()) - NN_FAIL("RowBandPacking::pack: packing 未初始化"); + if (!valid()) NN_FAIL("RowBandPacking::pack: packing 未初始化"); if (outputs.schema() != schema_) NN_FAIL("RowBandPacking::pack: 输出 schema 与 packing 不匹配"); @@ -110,23 +100,18 @@ class RowBandPacking for (const RowBand& band : bands_) { const Tensor* value = outputs.find(band.name); - if (value == nullptr) - NN_FAIL("RowBandPacking::pack: 缺少输出端口 '" + band.name + "'"); + if (value == nullptr) NN_FAIL("RowBandPacking::pack: 缺少输出端口 '" + band.name + "'"); if (first == nullptr) first = value; - if (value->cols() != first->cols()) - NN_FAIL("RowBandPacking::pack: 端口列数不一致"); + if (value->cols() != first->cols()) NN_FAIL("RowBandPacking::pack: 端口列数不一致"); if (value->precision() != first->precision()) NN_FAIL("RowBandPacking::pack: 端口精度不一致"); if (value->rows() != band.row_count) NN_FAIL("RowBandPacking::pack: 端口 '" + band.name + "' 行数与 schema/packing 不一致"); } - if (first == nullptr) - NN_FAIL("RowBandPacking::pack: schema 没有输出端口"); + if (first == nullptr) NN_FAIL("RowBandPacking::pack: schema 没有输出端口"); - NN_TRY(packed, - engine.try_create_tensor(total_rows_, first->cols(), - first->precision())); + NN_TRY(packed, engine.try_create_tensor(total_rows_, first->cols(), first->precision())); for (const RowBand& band : bands_) { const Tensor* value = outputs.find(band.name); @@ -136,22 +121,18 @@ class RowBandPacking return packed; } - [[nodiscard]] Result unpack( - ComputeEngine& engine, const Tensor& packed) const + [[nodiscard]] Result unpack(ComputeEngine& engine, const Tensor& packed) const { - if (!valid()) - NN_FAIL("RowBandPacking::unpack: packing 未初始化"); + if (!valid()) NN_FAIL("RowBandPacking::unpack: packing 未初始化"); if (packed.rows() != total_rows_) NN_FAIL("RowBandPacking::unpack: 打包张量行数不匹配（期望 " + - std::to_string(total_rows_) + "，实际 " + - std::to_string(packed.rows()) + "）"); + std::to_string(total_rows_) + "，实际 " + std::to_string(packed.rows()) + "）"); std::vector values; values.reserve(bands_.size()); for (const RowBand& band : bands_) { - NN_TRY(part, - engine.slice_rows(packed, band.first_row, band.row_count)); + NN_TRY(part, engine.slice_rows(packed, band.first_row, band.row_count)); values.push_back(NamedTensor{ .name = band.name, .value = std::move(*part), @@ -160,22 +141,20 @@ class RowBandPacking return NamedTensorSet::create(schema_, std::move(values)); } - [[nodiscard]] Result unpack_gradients( - ComputeEngine& engine, const Tensor& packed) const + [[nodiscard]] Result unpack_gradients(ComputeEngine& engine, + const Tensor& packed) const { - if (!valid()) - NN_FAIL("RowBandPacking::unpack_gradients: packing 未初始化"); + if (!valid()) NN_FAIL("RowBandPacking::unpack_gradients: packing 未初始化"); if (packed.rows() != total_rows_) NN_FAIL("RowBandPacking::unpack_gradients: 打包张量行数不匹配" - "（期望 " + std::to_string(total_rows_) + "，实际 " + - std::to_string(packed.rows()) + "）"); + "（期望 " + + std::to_string(total_rows_) + "，实际 " + std::to_string(packed.rows()) + "）"); std::vector values; values.reserve(bands_.size()); for (const RowBand& band : bands_) { - NN_TRY(part, - engine.slice_rows(packed, band.first_row, band.row_count)); + NN_TRY(part, engine.slice_rows(packed, band.first_row, band.row_count)); values.push_back(NamedTensor{ .name = band.name, .value = std::move(*part), @@ -214,11 +193,10 @@ class OutputBundleLayer final : public Layer return nullptr; } - [[nodiscard]] Result run_forward_( - const Tensor& input, const TensorSchema& requested_schema) + [[nodiscard]] Result run_forward_(const Tensor& input, + const TensorSchema& requested_schema) { - if (requested_schema != schema_) - NN_FAIL("OutputBundleLayer: named forward schema 不匹配"); + if (requested_schema != schema_) NN_FAIL("OutputBundleLayer: named forward schema 不匹配"); // 新一轮 forward 开始后，旧的配对状态立刻失效。否则 trunk 或分支 // 中途失败时，调用方仍可能拿上一次成功的缓存执行 backward。 @@ -234,8 +212,7 @@ class OutputBundleLayer final : public Layer if (port == nullptr) { if (field.required) - NN_FAIL("OutputBundleLayer: 必需输出端口 '" + field.name + - "' 没有分支"); + NN_FAIL("OutputBundleLayer: 必需输出端口 '" + field.name + "' 没有分支"); continue; } NN_TRY(out, port->branch->forward(input)); @@ -245,50 +222,35 @@ class OutputBundleLayer final : public Layer .value = std::move(*out), }); } - if (values.empty()) - NN_FAIL("OutputBundleLayer: 本次 forward 没有产生任何输出"); + if (values.empty()) NN_FAIL("OutputBundleLayer: 本次 forward 没有产生任何输出"); NN_TRY(set, NamedTensorSet::create(schema_, std::move(values))); forward_ready_ = true; return set; } public: - explicit OutputBundleLayer(TensorSchema schema) - : schema_(std::move(schema)) - { - NN_ASSERT_ALWAYS(schema_.valid(), - "OutputBundleLayer: schema 必须有效"); - } + explicit OutputBundleLayer(TensorSchema schema) : schema_(std::move(schema)) + { NN_ASSERT_ALWAYS(schema_.valid(), "OutputBundleLayer: schema 必须有效"); } - [[nodiscard]] const char* layer_name() const noexcept override - { - return "OutputBundleLayer"; - } + [[nodiscard]] const char* layer_name() const noexcept override { return "OutputBundleLayer"; } - [[nodiscard]] Result output_schema() const override - { - return schema_; - } + [[nodiscard]] Result output_schema() const override { return schema_; } - [[nodiscard]] Result add_port( - std::string name, std::unique_ptr branch) + [[nodiscard]] Result add_port(std::string name, std::unique_ptr branch) { if (binding_state() == BindingState::ready) NN_FAIL("OutputBundleLayer::add_port: 层已 init，不能增加端口"); - if (!branch) - NN_FAIL("OutputBundleLayer::add_port: 分支为空"); + if (!branch) NN_FAIL("OutputBundleLayer::add_port: 分支为空"); NN_TRY_CHECK(schema_.field(name)); if (find_port_(name) != nullptr) - NN_FAIL("OutputBundleLayer::add_port: 输出端口 '" + name + - "' 重复"); + NN_FAIL("OutputBundleLayer::add_port: 输出端口 '" + name + "' 重复"); branch->set_precision_profile(p_); ports_.push_back(Port{std::move(name), std::move(branch)}); return {}; } - [[nodiscard]] Result set_row_band_packing( - std::vector order = {}) + [[nodiscard]] Result set_row_band_packing(std::vector order = {}) { if (binding_state() == BindingState::ready) NN_FAIL("OutputBundleLayer::set_row_band_packing: 层已 init，" @@ -298,10 +260,7 @@ class OutputBundleLayer final : public Layer return {}; } - [[nodiscard]] bool has_row_band_packing() const noexcept - { - return packing_.has_value(); - } + [[nodiscard]] bool has_row_band_packing() const noexcept { return packing_.has_value(); } void set_precision_profile(const PrecisionProfile& profile) override { @@ -326,43 +285,30 @@ class OutputBundleLayer final : public Layer return packing_->pack(engine_ref(), *set); } - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { if (!packing_) NN_FAIL("OutputBundleLayer::backward: 多输出 backwards 必须使用 " "backward_named()，或先配置 RowBandPacking"); - NN_TRY(grads, - packing_->unpack_gradients(engine_ref(), grad_output)); + NN_TRY(grads, packing_->unpack_gradients(engine_ref(), grad_output)); return backward_named(*grads); } - [[nodiscard]] Result forward_named( - const Tensor& input) override - { - return run_forward_(input, schema_); - } + [[nodiscard]] Result forward_named(const Tensor& input) override + { return run_forward_(input, schema_); } - [[nodiscard]] Result forward_named( - const Tensor& input, const TensorSchema& schema) override - { - return run_forward_(input, schema); - } + [[nodiscard]] Result forward_named(const Tensor& input, + const TensorSchema& schema) override + { return run_forward_(input, schema); } - [[nodiscard]] Result backward_named( - const GradientBundle& gradients) override - { - return backward_named(gradients, schema_); - } + [[nodiscard]] Result backward_named(const GradientBundle& gradients) override + { return backward_named(gradients, schema_); } - [[nodiscard]] Result backward_named( - const GradientBundle& gradients, - const TensorSchema& requested_schema) override + [[nodiscard]] Result backward_named(const GradientBundle& gradients, + const TensorSchema& requested_schema) override { - if (!forward_ready_) - NN_FAIL("OutputBundleLayer::backward_named: 没有匹配的 forward"); - if (requested_schema != schema_ || - gradients.schema() != schema_) + if (!forward_ready_) NN_FAIL("OutputBundleLayer::backward_named: 没有匹配的 forward"); + if (requested_schema != schema_ || gradients.schema() != schema_) NN_FAIL("OutputBundleLayer::backward_named: schema 不匹配"); // backward 只能消费一次；错误路径也不保留半配对状态。 @@ -379,15 +325,14 @@ class OutputBundleLayer final : public Layer if (!produced_[i]) { if (port == nullptr && grad != nullptr) - NN_FAIL("OutputBundleLayer::backward_named: 未激活端口 '" + - field.name + "' 收到了梯度"); + NN_FAIL("OutputBundleLayer::backward_named: 未激活端口 '" + field.name + + "' 收到了梯度"); continue; } if (grad == nullptr) - NN_FAIL("OutputBundleLayer::backward_named: 激活端口 '" + - field.name + "' 缺少梯度"); - NN_ASSERT_ALWAYS(port != nullptr, - "OutputBundleLayer backward lost active port"); + NN_FAIL("OutputBundleLayer::backward_named: 激活端口 '" + field.name + + "' 缺少梯度"); + NN_ASSERT_ALWAYS(port != nullptr, "OutputBundleLayer backward lost active port"); NN_TRY(branch_grad, port->branch->backward(*grad)); if (!have_input_grad) { @@ -404,8 +349,7 @@ class OutputBundleLayer final : public Layer NN_TRY_CHECK(engine_ref().accumulate(input_grad, *branch_grad)); } } - if (!have_input_grad) - NN_FAIL("OutputBundleLayer::backward_named: 没有可反向的输出端口"); + if (!have_input_grad) NN_FAIL("OutputBundleLayer::backward_named: 没有可反向的输出端口"); return input_grad; } @@ -444,33 +388,39 @@ class OutputBundleLayer final : public Layer void set_training(bool training) override { - for (Port& p : ports_) p.branch->set_training(training); + for (Port& p : ports_) + p.branch->set_training(training); } void set_flush_interval(std::size_t interval) override { - for (Port& p : ports_) p.branch->set_flush_interval(interval); + for (Port& p : ports_) + p.branch->set_flush_interval(interval); } void set_checkpoint_every(std::size_t stride) override { - for (Port& p : ports_) p.branch->set_checkpoint_every(stride); + for (Port& p : ports_) + p.branch->set_checkpoint_every(stride); } void set_activation_offload(bool enabled) override { - for (Port& p : ports_) p.branch->set_activation_offload(enabled); + for (Port& p : ports_) + p.branch->set_activation_offload(enabled); } void set_doc_ids(std::span ids) override { - for (Port& p : ports_) p.branch->set_doc_ids(ids); + for (Port& p : ports_) + p.branch->set_doc_ids(ids); } [[nodiscard]] std::size_t offload_ram_bytes() override { std::size_t total = 0; - for (Port& p : ports_) total += p.branch->offload_ram_bytes(); + for (Port& p : ports_) + total += p.branch->offload_ram_bytes(); return total; } @@ -487,7 +437,8 @@ class OutputBundleLayer final : public Layer void clear_cache() override { - for (Port& p : ports_) p.branch->clear_cache(); + for (Port& p : ports_) + p.branch->clear_cache(); forward_ready_ = false; } @@ -499,8 +450,7 @@ class OutputBundleLayer final : public Layer for (const TensorFieldSpec& field : schema_.fields()) { if (field.required && find_port_(field.name) == nullptr) - NN_FAIL("OutputBundleLayer::init: 必需输出端口 '" + field.name + - "' 没有分支"); + NN_FAIL("OutputBundleLayer::init: 必需输出端口 '" + field.name + "' 没有分支"); } return {}; } @@ -509,7 +459,8 @@ class OutputBundleLayer final : public Layer { std::vector out; out.reserve(ports_.size()); - for (Port& p : ports_) out.push_back(p.branch.get()); + for (Port& p : ports_) + out.push_back(p.branch.get()); return out; } }; diff --git a/include/neuralnet.cpp/compute_layer_conv.hpp b/include/neuralnet.cpp/compute_layer_conv.hpp index e48eea2c..9af69f46 100644 --- a/include/neuralnet.cpp/compute_layer_conv.hpp +++ b/include/neuralnet.cpp/compute_layer_conv.hpp @@ -26,15 +26,15 @@ namespace nn // rearrange_3d + gather_rows/scatter_add_rows（配一份形状相关的置换索引） // 组合出这两个方向，置换索引与数据无关，由各层缓存一次即可。 // ══════════════════════════════════════════════════════════════════════════ -namespace conv_engine { +namespace conv_engine +{ // 构造置换索引 perm（长度 C*P，一列）：perm[c*P + p] = p*C + c // 说明：rearrange_3d(Z, C, P, B) 得到行序 (p*C + c)，按 perm 用 gather 重排即得 (c*P + p)。 -[[nodiscard]] inline Result make_layout_perm( - ComputeEngine& engine, std::size_t C, std::size_t P) +[[nodiscard]] inline Result make_layout_perm(ComputeEngine& engine, std::size_t C, + std::size_t P) { - if (C == 0 || P == 0) - NN_FAIL("make_layout_perm: C/P must be > 0"); + if (C == 0 || P == 0) NN_FAIL("make_layout_perm: C/P must be > 0"); std::vector perm(C * P); // 宿主桥（17 §3 D11）：不经 Matrix for (std::size_t c = 0; c < C; ++c) for (std::size_t p = 0; p < P; ++p) @@ -43,9 +43,9 @@ namespace conv_engine { } // (C, B*P) → (C*P, B)：cols 布局转 samples 布局（Conv/Pool 层输出） -[[nodiscard]] inline Result cols_to_samples( - ComputeEngine& engine, const Tensor& Z, - std::size_t C, std::size_t P, std::size_t B, const Tensor& perm) +[[nodiscard]] inline Result cols_to_samples(ComputeEngine& engine, const Tensor& Z, + std::size_t C, std::size_t P, std::size_t B, + const Tensor& perm) { auto w = engine.rearrange_3d(Z, C, P, B, /*inverse=*/false); // (P*C, B) NN_TRY_CHECK(w); @@ -53,15 +53,18 @@ namespace conv_engine { } // (C*P, B) → (C, B*P)：samples 布局转回 cols 布局（Conv/Pool 层反向入口） -[[nodiscard]] inline Result samples_to_cols( - ComputeEngine& engine, const Tensor& g, - std::size_t C, std::size_t P, std::size_t B, const Tensor& perm) +[[nodiscard]] inline Result samples_to_cols(ComputeEngine& engine, const Tensor& g, + std::size_t C, std::size_t P, std::size_t B, + const Tensor& perm) { Tensor x2 = engine.create_tensor(P * C, B); - if (!x2.valid()) - NN_FAIL("samples_to_cols: 张量分配失败"); - { NN_TRY(r, engine.zero(x2)); } - { NN_TRY(r, engine.scatter_add_rows(x2, perm, g)); } + if (!x2.valid()) NN_FAIL("samples_to_cols: 张量分配失败"); + { + NN_TRY(r, engine.zero(x2)); + } + { + NN_TRY(r, engine.scatter_add_rows(x2, perm, g)); + } return engine.rearrange_3d(x2, C, P, B, /*inverse=*/true); // (C, B*P) } @@ -115,12 +118,11 @@ class Conv2D final : public Layer public: [[nodiscard]] const char* layer_name() const noexcept override { return "Conv2D"; } - 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), - kernel_(kernel), stride_(stride != 0 ? stride : 1), padding_(padding), - in_h_(in_h), in_w_(in_w) + 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), kernel_(kernel), + stride_(stride != 0 ? stride : 1), padding_(padding), in_h_(in_h), in_w_(in_w) { // API 级守卫：kernel 过大时 (in + 2*pad - kernel) 无符号下溢 → // out_h_/out_w_ 巨值 → 分配 abort。记录无效，init() 时报错。 @@ -139,8 +141,7 @@ class Conv2D final : public Layer [[nodiscard]] Result init_impl(ComputeEngine& engine) override { - if (shape_invalid_) - NN_FAIL("Conv2D: kernel 过大 (kernel > in + 2*padding)"); + if (shape_invalid_) NN_FAIL("Conv2D: kernel 过大 (kernel > in + 2*padding)"); const std::size_t fan_in = in_channels_ * kernel_ * kernel_; // ── 声明式初始化（M2，17 §4.4）：He 风格均匀分布，层算 limit、引擎 @@ -154,8 +155,7 @@ class Conv2D final : public Layer if (!b_.valid()) NN_FAIL("Conv2D: 偏置初始化失败"); grad_w_ = engine.create_tensor(out_channels_, fan_in, p_.param, InitSpec::zero()); grad_b_ = engine.create_tensor(out_channels_, 1, p_.param, InitSpec::zero()); - if (!grad_w_.valid() || !grad_b_.valid()) - NN_FAIL("Conv2D: 梯度缓冲初始化失败"); + if (!grad_w_.valid() || !grad_b_.valid()) NN_FAIL("Conv2D: 梯度缓冲初始化失败"); return {}; } @@ -174,23 +174,22 @@ class Conv2D final : public Layer // ── forward: Z = W × im2col(x) + b，再转 samples 布局 ─────────────── // 全程引擎原语 + DSL：窗口展开在 im2col 原语内完成（CPU/GPU 各自实现）， // 无 to_matrix/from_matrix 往返、无 CPU 标量循环。 - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); - if (shape_invalid_) - NN_FAIL("Conv2D: kernel 过大 (kernel > in + 2*padding)"); + if (shape_invalid_) NN_FAIL("Conv2D: kernel 过大 (kernel > in + 2*padding)"); if (input.rows() != in_channels_ * in_h_ * in_w_) NN_FAIL("conv forward: input shape mismatch"); const std::size_t batch = input.cols(); - if (batch == 0) - NN_FAIL("conv forward: batch must be > 0"); - { NN_TRY(r, ensure_perm_(engine)); } + if (batch == 0) NN_FAIL("conv forward: batch must be > 0"); + { + NN_TRY(r, ensure_perm_(engine)); + } const std::size_t P = out_h_ * out_w_; // 1) 窗口展开（引擎原语）：x (C_in*H*W, B) → col (C_in*k*k, P*B) - auto col = engine.im2col(input, in_channels_, in_h_, in_w_, - kernel_, stride_, padding_, out_h_, out_w_); + auto col = engine.im2col(input, in_channels_, in_h_, in_w_, kernel_, stride_, padding_, + out_h_, out_w_); NN_TRY_CHECK(col); // checkpoint 模式不驻留：显式清空，避免 size 相同导致静默用陈旧 im2col col_cache_ = checkpoint_mode_ ? Tensor{} : *col; @@ -198,9 +197,8 @@ class Conv2D final : public Layer // 2) Z = W × col + b → (C_out, P*B)：matmul 段与行广播偏置**融合为单次 // dispatch**（不物化 matmul 中间结果；与 Linear::forward / // matmul_with_bias 同一结构）。 - auto Z = dsl::compute(engine, - dsl::matmul(w_, *col, false, false) + dsl::row_broadcast(b_), - out_channels_, batch * P, p_.compute); + auto Z = dsl::compute(engine, dsl::matmul(w_, *col, false, false) + dsl::row_broadcast(b_), + out_channels_, batch * P, p_.compute); NN_TRY_CHECK(Z); // 3) 布局转换 (C_out, P*B) → (C_out*P, B)（rearrange_3d + gather 置换） @@ -208,53 +206,50 @@ class Conv2D final : public Layer } // ── backward ─────────────────────────────────────────────────────── - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { ComputeEngine& engine = engine_ref(); - if (shape_invalid_) - NN_FAIL("Conv2D: kernel 过大 (kernel > in + 2*padding)"); + if (shape_invalid_) NN_FAIL("Conv2D: kernel 过大 (kernel > in + 2*padding)"); if (grad_output.rows() != out_channels_ * out_h_ * out_w_) NN_FAIL("conv backward: grad_output shape mismatch"); const std::size_t batch = grad_output.cols(); - if (batch == 0) - NN_FAIL("conv backward: batch must be > 0"); + if (batch == 0) NN_FAIL("conv backward: batch must be > 0"); const std::size_t kk = kernel_ * kernel_; const std::size_t P = out_h_ * out_w_; // 缓存前置校验：checkpoint 模式（尚未 forward_recompute）下 forward 不驻留 // col_cache_；backward 必须校验缓存命中且形状匹配，否则会静默算出错误梯度。 if (col_cache_.rows() != in_channels_ * kk || col_cache_.cols() != P * batch) - NN_FAIL("conv backward: im2col 缓存缺失或不匹配" "（checkpoint 模式需先 forward_recompute；batch 变化后需重新 forward）"); - { NN_TRY(r, ensure_perm_(engine)); } + NN_FAIL("conv backward: im2col 缓存缺失或不匹配" + "（checkpoint 模式需先 forward_recompute；batch 变化后需重新 forward）"); + { + NN_TRY(r, ensure_perm_(engine)); + } // 1) grad_output (C_out*P, B) → gZ (C_out, P*B) - auto gZ = conv_engine::samples_to_cols(engine, grad_output, - out_channels_, P, batch, perm_cache_); + auto gZ = + conv_engine::samples_to_cols(engine, grad_output, out_channels_, P, batch, perm_cache_); NN_TRY_CHECK(gZ); // 2) grad_W += gZ × col^T → (C_out, C_in*k*k)：matmul 段与累加**融合为单次 // dispatch**并原地写入 grad_w_（不物化 gw 中间结果） - auto r1 = dsl::compute_into(engine, - dsl::leaf(grad_w_) + dsl::matmul(*gZ, col_cache_, false, true), grad_w_); + auto r1 = dsl::compute_into( + engine, dsl::leaf(grad_w_) + dsl::matmul(*gZ, col_cache_, false, true), grad_w_); NN_TRY_CHECK(r1); // 3) grad_b += row_reduce_sum(gZ) → (C_out, 1)（对全部位置与样本求和） // 归约步用 dsl::compute_reduce、累加步用 dsl::compute_into（同 Linear） - auto gb = dsl::compute_reduce(engine, - dsl::row_reduce_sum(dsl::leaf(*gZ)), gZ->rows(), gZ->cols(), - p_.compute); + auto gb = dsl::compute_reduce(engine, dsl::row_reduce_sum(dsl::leaf(*gZ)), gZ->rows(), + gZ->cols(), p_.compute); NN_TRY_CHECK(gb); - auto r2 = dsl::compute_into(engine, - dsl::leaf(grad_b_) + dsl::leaf(*gb), grad_b_); + auto r2 = dsl::compute_into(engine, dsl::leaf(grad_b_) + dsl::leaf(*gb), grad_b_); NN_TRY_CHECK(r2); // 4) gcol = W^T × gZ → (C_in*k*k, P*B) ；5) grad_x = col2im(gcol) // 纯 matmul 段由 dsl::compute 直写（scan 的 Conv2D backward dry-run 自动登记） - auto gcol = dsl::compute(engine, - dsl::matmul(w_, *gZ, true, false), - w_.cols(), gZ->cols(), p_.compute); + auto gcol = dsl::compute(engine, dsl::matmul(w_, *gZ, true, false), w_.cols(), gZ->cols(), + p_.compute); NN_TRY_CHECK(gcol); - return engine.col2im(*gcol, in_channels_, in_h_, in_w_, - kernel_, stride_, padding_, out_h_, out_w_); + return engine.col2im(*gcol, in_channels_, in_h_, in_w_, kernel_, stride_, padding_, out_h_, + out_w_); } }; @@ -313,8 +308,7 @@ class MaxPool2D final : public Layer for (std::size_t c = 0; c < channels_; ++c) for (std::size_t i = 0; i < kk; ++i) expand[c * kk + i] = static_cast(c); - auto t = detail::upload_span(engine, channels_ * kk, 1, Precision::F32, - std::span(expand)); + auto t = detail::upload_span(engine, channels_ * kk, 1, Precision::F32, std::span(expand)); NN_TRY_CHECK(t); expand_cache_ = std::move(*t); return {}; @@ -323,10 +317,10 @@ class MaxPool2D final : public Layer public: [[nodiscard]] const char* layer_name() const noexcept override { return "MaxPool2D"; } - MaxPool2D(std::size_t channels, std::size_t in_h, std::size_t in_w, - std::size_t pool = 2, std::size_t stride = 0) - : channels_(channels), in_h_(in_h), in_w_(in_w), - pool_(pool), stride_(stride != 0 ? stride : (pool != 0 ? pool : 1)) + MaxPool2D(std::size_t channels, std::size_t in_h, std::size_t in_w, std::size_t pool = 2, + std::size_t stride = 0) + : channels_(channels), in_h_(in_h), in_w_(in_w), pool_(pool), + stride_(stride != 0 ? stride : (pool != 0 ? pool : 1)) { // API 级守卫：pool 过大（或 pool=0）时 (in - pool) 无符号下溢 → // out_h_/out_w_ 巨值 → 分配 abort。记录无效，forward/backward 时报错。 @@ -354,24 +348,22 @@ class MaxPool2D final : public Layer [[nodiscard]] bool recompute_supported() const override { return true; } // ── forward ───────────────────────────────────────────────────────── - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); - if (shape_invalid_) - NN_FAIL("MaxPool2D: pool 窗口大于输入尺寸"); + if (shape_invalid_) NN_FAIL("MaxPool2D: pool 窗口大于输入尺寸"); if (input.rows() != channels_ * in_h_ * in_w_) NN_FAIL("maxpool forward: input shape mismatch"); const std::size_t batch = input.cols(); - if (batch == 0) - NN_FAIL("maxpool forward: batch must be > 0"); - { NN_TRY(r, ensure_perm_(engine)); } + if (batch == 0) NN_FAIL("maxpool forward: batch must be > 0"); + { + NN_TRY(r, ensure_perm_(engine)); + } const std::size_t kk = pool_ * pool_; const std::size_t P = out_h_ * out_w_; // 1) 窗口展开（引擎原语）：x (C*H*W, B) → col (C*kk, P*B) - auto col = engine.im2col(input, channels_, in_h_, in_w_, - pool_, stride_, 0, out_h_, out_w_); + auto col = engine.im2col(input, channels_, in_h_, in_w_, pool_, stride_, 0, out_h_, out_w_); NN_TRY_CHECK(col); // 2) 分组归约求窗口 max：每 kk 行一组 → (C, P*B)。 @@ -382,9 +374,9 @@ class MaxPool2D final : public Layer // 与 RAPT 的 row_broadcast(x) + rparam(0) 同一处理：与一个**运行时 0** // 相加使表达式合法，语义不变（rparam 值不进 expr_spec_key）。 auto pooled = dsl::compute(engine, - dsl::grouped_reduce_max(*col, static_cast(kk)) - + dsl::rparam(Scalar{0}), - channels_, col->cols()); + dsl::grouped_reduce_max(*col, static_cast(kk)) + + dsl::rparam(Scalar{0}), + channels_, col->cols()); NN_TRY_CHECK(pooled); // 3) checkpoint 模式不驻留：显式清空（避免 size 相同静默用陈旧数据） @@ -404,17 +396,14 @@ class MaxPool2D final : public Layer } // ── backward: 按窗口 mask 散射（重叠窗在 col2im 内累加）────────────── - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { ComputeEngine& engine = engine_ref(); - if (shape_invalid_) - NN_FAIL("MaxPool2D: pool 窗口大于输入尺寸"); + if (shape_invalid_) NN_FAIL("MaxPool2D: pool 窗口大于输入尺寸"); if (grad_output.rows() != channels_ * out_h_ * out_w_) NN_FAIL("maxpool backward: grad_output shape mismatch"); const std::size_t batch = grad_output.cols(); - if (batch == 0) - NN_FAIL("maxpool backward: batch must be > 0"); + if (batch == 0) NN_FAIL("maxpool backward: batch must be > 0"); const std::size_t kk = pool_ * pool_; const std::size_t P = out_h_ * out_w_; // 缓存前置校验：checkpoint（尚未 forward_recompute）或 clear_cache 后缓存为空， @@ -422,52 +411,52 @@ class MaxPool2D final : public Layer if (col_cache_.rows() != channels_ * kk || col_cache_.cols() != P * batch || pooled_cache_.rows() != channels_ || pooled_cache_.cols() != P * batch || !perm_cache_.valid()) - NN_FAIL("maxpool backward: 展开/max 缓存缺失或不匹配" "（checkpoint 模式需先 forward_recompute；clear_cache 后需重新 forward）"); + NN_FAIL("maxpool backward: 展开/max 缓存缺失或不匹配" + "（checkpoint 模式需先 forward_recompute；clear_cache 后需重新 forward）"); // 1) grad_output (C*P, B) → g_pooled (C, P*B) - auto g_pooled = conv_engine::samples_to_cols(engine, grad_output, - channels_, P, batch, perm_cache_); + auto g_pooled = + conv_engine::samples_to_cols(engine, grad_output, channels_, P, batch, perm_cache_); NN_TRY_CHECK(g_pooled); // 2) 并列均分（全部为整张量原语 + DSL，无逐通道循环） // mx_exp = 组内广播(pooled)：gather_rows 把每个通道的 max 复制 kk 次 // 总梯度守恒（Σ share = 窗口梯度）；无并列时（cnt=1）全部梯度落在该 // 唯一最大值位置。 - { NN_TRY(r, ensure_expand_(engine)); } + { + NN_TRY(r, ensure_expand_(engine)); + } auto mx_exp = engine.gather_rows(pooled_cache_, expand_cache_); // (C*kk, P*B) NN_TRY_CHECK(mx_exp); // 并列个数 cnt = 分组求和([x == 窗口 max]) → (C, P*B) // （grouped_reduce_sum 视图，同 forward） - auto eq = dsl::compute(engine, - dsl::select(dsl::leaf(col_cache_) == dsl::leaf(*mx_exp), - Scalar{1}, Scalar{0}), + auto eq = dsl::compute( + engine, dsl::select(dsl::leaf(col_cache_) == dsl::leaf(*mx_exp), Scalar{1}, Scalar{0}), channels_ * kk, P * batch); // (C*kk, P*B) NN_TRY_CHECK(eq); // 同 forward：裸分组归约视图根 = 空指令表（GPU 校验拒绝）→ + rparam(0) auto cnt = dsl::compute(engine, - dsl::grouped_reduce_sum(*eq, static_cast(kk)) - + dsl::rparam(Scalar{0}), - channels_, P * batch); + dsl::grouped_reduce_sum(*eq, static_cast(kk)) + + dsl::rparam(Scalar{0}), + channels_, P * batch); NN_TRY_CHECK(cnt); // 每个并列元素分到的梯度 g/cnt → 再广播回窗口内 - auto gdiv = dsl::compute(engine, - dsl::leaf(*g_pooled) / dsl::leaf(*cnt), - channels_, P * batch); // (C, P*B) + auto gdiv = dsl::compute(engine, dsl::leaf(*g_pooled) / dsl::leaf(*cnt), channels_, + P * batch); // (C, P*B) NN_TRY_CHECK(gdiv); auto gdiv_exp = engine.gather_rows(*gdiv, expand_cache_); // (C*kk, P*B) NN_TRY_CHECK(gdiv_exp); auto gc = dsl::compute(engine, - dsl::select(dsl::leaf(col_cache_) == dsl::leaf(*mx_exp), - dsl::leaf(*gdiv_exp), Scalar{0}), - channels_ * kk, P * batch); // (C*kk, P*B) + dsl::select(dsl::leaf(col_cache_) == dsl::leaf(*mx_exp), + dsl::leaf(*gdiv_exp), Scalar{0}), + channels_ * kk, P * batch); // (C*kk, P*B) NN_TRY_CHECK(gc); // 3) 反向散射到输入像素（重叠窗口在 col2im 内累加） - return engine.col2im(*gc, channels_, in_h_, in_w_, - pool_, stride_, 0, out_h_, out_w_); + return engine.col2im(*gc, channels_, in_h_, in_w_, pool_, stride_, 0, out_h_, out_w_); } }; } // namespace nn diff --git a/include/neuralnet.cpp/compute_layer_embedding.hpp b/include/neuralnet.cpp/compute_layer_embedding.hpp index 2b1fc254..d3b5f4ea 100644 --- a/include/neuralnet.cpp/compute_layer_embedding.hpp +++ b/include/neuralnet.cpp/compute_layer_embedding.hpp @@ -59,36 +59,31 @@ class Embedding final : public Layer { if (grad_table_.valid()) return; ComputeEngine& engine = engine_ref(); - grad_table_ = engine.create_tensor(vocab_size_, d_model_, p_.param, - InitSpec::zero()); - NN_UNWRAP(grad_table_.valid(), - "Embedding: 梯度缓冲（惰性）初始化失败"); + grad_table_ = engine.create_tensor(vocab_size_, d_model_, p_.param, InitSpec::zero()); + NN_UNWRAP(grad_table_.valid(), "Embedding: 梯度缓冲（惰性）初始化失败"); } // 清零 padding 行（覆盖语义：insert_rows 直接写目标行） [[nodiscard]] Result clear_padding_row_(Tensor& t) { if (padding_idx_ < 0) return {}; - return engine_ref().insert_rows( - t, static_cast(padding_idx_), padding_zero_row_); + return engine_ref().insert_rows(t, static_cast(padding_idx_), + padding_zero_row_); } public: // padding_idx：<0 = 无；否则 ∈ [0, vocab) - Embedding(std::size_t vocab_size, std::size_t d_model, - std::ptrdiff_t padding_idx = -1) + Embedding(std::size_t vocab_size, std::size_t d_model, std::ptrdiff_t padding_idx = -1) : vocab_size_(vocab_size), d_model_(d_model), padding_idx_(padding_idx) { // 前置条件校验（铁律 #1）：构造后仍会使用这些成员 → 无条件终止口径 NN_ASSERT_ALWAYS(vocab_size > 0, "Embedding: vocab_size 必须 > 0"); NN_ASSERT_ALWAYS(d_model > 0, "Embedding: d_model 必须 > 0"); - NN_ASSERT_ALWAYS(padding_idx < 0 || - static_cast(padding_idx) < vocab_size, + NN_ASSERT_ALWAYS(padding_idx < 0 || static_cast(padding_idx) < vocab_size, "Embedding: padding_idx 越界（须 < vocab_size 或负数表示无）"); } - [[nodiscard]] const char* layer_name() const noexcept override - { return "Embedding"; } + [[nodiscard]] const char* layer_name() const noexcept override { return "Embedding"; } [[nodiscard]] std::size_t vocab_size() const noexcept { return vocab_size_; } [[nodiscard]] std::size_t d_model() const noexcept { return d_model_; } @@ -100,10 +95,8 @@ class Embedding final : public Layer table_ = engine.create_tensor(vocab_size_, d_model_, p_.param, InitSpec::normal(Scalar{0}, Scalar{1}, kInitSeed)); if (!table_.valid()) NN_FAIL("Embedding: 嵌入表初始化失败"); - padding_zero_row_ = - engine.create_tensor(1, d_model_, p_.param, InitSpec::zero()); - if (!padding_zero_row_.valid()) - NN_FAIL("Embedding: padding 零行缓冲初始化失败"); + padding_zero_row_ = engine.create_tensor(1, d_model_, p_.param, InitSpec::zero()); + if (!padding_zero_row_.valid()) NN_FAIL("Embedding: padding 零行缓冲初始化失败"); // padding 行强制清零（分布初始化会填值） return clear_padding_row_(table_); } @@ -154,8 +147,7 @@ class Embedding final : public Layer } std::vector f(ids.begin(), ids.end()); ComputeEngine& engine = engine_ref(); - NN_TRY(idx, upload_span(engine, 1, ids.size(), Precision::F32, - std::span(f))); + NN_TRY(idx, upload_span(engine, 1, ids.size(), Precision::F32, std::span(f))); return forward(*idx); } @@ -178,8 +170,8 @@ class Embedding final : public Layer // padding 行梯度清零（每步；torch 同语义） NN_TRY(cpr, clear_padding_row_(grad_table_)); // 索引梯度 = 零（形状镜像输入） - Tensor zero = engine.create_tensor(ids_cache_.rows(), ids_cache_.cols(), - Precision::F32, InitSpec::zero()); + Tensor zero = engine.create_tensor(ids_cache_.rows(), ids_cache_.cols(), Precision::F32, + InitSpec::zero()); if (!zero.valid()) NN_FAIL("Embedding backward: 零梯度张量分配失败"); return zero; } diff --git a/include/neuralnet.cpp/compute_layer_feedforward.hpp b/include/neuralnet.cpp/compute_layer_feedforward.hpp index 55566144..13a1c50d 100644 --- a/include/neuralnet.cpp/compute_layer_feedforward.hpp +++ b/include/neuralnet.cpp/compute_layer_feedforward.hpp @@ -16,7 +16,8 @@ #include "model_spec.hpp" #include "expr_dsl.hpp" -namespace nn { +namespace nn +{ class FeedForward final : public Layer { @@ -27,7 +28,7 @@ class FeedForward final : public Layer std::unique_ptr fc2_; // (d_ff → d_model) // 三值模式（构造期定型；决定 fc1_/fc2_ 的**具体类型**与 profile 分流） WeightQuant weight_quant_ = WeightQuant::None; - GeLU gelu_; + GeLU gelu_; SwiGLU swiglu_; // SwiGLU 激活（含 split/merge） ReLU2GLU relu2glu_; // 门控平方 ReLU 激活（BitNet 2B4T；同 split/merge 接线） bool use_swiglu_ = false; @@ -47,22 +48,20 @@ class FeedForward final : public Layer [[nodiscard]] const char* layer_name() const noexcept override { return "FeedForward"; } FeedForward(std::size_t d_model, std::size_t d_ff, - ActivationType activation = ActivationType::GeLU, - bool subln = false, + ActivationType activation = ActivationType::GeLU, bool subln = false, NormType subln_norm_type = NormType::LayerNorm, WeightQuant weight_quant = WeightQuant::None) - : fc1_(make_linear_layer(d_model, - (activation == ActivationType::SwiGLU || - activation == ActivationType::ReLU2) ? 2 * d_ff : d_ff, weight_quant)), - fc2_(make_linear_layer(d_ff, d_model, weight_quant)), - weight_quant_(weight_quant), - swiglu_(d_ff), - relu2glu_(d_ff), - use_swiglu_(activation == ActivationType::SwiGLU), + : fc1_(make_linear_layer( + d_model, + (activation == ActivationType::SwiGLU || activation == ActivationType::ReLU2) + ? 2 * d_ff + : d_ff, + weight_quant)), + fc2_(make_linear_layer(d_ff, d_model, weight_quant)), weight_quant_(weight_quant), + swiglu_(d_ff), relu2glu_(d_ff), use_swiglu_(activation == ActivationType::SwiGLU), use_relu2_(activation == ActivationType::ReLU2) { - if (subln) - sub_norm_ = make_norm_layer(d_ff, subln_norm_type); + if (subln) sub_norm_ = make_norm_layer(d_ff, subln_norm_type); } [[nodiscard]] Result init_impl(ComputeEngine& engine) override @@ -71,9 +70,15 @@ class FeedForward final : public Layer NN_TRY(r2, fc2_->init(engine)); // M6 段 C：激活子层也是 Layer（engine 由 init 绑定），必须一并 init—— // 否则 forward 里调 gelu_/swiglu_.forward() 会在 engine_ref() 处 fail-fast。 - { NN_TRY(r, gelu_.init(engine)); } - { NN_TRY(r, swiglu_.init(engine)); } - { NN_TRY(r, relu2glu_.init(engine)); } + { + NN_TRY(r, gelu_.init(engine)); + } + { + NN_TRY(r, swiglu_.init(engine)); + } + { + NN_TRY(r, relu2glu_.init(engine)); + } if (sub_norm_) { NN_TRY(r, sub_norm_->init(engine)); } return {}; } @@ -131,8 +136,8 @@ class FeedForward final : public Layer // 三值权重）。非三值模式逐位不变（split 退化为恒等）。 void set_precision_profile(const PrecisionProfile& profile) override { - const TernaryProfileSplit sp = split_ternary_profile( - layer_name(), profile, weight_quant_ == WeightQuant::T1_58); + const TernaryProfileSplit sp = + split_ternary_profile(layer_name(), profile, weight_quant_ == WeightQuant::T1_58); Layer::set_precision_profile(sp.self); fc1_->set_precision_profile(sp.linear); fc2_->set_precision_profile(sp.linear); @@ -155,11 +160,16 @@ class FeedForward final : public Layer std::vector activation_cache() override { std::vector r; - auto a = fc1_->activation_cache(); r.insert(r.end(), a.begin(), a.end()); - auto b = fc2_->activation_cache(); r.insert(r.end(), b.begin(), b.end()); - auto g = gelu_.activation_cache(); r.insert(r.end(), g.begin(), g.end()); - auto s = swiglu_.activation_cache(); r.insert(r.end(), s.begin(), s.end()); - auto r2 = relu2glu_.activation_cache(); r.insert(r.end(), r2.begin(), r2.end()); + auto a = fc1_->activation_cache(); + r.insert(r.end(), a.begin(), a.end()); + auto b = fc2_->activation_cache(); + r.insert(r.end(), b.begin(), b.end()); + auto g = gelu_.activation_cache(); + r.insert(r.end(), g.begin(), g.end()); + auto s = swiglu_.activation_cache(); + r.insert(r.end(), s.begin(), s.end()); + auto r2 = relu2glu_.activation_cache(); + r.insert(r.end(), r2.begin(), r2.end()); if (sub_norm_) { auto sn = sub_norm_->activation_cache(); @@ -168,8 +178,7 @@ class FeedForward final : public Layer return r; } - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { NN_TRY(h1, fc1_->forward(input)); Tensor h2_out; @@ -197,8 +206,7 @@ class FeedForward final : public Layer return fc2_->forward(h2_out); } - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { NN_TRY(b2, fc2_->backward(grad_output)); Tensor g = std::move(*b2); diff --git a/include/neuralnet.cpp/compute_layer_gpt.hpp b/include/neuralnet.cpp/compute_layer_gpt.hpp index 6ade6f8b..40512ee8 100644 --- a/include/neuralnet.cpp/compute_layer_gpt.hpp +++ b/include/neuralnet.cpp/compute_layer_gpt.hpp @@ -41,22 +41,18 @@ class GPTBlock final : public Layer ActivationOffloader offloader_; public: - 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, + 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, - PrecisionProfile precision = PrecisionProfile{}, - bool subln = false, - std::size_t n_head_kv = 0, - WeightQuant weight_quant = WeightQuant::None) - : self_attn_(d_model, num_heads, max_len, seq_len, pos_enc, subln, norm_type, - n_head_kv, weight_quant), + PrecisionProfile precision = PrecisionProfile{}, bool subln = false, + std::size_t n_head_kv = 0, WeightQuant weight_quant = WeightQuant::None) + : self_attn_(d_model, num_heads, max_len, seq_len, pos_enc, subln, norm_type, n_head_kv, + weight_quant), norm1_(make_norm_layer(d_model, norm_type)), ff_(d_model, d_ff, activation, subln, norm_type, weight_quant), - norm2_(make_norm_layer(d_model, norm_type)), - weight_quant_(weight_quant) + norm2_(make_norm_layer(d_model, norm_type)), weight_quant_(weight_quant) { // D7：将精度配置注入所有子层（§9.2）—— override 内已逐子层下传 set_precision_profile(precision); @@ -71,8 +67,8 @@ class GPTBlock final : public Layer // 那一份；attention/FF 是"线性层持有者"，拿三值那一份（它们各自再分流）。 void set_precision_profile(const PrecisionProfile& profile) override { - const TernaryProfileSplit sp = split_ternary_profile( - layer_name(), profile, weight_quant_ == WeightQuant::T1_58); + const TernaryProfileSplit sp = + split_ternary_profile(layer_name(), profile, weight_quant_ == WeightQuant::T1_58); Layer::set_precision_profile(sp.self); self_attn_.set_precision_profile(sp.linear); if (norm1_) norm1_->set_precision_profile(sp.self); @@ -91,30 +87,21 @@ class GPTBlock final : public Layer std::vector parameters() override { - return collect_refs(self_attn_.parameters(), - norm1_->parameters(), - ff_.parameters(), + return collect_refs(self_attn_.parameters(), norm1_->parameters(), ff_.parameters(), norm2_->parameters()); } // 子层结构（部署格式的完整遍历；顺序 = 上面 `parameters()` 的拼接顺序） [[nodiscard]] std::vector param_children() noexcept override - { - return {&self_attn_, norm1_.get(), &ff_, norm2_.get()}; - } + { return {&self_attn_, norm1_.get(), &ff_, norm2_.get()}; } // 文档感知：把每样本文档 id 转发给内部自注意力（用于块对角掩码） - void set_doc_ids(std::span ids) override - { - self_attn_.set_doc_ids(ids); - } + void set_doc_ids(std::span ids) override { self_attn_.set_doc_ids(ids); } std::vector param_gradients() override { - return collect_refs(self_attn_.param_gradients(), - norm1_->param_gradients(), - ff_.param_gradients(), - norm2_->param_gradients()); + return collect_refs(self_attn_.param_gradients(), norm1_->param_gradients(), + ff_.param_gradients(), norm2_->param_gradients()); } // 梯度检查点：把模式传播给内部注意力/归一化/FFN @@ -130,8 +117,7 @@ class GPTBlock final : public Layer // GPTBlock 可作为“重计算单元”：从保存的块输入重算 forward 重建缓存 [[nodiscard]] bool recompute_supported() const override { return true; } - [[nodiscard]] Result forward_recompute( - const Tensor& saved_input) override + [[nodiscard]] Result forward_recompute(const Tensor& saved_input) override { // 临时关闭本块及其子层的 checkpoint 模式，使 forward 重建缓存 set_checkpoint_mode(false); @@ -152,10 +138,14 @@ class GPTBlock final : public Layer std::vector activation_cache() override { std::vector r; - auto a = self_attn_.activation_cache(); r.insert(r.end(), a.begin(), a.end()); - auto n1 = norm1_->activation_cache(); r.insert(r.end(), n1.begin(), n1.end()); - auto f = ff_.activation_cache(); r.insert(r.end(), f.begin(), f.end()); - auto n2 = norm2_->activation_cache(); r.insert(r.end(), n2.begin(), n2.end()); + auto a = self_attn_.activation_cache(); + r.insert(r.end(), a.begin(), a.end()); + auto n1 = norm1_->activation_cache(); + r.insert(r.end(), n1.begin(), n1.end()); + auto f = ff_.activation_cache(); + r.insert(r.end(), f.begin(), f.end()); + auto n2 = norm2_->activation_cache(); + r.insert(r.end(), n2.begin(), n2.end()); if (residual2_cache_.valid()) r.emplace_back(residual2_cache_); return r; } @@ -167,50 +157,37 @@ class GPTBlock final : public Layer // 导出/导入：委托给通用 ActivationOffloader（见 compute_layer_base.hpp）。 // 参与 offload 的缓存集合 = 本块 activation_cache()。 [[nodiscard]] Result export_activations(ComputeEngine& engine) - { - return offloader_.export_activations(engine, activation_cache()); - } + { return offloader_.export_activations(engine, activation_cache()); } [[nodiscard]] Result import_activations(ComputeEngine& engine) - { - return offloader_.import_activations(engine); - } + { return offloader_.import_activations(engine); } // 实际 slab 字节数（诊断用；未创建时 0） [[nodiscard]] std::size_t offload_slab_bytes() const noexcept - { - return offloader_.slab_bytes(); - } + { return offloader_.slab_bytes(); } // 新一轮 forward 入口：上一轮导出未被 backward 消费（推理/评估）时作废 void discard_offloaded_activations() noexcept { offloader_.discard(); } - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); // 残差分支 1（Pre-Norm + 自注意力），与 TransformerEncoderLayer/RAPTBlock 同骨架 NN_TRY(r2, prenorm_residual_forward_(engine, input, *norm1_, self_attn_, p_.compute)); Tensor res2 = std::move(*r2); - if (!checkpoint_mode_) - residual2_cache_ = res2; + if (!checkpoint_mode_) residual2_cache_ = res2; // 残差分支 2（Pre-Norm + FFN） return prenorm_residual_forward_(engine, res2, *norm2_, ff_, p_.compute); } - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { ComputeEngine& engine = engine_ref(); // activation offload：从 host 恢复激活再反向（替代重计算） - if (offloader_.offloaded()) - { - NN_TRY(im, import_activations(engine)); - } + if (offloader_.offloaded()) { NN_TRY(im, import_activations(engine)); } // 两处残差分流 + 子层正反向（与 TransformerEncoderLayer / RAPTBlock 同一骨架） - return prenorm_residual_backward_(engine, grad_output, - *norm1_, self_attn_, *norm2_, ff_, + return prenorm_residual_backward_(engine, grad_output, *norm1_, self_attn_, *norm2_, ff_, p_.compute); } @@ -220,29 +197,23 @@ class GPTBlock final : public Layer // x = x + FFN(LN₂(x)) // 输入: x_new (d_model, 1) // 输出: (d_model, 1) - [[nodiscard]] Result forward_step( - ComputeEngine& engine, - const Tensor& x_new, - Tensor& k_cache, - Tensor& v_cache, - std::size_t cur_len) + [[nodiscard]] Result forward_step(ComputeEngine& engine, const Tensor& x_new, + Tensor& k_cache, Tensor& v_cache, std::size_t cur_len) { NN_TRY(n1, norm1_->forward(x_new)); NN_TRY(a, self_attn_.forward_step(engine, *n1, k_cache, v_cache, cur_len)); - auto r2 = dsl::compute(engine, - dsl::leaf(x_new) + dsl::leaf(*a), - x_new.rows(), x_new.cols(), p_.compute); + auto r2 = dsl::compute(engine, dsl::leaf(x_new) + dsl::leaf(*a), x_new.rows(), x_new.cols(), + p_.compute); NN_TRY_CHECK(r2); NN_TRY(n2, norm2_->forward(*r2)); NN_TRY(f, ff_.forward(*n2)); - return dsl::compute(engine, - dsl::leaf(*r2) + dsl::leaf(*f), - r2->rows(), r2->cols(), p_.compute); + return dsl::compute(engine, dsl::leaf(*r2) + dsl::leaf(*f), r2->rows(), r2->cols(), + p_.compute); } }; @@ -319,11 +290,9 @@ class GPTModel final : public Layer // 多份 f32 副本 → OOM）。 [[nodiscard]] Result head_forward_(ComputeEngine& engine, const Tensor& x) { - if (!tie_embeddings_) - return lm_head_->forward(x); - return dsl::compute(engine, - dsl::matmul(token_emb_, x, /*transA=*/false, /*transB=*/false), - vocab_size_, x.cols(), p_.stable); + if (!tie_embeddings_) return lm_head_->forward(x); + return dsl::compute(engine, dsl::matmul(token_emb_, x, /*transA=*/false, /*transB=*/false), + vocab_size_, x.cols(), p_.stable); } public: @@ -332,16 +301,14 @@ class GPTModel final : public Layer PosEncodingType pos_enc_type = PosEncodingType::Learned, ActivationType activation = ActivationType::GeLU, NormType norm_type = NormType::LayerNorm, - PrecisionProfile precision = PrecisionProfile{}, - bool subln = false, - std::size_t n_head_kv = 0, - bool tie_embeddings = false, + PrecisionProfile precision = PrecisionProfile{}, bool subln = false, + std::size_t n_head_kv = 0, bool tie_embeddings = false, WeightQuant weight_quant = WeightQuant::None) : vocab_size_(vocab_size), d_model_(d_model), seq_len_(seq_len), ln_f_(make_norm_layer(d_model, norm_type)), // 三值模式：`precision.param == T1_58` 与显式 weight_quant 等价（唯一判据） - lm_head_(make_linear_layer( - d_model, vocab_size, effective_weight_quant(precision, weight_quant))), + lm_head_(make_linear_layer(d_model, vocab_size, + effective_weight_quant(precision, weight_quant))), weight_quant_(effective_weight_quant(precision, weight_quant)), tie_embeddings_(tie_embeddings) { @@ -350,8 +317,8 @@ class GPTModel final : public Layer //（param=T1_58 只对 BitLinear 有意义）——构造器里两处"辅助对象/LM head" // 的注入必须用同一份口径，否则会把 T1_58 打到 norm / 位置编码表上 //（它们要建张量 → 引擎拒绝"非存储精度"）。 - const TernaryProfileSplit sp = split_ternary_profile( - layer_name(), precision, weight_quant_ == WeightQuant::T1_58); + const TernaryProfileSplit sp = + split_ternary_profile(layer_name(), precision, weight_quant_ == WeightQuant::T1_58); set_precision_profile(precision); // ── LM head（词表投影）：计算精度强制 = stable ──────────────────── // logits 是 (vocab, batch·seq) —— 全模型最大的张量，且被 loss 链**多次** @@ -377,9 +344,8 @@ class GPTModel final : public Layer 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, precision, - subln, n_head_kv, weight_quant_); + blocks_.emplace_back(d_model, num_heads, d_ff, seq_len, seq_len, pos_enc_type, + activation, norm_type, precision, subln, n_head_kv, weight_quant_); } // 位置编码器是辅助对象（非 Layer）→ profile 需单独下传： @@ -401,8 +367,8 @@ class GPTModel final : public Layer // P1.5 三值分流（§4.8.3）：本模型自身（token 嵌入按 param 建张量）与 // 各 norm/位置编码拿 param 归一化为 f32 的那一份；LM head 是线性层， // 三值模式下拿 param=T1_58；块是"线性层持有者"，拿同一份（块内再分流）。 - const TernaryProfileSplit sp = split_ternary_profile( - layer_name(), profile, weight_quant_ == WeightQuant::T1_58); + const TernaryProfileSplit sp = + split_ternary_profile(layer_name(), profile, weight_quant_ == WeightQuant::T1_58); Layer::set_precision_profile(sp.self); if (ln_f_) ln_f_->set_precision_profile(sp.self); if (!tie_embeddings_) @@ -411,7 +377,8 @@ class GPTModel final : public Layer head_prof.compute = sp.self.stable; lm_head_->set_precision_profile(head_prof); } - for (auto& b : blocks_) b.set_precision_profile(sp.self); + for (auto& b : blocks_) + b.set_precision_profile(sp.self); if (pos_encoder_) pos_encoder_->set_precision_profile(sp.self); } @@ -421,34 +388,25 @@ class GPTModel final : public Layer constexpr Scalar emb_init_std = 0.02; token_emb_ = engine.create_tensor(vocab_size_, d_model_, p_.param, InitSpec::normal(0, emb_init_std, kInitSeed)); - if (!token_emb_.valid()) - NN_FAIL("GPTModel: token_emb 初始化失败"); + if (!token_emb_.valid()) NN_FAIL("GPTModel: token_emb 初始化失败"); grad_token_emb_ = engine.create_tensor(vocab_size_, d_model_, p_.param, InitSpec::zero()); - if (!grad_token_emb_.valid()) - NN_FAIL("GPTModel: token_emb 梯度缓冲初始化失败"); + if (!grad_token_emb_.valid()) NN_FAIL("GPTModel: token_emb 梯度缓冲初始化失败"); // 初始化子层 - if (pos_encoder_) - { - NN_TRY(r, pos_encoder_->init(engine)); - } + if (pos_encoder_) { NN_TRY(r, pos_encoder_->init(engine)); } for (auto& block : blocks_) { NN_TRY(r, block.init(engine)); } - if (ln_f_) - { - NN_TRY(r, ln_f_->init(engine)); - } + if (ln_f_) { NN_TRY(r, ln_f_->init(engine)); } if (!tie_embeddings_) { NN_TRY(r, lm_head_->init(engine)); } return {}; } std::vector parameters() override { - auto refs = collect_refs(token_emb_, - pos_encoder_->parameters(), - collect_block_refs_(blocks_, &GPTBlock::parameters), - ln_f_->parameters()); + auto refs = + collect_refs(token_emb_, pos_encoder_->parameters(), + collect_block_refs_(blocks_, &GPTBlock::parameters), ln_f_->parameters()); if (!tie_embeddings_) { auto h = lm_head_->parameters(); @@ -464,17 +422,16 @@ class GPTModel final : public Layer { std::vector r; r.reserve(r.size() + blocks_.size() + 2); - for (auto& b : blocks_) r.push_back(static_cast(&b)); + for (auto& b : blocks_) + r.push_back(static_cast(&b)); if (ln_f_) r.push_back(static_cast(ln_f_.get())); - if (!tie_embeddings_ && lm_head_) - r.push_back(static_cast(lm_head_.get())); + if (!tie_embeddings_ && lm_head_) r.push_back(static_cast(lm_head_.get())); return r; } std::vector param_gradients() override { - auto refs = collect_refs(grad_token_emb_, - pos_encoder_->param_gradients(), + auto refs = collect_refs(grad_token_emb_, pos_encoder_->param_gradients(), collect_block_refs_(blocks_, &GPTBlock::param_gradients), ln_f_->param_gradients()); if (!tie_embeddings_) @@ -489,19 +446,23 @@ class GPTModel final : public Layer // 传入空 span 清除文档感知（退化为纯因果）。在 forward 前调用。 void set_doc_ids(std::span ids) override { - if (ids.empty()) { doc_ids_.clear(); return; } + if (ids.empty()) + { + doc_ids_.clear(); + return; + } doc_ids_.assign(ids.begin(), ids.end()); } - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); // 新一轮 forward：上一轮 offload 的激活若未被 backward 消费（推理/评估 // 路径），其 slab 数据已作废——入口统一丢弃，保证 export 状态机从 // Idle 开始；否则本轮 export 会被 fail-fast 拒绝 / backward 导回旧激活。 if (activation_offload_) - for (auto& b : blocks_) b.discard_offloaded_activations(); + for (auto& b : blocks_) + b.discard_offloaded_activations(); const std::size_t seq_len = input.rows(); batch_size_ = input.cols(); // ── 1. gather 所有 token 的 embedding（统一采用 batch-major 列序） ── @@ -578,8 +539,7 @@ class GPTModel final : public Layer return head_forward_(engine, x); } - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { ComputeEngine& engine = engine_ref(); const std::size_t seq_len = seq_len_; @@ -596,18 +556,19 @@ class GPTModel final : public Layer if (tie_embeddings_) { // dX = W_embᵀ·g - auto dx = dsl::compute(engine, - dsl::matmul(token_emb_, grad_output, /*transA=*/true, /*transB=*/false), + auto dx = dsl::compute( + engine, dsl::matmul(token_emb_, grad_output, /*transA=*/true, /*transB=*/false), d_model_, grad_output.cols(), p_.stable); NN_TRY_CHECK(dx); grad_x = std::move(*dx); // dW_emb += g·xᵀ —— 与第 5 步的 scatter_add_rows **累加到同一张 // grad_token_emb_**（token_emb_ 同时充当输入嵌入与输出权重）。 - auto dw = dsl::compute_into(engine, - dsl::leaf(grad_token_emb_) - + dsl::matmul(grad_output, head_input_cache_, /*transA=*/false, - /*transB=*/true), - grad_token_emb_); + auto dw = + dsl::compute_into(engine, + dsl::leaf(grad_token_emb_) + + dsl::matmul(grad_output, head_input_cache_, /*transA=*/false, + /*transB=*/true), + grad_token_emb_); NN_TRY_CHECK(dw); } else @@ -665,10 +626,9 @@ class GPTModel final : public Layer NN_TRY(sr, engine.scatter_add_rows(grad_token_emb_, stored_tokens_tensor_, *grad_T)); // grad_input: token IDs 无梯度，返回零张量（仅用于接口一致性） - Tensor grad_input = engine.create_tensor(seq_len, batch_size_, Precision::F32, - InitSpec::zero()); - if (!grad_input.valid()) - NN_FAIL("GPT token_emb backward: 梯度张量分配失败"); + Tensor grad_input = + engine.create_tensor(seq_len, batch_size_, Precision::F32, InitSpec::zero()); + if (!grad_input.valid()) NN_FAIL("GPT token_emb backward: 梯度张量分配失败"); return grad_input; } @@ -688,7 +648,8 @@ class GPTModel final : public Layer void set_activation_offload(bool enabled) override { activation_offload_ = enabled; - for (auto& b : blocks_) b.set_offload_enabled(enabled); + for (auto& b : blocks_) + b.set_offload_enabled(enabled); } [[nodiscard]] bool activation_offload() const noexcept { return activation_offload_; } @@ -706,14 +667,16 @@ class GPTModel final : public Layer void set_checkpoint_mode(bool enabled) override { Layer::set_checkpoint_mode(enabled); - for (auto& b : blocks_) b.set_checkpoint_mode(enabled); + for (auto& b : blocks_) + b.set_checkpoint_mode(enabled); ln_f_->set_checkpoint_mode(enabled); lm_head_->set_checkpoint_mode(enabled); } void clear_cache() override { - for (auto& b : blocks_) b.clear_cache(); + for (auto& b : blocks_) + b.clear_cache(); ln_f_->clear_cache(); lm_head_->clear_cache(); head_input_cache_ = Tensor{}; @@ -724,13 +687,9 @@ class GPTModel final : public Layer // 流程: token_emb[token] [+ pos_emb[pos]] → N × GPTBlock.forward_step // → ln_f → lm_head → (vocab_size, 1) // 每层的 KV cache 由调用方维护，forward_step 只负责写入和计算。 - [[nodiscard]] Result forward_step( - ComputeEngine& engine, - std::size_t token_id, - std::size_t pos, - std::vector& k_caches, - std::vector& v_caches, - std::size_t cur_len) + [[nodiscard]] Result forward_step(ComputeEngine& engine, std::size_t token_id, + std::size_t pos, std::vector& k_caches, + std::vector& v_caches, std::size_t cur_len) { // 1. token embedding 查表 → (1, d_model) → transpose → (d_model, 1) std::vector id_v(1); // 宿主桥（17 §3 D11） @@ -747,13 +706,11 @@ class GPTModel final : public Layer Tensor x = std::move(*x_new); for (std::size_t i = 0; i < blocks_.size(); ++i) { - auto r = blocks_[i].forward_step( - engine, x, k_caches[i], v_caches[i], cur_len); + auto r = blocks_[i].forward_step(engine, x, k_caches[i], v_caches[i], cur_len); NN_TRY_CHECK(r); x = std::move(*r); // 按 flush_interval 拆分提交（防 TDR） - if (flush_interval_ > 0 && - (i + 1) % flush_interval_ == 0 && i + 1 < blocks_.size()) + if (flush_interval_ > 0 && (i + 1) % flush_interval_ == 0 && i + 1 < blocks_.size()) { NN_TRY(fr, engine.flush_batch()); } @@ -807,8 +764,7 @@ class GPTModel final : public Layer // ── prefill: 截断到 seq_len_ 长度（滑动窗口初始） ────────────── std::size_t start_init = 0; - if (context.size() > seq_len_) - start_init = context.size() - seq_len_; + if (context.size() > seq_len_) start_init = context.size() - seq_len_; std::size_t cur_len = 0; Tensor last_logits_t; @@ -822,8 +778,7 @@ class GPTModel final : public Layer NN_TRY(br, engine.begin_batch()); for (std::size_t i = start; i < context.size(); ++i) { - auto r = forward_step(engine, context[i], cur_len, - k_caches, v_caches, cur_len); + auto r = forward_step(engine, context[i], cur_len, k_caches, v_caches, cur_len); if (!r) { (void)engine.end_batch(); // 出错也收尾，避免录制状态泄漏 @@ -845,8 +800,14 @@ class GPTModel final : public Layer // Sliding window: rebuild cache when full (keep last seq_len_-1 tokens) if (cur_len >= seq_len_) { - for (auto& kc : k_caches) { NN_TRY(r, engine.zero(kc)); } - for (auto& vc : v_caches) { NN_TRY(r, engine.zero(vc)); } + for (auto& kc : k_caches) + { + NN_TRY(r, engine.zero(kc)); + } + for (auto& vc : v_caches) + { + NN_TRY(r, engine.zero(vc)); + } cur_len = 0; const std::size_t keep = seq_len_ - 1; const std::size_t start_new = (context.size() > keep) ? (context.size() - keep) : 0; @@ -868,16 +829,15 @@ class GPTModel final : public Layer context.push_back(next_token); - if (step >= min_new_tokens && next_token == eos_token_id) - break; + if (step >= min_new_tokens && next_token == eos_token_id) break; generated.push_back(next_token); // Run forward_step for next_token to write KV cache and get new logits NN_TRY(br, engine.begin_batch()); - auto logits_res = forward_step(engine, next_token, cur_len, - k_caches, v_caches, cur_len); + auto logits_res = + forward_step(engine, next_token, cur_len, k_caches, v_caches, cur_len); NN_TRY_CHECK(logits_res); NN_TRY(er, engine.end_batch()); diff --git a/include/neuralnet.cpp/compute_layer_mlp.hpp b/include/neuralnet.cpp/compute_layer_mlp.hpp index d3ba3159..93acd9df 100644 --- a/include/neuralnet.cpp/compute_layer_mlp.hpp +++ b/include/neuralnet.cpp/compute_layer_mlp.hpp @@ -47,17 +47,17 @@ class Linear final : public Layer { if (grad_w_.valid()) return; ComputeEngine& engine = engine_ref(); - grad_w_ = engine.create_tensor(out_features_, in_features_, p_.param, - InitSpec::zero()); + grad_w_ = engine.create_tensor(out_features_, in_features_, p_.param, InitSpec::zero()); grad_b_ = engine.create_tensor(out_features_, 1, p_.param, InitSpec::zero()); // 终止式检查（本函数返回 void，不能用 NN_FAIL 的 return 语义） - NN_UNWRAP(grad_w_.valid() && grad_b_.valid(), - "Linear: 梯度缓冲（惰性）初始化失败"); + NN_UNWRAP(grad_w_.valid() && grad_b_.valid(), "Linear: 梯度缓冲（惰性）初始化失败"); } public: Linear(std::size_t in_features, std::size_t out_features) - : in_features_(in_features), out_features_(out_features) {} + : in_features_(in_features), out_features_(out_features) + { + } [[nodiscard]] const char* layer_name() const noexcept override { return "Linear"; } @@ -76,10 +76,7 @@ class Linear final : public Layer return {}; } - [[nodiscard]] std::vector parameters() override - { - return {w_, b_}; - } + [[nodiscard]] std::vector parameters() override { return {w_, b_}; } [[nodiscard]] std::vector param_gradients() override { @@ -87,10 +84,7 @@ class Linear final : public Layer return {grad_w_, grad_b_}; } - [[nodiscard]] bool grad_buffers_allocated() const noexcept override - { - return grad_w_.valid(); - } + [[nodiscard]] bool grad_buffers_allocated() const noexcept override { return grad_w_.valid(); } void clear_cache() override { input_cache_ = Tensor{}; } @@ -103,39 +97,32 @@ class Linear final : public Layer // ── forward: 一行代码，精度由 p_.compute 决定 ────────────────────── // 引擎内部处理 matmul + broadcast bias 的精度问题 - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); - if (input.rows() != w_.cols()) - NN_FAIL("linear forward: input shape mismatch"); - if (!checkpoint_mode_) - input_cache_ = input; + if (input.rows() != w_.cols()) NN_FAIL("linear forward: input shape mismatch"); + if (!checkpoint_mode_) input_cache_ = input; // 计算精度 = p_.compute。Layer 直写该表达式：matmul 段 + row_broadcast // bias 融合为单 kernel（与 engine.matmul_with_bias 的 DSL 融合结构一致， // 见 compute_cpu_engine.hpp）；f16 存储经基类边界 cast 入口（NVI） // （边界 cast 或 in-kernel f16 变体，内部按 f32 计算、输出按目标精度舍入）。 - return dsl::compute(engine, - dsl::matmul(w_, input, false, false) + dsl::row_broadcast(b_), - p_.compute); // 输出形状 (M,N) 由表达式推导（matmul 段 + 行广播） + return dsl::compute(engine, dsl::matmul(w_, input, false, false) + dsl::row_broadcast(b_), + p_.compute); // 输出形状 (M,N) 由表达式推导（matmul 段 + 行广播） } // ── backward: 同样简洁，精度由引擎处理 ──────────────────────────── - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { ComputeEngine& engine = engine_ref(); - if (grad_output.rows() != w_.rows()) - NN_FAIL("linear backward: grad_output shape mismatch"); + if (grad_output.rows() != w_.rows()) NN_FAIL("linear backward: grad_output shape mismatch"); ensure_grad_(); // 训练路径：首次反向才建梯度缓冲（F2） // 计算精度 = p_.compute（in-place 累加的目标精度 = grad_w_ 的存储精度， // compute_into 无需 P：§8.3 in-place 存储精度不可变） // grad_input = W^T × grad_output：纯 matmul 段（无尾链），由 dsl::compute // 直写（scan 的 Linear backward dry-run 自动登记该结构）。 - auto grad_input = dsl::compute(engine, - dsl::matmul(w_, grad_output, true, false), - p_.compute); // (K, N) 由表达式推导 + auto grad_input = dsl::compute(engine, dsl::matmul(w_, grad_output, true, false), + p_.compute); // (K, N) 由表达式推导 NN_TRY_CHECK(grad_input); nn_dbg_scan("lin.grad_out", engine, grad_output); nn_dbg_scan("lin.cache", engine, input_cache_); @@ -147,8 +134,8 @@ class Linear final : public Layer // 并原地写入 grad_w_（GPU 上 1 个融合 kernel：不物化 gw (out,in)， // 也不额外分配输出缓冲）。 // k（求和维度 = batch 大小）是形状参数，不进 key → 同一 shader 适配任意 batch。 - auto grad_w_acc = dsl::compute_into(engine, - dsl::leaf(grad_w_) + dsl::matmul(grad_output, input_cache_, false, true), + auto grad_w_acc = dsl::compute_into( + engine, dsl::leaf(grad_w_) + dsl::matmul(grad_output, input_cache_, false, true), grad_w_); NN_TRY_CHECK(grad_w_acc); nn_dbg_scan("lin.grad_w(post-accum)", engine, grad_w_); @@ -160,13 +147,11 @@ class Linear final : public Layer // 故保留"归约 + 累加"两步：累加步用 dsl::compute_into（原地 dst += src， // f16 src 由叶子精度签名处理），归约步用 dsl::compute_reduce（输出 // (out,1) 归约向量同形）。 - auto gb = dsl::compute_reduce(engine, - dsl::row_reduce_sum(dsl::leaf(grad_output)), - grad_output.rows(), grad_output.cols(), p_.compute); + auto gb = dsl::compute_reduce(engine, dsl::row_reduce_sum(dsl::leaf(grad_output)), + grad_output.rows(), grad_output.cols(), p_.compute); NN_TRY_CHECK(gb); nn_dbg_scan("lin.row_sum(grad_out)", engine, *gb); - auto r2 = dsl::compute_into(engine, - dsl::leaf(grad_b_) + dsl::leaf(*gb), grad_b_); + auto r2 = dsl::compute_into(engine, dsl::leaf(grad_b_) + dsl::leaf(*gb), grad_b_); NN_TRY_CHECK(r2); nn_dbg_scan("lin.grad_b(post-accum)", engine, grad_b_); @@ -196,30 +181,25 @@ class ReLU final : public Layer // ── forward: out = max(x, 0) ────────────────────────────────────────── // ReLU 算法由 Layer 表达为 Max 原语 + 标量 0 // Engine/Shader 只提供 Max 原语，不知道 "ReLU" 是什么 - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); - if (!checkpoint_mode_) - input_cache_ = input; - return dsl::compute(engine, - dsl::max(dsl::leaf(input), Scalar{0}), - p_.compute); // 输出网格 = 输入网格（由表达式推导） + if (!checkpoint_mode_) input_cache_ = input; + return dsl::compute(engine, dsl::max(dsl::leaf(input), Scalar{0}), + p_.compute); // 输出网格 = 输入网格（由表达式推导） } // ── backward: grad_x = (x > 0) ? grad_out : 0 ──────────────────────── // ReLU 反向算法由 Layer 表达为 Select + Gt 原语 // Engine/Shader 只提供 Select/Gt 原语，不知道 "ReLU backward" 是什么 - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { ComputeEngine& engine = engine_ref(); - if (input_cache_.rows() != grad_output.rows() || - input_cache_.cols() != grad_output.cols()) + if (input_cache_.rows() != grad_output.rows() || input_cache_.cols() != grad_output.cols()) NN_FAIL("relu backward: shape mismatch"); - return dsl::compute(engine, - dsl::select(dsl::leaf(input_cache_) > Scalar{0}, - dsl::leaf(grad_output), Scalar{0}), + return dsl::compute( + engine, + dsl::select(dsl::leaf(input_cache_) > Scalar{0}, dsl::leaf(grad_output), Scalar{0}), p_.compute); // 两操作数同形状（上方已校验），网格由表达式推导 } }; @@ -247,10 +227,7 @@ class GeLU final : public Layer GeLU() = default; - void clear_cache() override - { - input_cache_ = Tensor{}; - } + void clear_cache() override { input_cache_ = Tensor{}; } std::vector activation_cache() override { @@ -263,38 +240,33 @@ class GeLU final : public Layer // 单表达式 DSL 融合（M 融合）：sigmoid(βx) = 1/(1+exp(-βx))，全部折叠为 // 单个 GPU 融合 kernel（仅 input/output 落显存，无中间 Tensor）。表达式 // 文本只写在本 Layer；AOT 收集由 scan_exprs dry-run 本方法完成。 - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); - if (!checkpoint_mode_) - input_cache_ = input; + if (!checkpoint_mode_) input_cache_ = input; const Scalar beta = BETA; // out = x * sigmoid(βx) = x / (1 + exp(-β·x)) return dsl::compute(engine, - dsl::leaf(input) / (Scalar{1} + dsl::exp(-(dsl::leaf(input) * beta))), - input.rows(), input.cols(), p_.compute); + dsl::leaf(input) / (Scalar{1} + dsl::exp(-(dsl::leaf(input) * beta))), + input.rows(), input.cols(), p_.compute); } // ── backward: grad_x = grad_out * factor ───────────────────────────── // factor = s * (1 + βx * (1 - s)), s = sigmoid(βx) = 1/(1+exp(-βx)) // 单表达式 DSL 融合：sigmoid 用 input_cache_ 重算（不缓存，省显存）， // βx 子表达式由 DSL 的 CSE 复用，消除中间 Tensor。 - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { ComputeEngine& engine = engine_ref(); - if (input_cache_.rows() != grad_output.rows() || - input_cache_.cols() != grad_output.cols()) + if (input_cache_.rows() != grad_output.rows() || input_cache_.cols() != grad_output.cols()) NN_FAIL("gelu backward: shape mismatch"); const Scalar beta = BETA; // s = 1/(1+exp(-βx))；factor = s * (1 + βx*(1-s))；out = grad_out * factor auto bx = dsl::leaf(input_cache_) * beta; // βx - auto s = Scalar{1} / (Scalar{1} + dsl::exp(-bx)); // sigmoid(βx) + auto s = Scalar{1} / (Scalar{1} + dsl::exp(-bx)); // sigmoid(βx) auto factor = s * (Scalar{1} + bx * (Scalar{1} - s)); // 1+βx(1-s) - return dsl::compute(engine, - dsl::leaf(grad_output) * factor, - grad_output.rows(), grad_output.cols(), p_.compute); + return dsl::compute(engine, dsl::leaf(grad_output) * factor, grad_output.rows(), + grad_output.cols(), p_.compute); } }; @@ -320,7 +292,8 @@ class SwiGLU final : public Layer { private: std::size_t d_ff_ = 0; - Tensor input_cache_; // 前向输入 (2*d_ff, batch)：backward 据此重算 gate/up/s（全融合，不缓存中间张量） + Tensor + input_cache_; // 前向输入 (2*d_ff, batch)：backward 据此重算 gate/up/s（全融合，不缓存中间张量） public: [[nodiscard]] const char* layer_name() const noexcept override { return "SwiGLU"; } @@ -328,10 +301,7 @@ class SwiGLU final : public Layer SwiGLU() = default; explicit SwiGLU(std::size_t d_ff) : d_ff_(d_ff) {} - void clear_cache() override - { - input_cache_ = Tensor{}; - } + void clear_cache() override { input_cache_ = Tensor{}; } std::vector activation_cache() override { @@ -344,21 +314,18 @@ class SwiGLU final : public Layer // 单表达式 DSL 融合：gate/up 用 RowAccess 行视图读取同一 (2*d_ff, batch) // 输入（gate = row_access(in, 0, d_ff)，up = row_access(in, d_ff, d_ff)）， // 行视图零拷贝、不物化半张量（无 D2D 拷贝）。输出 (d_ff, batch)。 - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); const std::uint32_t dff = static_cast(d_ff_); - if (!checkpoint_mode_) - input_cache_ = input; + if (!checkpoint_mode_) input_cache_ = input; const std::size_t cols = input.cols(); // gate[r] = in[r % d_ff]（r backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { ComputeEngine& engine = engine_ref(); const std::uint32_t dff = static_cast(d_ff_); @@ -388,17 +354,17 @@ class SwiGLU final : public Layer NN_FAIL("swiglu backward: input_cache shape mismatch"); if (grad_output.rows() != d_ff_ || grad_output.cols() != cols) NN_FAIL("swiglu backward: grad_output shape mismatch"); - const auto go = dsl::row_access(grad_output, 0u, dff); // go[r % d_ff] - const auto gate = dsl::row_access(input_cache_, 0u, dff); // in[r % d_ff] - const auto up = dsl::row_access(input_cache_, dff, dff); // in[d_ff + r % d_ff] - const auto s = Scalar{1} / (Scalar{1} + dsl::exp(-gate)); // σ(gate) + const auto go = dsl::row_access(grad_output, 0u, dff); // go[r % d_ff] + const auto gate = dsl::row_access(input_cache_, 0u, dff); // in[r % d_ff] + const auto up = dsl::row_access(input_cache_, dff, dff); // in[d_ff + r % d_ff] + const auto s = Scalar{1} / (Scalar{1} + dsl::exp(-gate)); // σ(gate) // factor = s·(1 + gate·(1−s))；grad_gate = go·up·factor；grad_up = go·gate·s const auto factor = s * (Scalar{1} + gate * (Scalar{1} - s)); const auto gg = go * up * factor; const auto gu = go * gate * s; - return dsl::compute(engine, - dsl::select(dsl::row() < dsl::rparam(static_cast(d_ff_)), gg, gu), + return dsl::compute( + engine, dsl::select(dsl::row() < dsl::rparam(static_cast(d_ff_)), gg, gu), rows, cols, p_.compute); } }; @@ -419,7 +385,7 @@ class ReLU2GLU final : public Layer { private: std::size_t d_ff_ = 0; - Tensor input_cache_; // 前向输入 (2*d_ff, batch)：backward 据此重算 gate + Tensor input_cache_; // 前向输入 (2*d_ff, batch)：backward 据此重算 gate public: [[nodiscard]] const char* layer_name() const noexcept override { return "ReLU2GLU"; } @@ -427,10 +393,7 @@ class ReLU2GLU final : public Layer ReLU2GLU() = default; explicit ReLU2GLU(std::size_t d_ff) : d_ff_(d_ff) {} - void clear_cache() override - { - input_cache_ = Tensor{}; - } + void clear_cache() override { input_cache_ = Tensor{}; } std::vector activation_cache() override { @@ -439,22 +402,19 @@ class ReLU2GLU final : public Layer return r; } - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); const std::uint32_t dff = static_cast(d_ff_); - if (!checkpoint_mode_) - input_cache_ = input; + if (!checkpoint_mode_) input_cache_ = input; const std::size_t cols = input.cols(); - const auto gv = dsl::row_access(input, 0u, dff); // gate = in[r] - const auto uv = dsl::row_access(input, dff, dff); // up = in[d_ff + r] - const auto r = dsl::relu(gv); // max(g, 0) + const auto gv = dsl::row_access(input, 0u, dff); // gate = in[r] + const auto uv = dsl::row_access(input, dff, dff); // up = in[d_ff + r] + const auto r = dsl::relu(gv); // max(g, 0) return dsl::compute(engine, r * r * uv, d_ff_, cols, p_.compute); } - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { ComputeEngine& engine = engine_ref(); const std::uint32_t dff = static_cast(d_ff_); @@ -466,15 +426,15 @@ class ReLU2GLU final : public Layer if (grad_output.rows() != d_ff_ || grad_output.cols() != cols) NN_FAIL("relu2glu backward: grad_output shape mismatch"); - const auto go = dsl::row_access(grad_output, 0u, dff); + const auto go = dsl::row_access(grad_output, 0u, dff); const auto gate = dsl::row_access(input_cache_, 0u, dff); - const auto up = dsl::row_access(input_cache_, dff, dff); - const auto r = dsl::relu(gate); + const auto up = dsl::row_access(input_cache_, dff, dff); + const auto r = dsl::relu(gate); // d/dg relu²(g) = 2·relu(g)（g=0 处不可导，取次梯度 0） const auto gg = go * up * (Scalar{2} * r); const auto gu = go * r * r; - return dsl::compute(engine, - dsl::select(dsl::row() < dsl::rparam(static_cast(d_ff_)), gg, gu), + return dsl::compute( + engine, dsl::select(dsl::row() < dsl::rparam(static_cast(d_ff_)), gg, gu), rows, cols, p_.compute); } }; @@ -505,14 +465,14 @@ class LayerNorm final : public Layer Scalar epsilon_; // 可学习参数 - Tensor gamma_; // (normalized_shape, 1) - Tensor beta_; // (normalized_shape, 1) + Tensor gamma_; // (normalized_shape, 1) + Tensor beta_; // (normalized_shape, 1) Tensor grad_gamma_; Tensor grad_beta_; // backward 缓存 - Tensor normalized_cache_; // (features, batch) - Tensor std_cache_; // (1, batch) — std_inv + Tensor normalized_cache_; // (features, batch) + Tensor std_cache_; // (1, batch) — std_inv static constexpr Scalar EPSILON = 1e-5; @@ -520,31 +480,26 @@ class LayerNorm final : public Layer [[nodiscard]] const char* layer_name() const noexcept override { return "LayerNorm"; } explicit LayerNorm(std::size_t normalized_shape, Scalar epsilon = EPSILON) - : normalized_shape_(normalized_shape), epsilon_(epsilon) {} + : normalized_shape_(normalized_shape), epsilon_(epsilon) + { + } [[nodiscard]] Result init_impl(ComputeEngine& engine) override { // gamma 初始化为 1, beta 初始化为 0（M2 声明式：引擎填数） gamma_ = engine.create_tensor(normalized_shape_, 1, p_.param, InitSpec::constant(1)); beta_ = engine.create_tensor(normalized_shape_, 1, p_.param, InitSpec::zero()); - if (!gamma_.valid() || !beta_.valid()) - NN_FAIL("LayerNorm: 参数初始化失败"); + if (!gamma_.valid() || !beta_.valid()) NN_FAIL("LayerNorm: 参数初始化失败"); grad_gamma_ = engine.create_tensor(normalized_shape_, 1, p_.param, InitSpec::zero()); grad_beta_ = engine.create_tensor(normalized_shape_, 1, p_.param, InitSpec::zero()); - if (!grad_gamma_.valid() || !grad_beta_.valid()) - NN_FAIL("LayerNorm: 梯度缓冲初始化失败"); + if (!grad_gamma_.valid() || !grad_beta_.valid()) NN_FAIL("LayerNorm: 梯度缓冲初始化失败"); return {}; } - [[nodiscard]] std::vector parameters() override - { - return {gamma_, beta_}; - } + [[nodiscard]] std::vector parameters() override { return {gamma_, beta_}; } [[nodiscard]] std::vector param_gradients() override - { - return {grad_gamma_, grad_beta_}; - } + { return {grad_gamma_, grad_beta_}; } void clear_cache() override { @@ -571,64 +526,59 @@ class LayerNorm final : public Layer // 5. std_inv = rsqrt(var_raw*(1/F) + ε) → (1,B) 原语 // 6. normalized=diff * std_inv (col 广播) → (F,B) 融合逐元素 // 7. out = normalized*gamma + beta (row 广播) → (F,B) 融合逐元素 - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); - if (input.rows() != normalized_shape_) - NN_FAIL("layernorm forward: input shape mismatch"); + if (input.rows() != normalized_shape_) NN_FAIL("layernorm forward: input shape mismatch"); const Scalar inv_features = Scalar{1} / static_cast(normalized_shape_); const std::size_t F = normalized_shape_; const std::size_t B = input.cols(); // 1. mean_raw = col_reduce_sum(x) → (1,B) - auto mean_raw = dsl::compute_reduce(engine, - dsl::col_reduce_sum(dsl::leaf(input)), F, B, p_.stable); + auto mean_raw = + dsl::compute_reduce(engine, dsl::col_reduce_sum(dsl::leaf(input)), F, B, p_.stable); NN_TRY_CHECK(mean_raw); // 2. mean = mean_raw*(1/F) → (1,B) // （1/F 是形状相关标量，由 RParam 承载：**值不进 expr_spec_key**， // 融合表达式保持 F 无关 → 不同归一化维度共享同一 AOT 融合 shader） - auto mean = dsl::compute(engine, - dsl::leaf(*mean_raw) * dsl::rparam(inv_features), - mean_raw->rows(), mean_raw->cols(), p_.stable); + auto mean = dsl::compute(engine, dsl::leaf(*mean_raw) * dsl::rparam(inv_features), + mean_raw->rows(), mean_raw->cols(), p_.stable); NN_TRY_CHECK(mean); // 3. diff = x - mean (col 广播) → (F,B) - auto diff = dsl::compute(engine, - dsl::leaf(input) - dsl::col_broadcast(*mean), F, B, p_.stable); + auto diff = + dsl::compute(engine, dsl::leaf(input) - dsl::col_broadcast(*mean), F, B, p_.stable); NN_TRY_CHECK(diff); // 4. var_raw = col_reduce_sum(diff²) → (1,B) - auto var_raw = dsl::compute_reduce(engine, - dsl::col_reduce_sum(dsl::leaf(*diff) * dsl::leaf(*diff)), F, B, p_.stable); + auto var_raw = dsl::compute_reduce( + engine, dsl::col_reduce_sum(dsl::leaf(*diff) * dsl::leaf(*diff)), F, B, p_.stable); NN_TRY_CHECK(var_raw); // 5. std_inv = rsqrt(var_raw*(1/F) + ε) → (1,B) // （"乘 1/F → 加 ε → rsqrt"三步塌成单表达式、单次遍历；1/F、ε 由 // RParam 承载，表达式结构仍与 F / ε 取值无关） - auto std_inv = dsl::compute(engine, - dsl::rsqrt(dsl::leaf(*var_raw) * dsl::rparam(inv_features) - + dsl::rparam(epsilon_)), + auto std_inv = dsl::compute( + engine, + dsl::rsqrt(dsl::leaf(*var_raw) * dsl::rparam(inv_features) + dsl::rparam(epsilon_)), var_raw->rows(), var_raw->cols(), p_.stable); NN_TRY_CHECK(std_inv); Tensor std_inv_t = std::move(*std_inv); - if (!checkpoint_mode_) - std_cache_ = std_inv_t; + if (!checkpoint_mode_) std_cache_ = std_inv_t; // 6. normalized = diff * std_inv (col 广播) → (F,B) - auto normalized = dsl::compute(engine, - dsl::leaf(*diff) * dsl::col_broadcast(std_inv_t), F, B, p_.stable); + auto normalized = + dsl::compute(engine, dsl::leaf(*diff) * dsl::col_broadcast(std_inv_t), F, B, p_.stable); NN_TRY_CHECK(normalized); Tensor normalized_t = std::move(*normalized); - if (!checkpoint_mode_) - normalized_cache_ = normalized_t; + if (!checkpoint_mode_) normalized_cache_ = normalized_t; // 7. out = normalized*gamma + beta (row 广播) → (F,B) return dsl::compute(engine, - dsl::leaf(normalized_t) * dsl::row_broadcast(gamma_) - + dsl::row_broadcast(beta_), - F, B, p_.stable); + dsl::leaf(normalized_t) * dsl::row_broadcast(gamma_) + + dsl::row_broadcast(beta_), + F, B, p_.stable); } // ── backward ────────────────────────────────────────────────────────── @@ -640,8 +590,7 @@ class LayerNorm final : public Layer // grad_beta += row_sum(grad_out) // 单表达式融合（融合表达式 F 无关，1/F 在 (1,B) 上用原语施加）： // mean_g/mean_gn 为列归约向量输出；(F,B) 全尺寸中间量由融合 kernel 消解。 - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { ComputeEngine& engine = engine_ref(); const Scalar inv_features = Scalar{1} / static_cast(normalized_shape_); @@ -649,55 +598,50 @@ class LayerNorm final : public Layer const std::size_t B = grad_output.cols(); // 1. mean_g_raw = col_reduce_sum(gy) → (1,B) - auto mg_raw = dsl::compute_reduce(engine, - dsl::col_reduce_sum( - dsl::leaf(grad_output) * dsl::row_broadcast(gamma_)), - F, B, p_.stable); + auto mg_raw = dsl::compute_reduce( + engine, dsl::col_reduce_sum(dsl::leaf(grad_output) * dsl::row_broadcast(gamma_)), F, B, + p_.stable); NN_TRY_CHECK(mg_raw); - auto mean_g = dsl::compute(engine, - dsl::leaf(*mg_raw) * dsl::rparam(inv_features), - mg_raw->rows(), mg_raw->cols(), p_.stable); + auto mean_g = dsl::compute(engine, dsl::leaf(*mg_raw) * dsl::rparam(inv_features), + mg_raw->rows(), mg_raw->cols(), p_.stable); NN_TRY_CHECK(mean_g); // 2. mean_gn_raw = col_reduce_sum(gy ⊙ normalized) → (1,B) - auto mgn_raw = dsl::compute_reduce(engine, - dsl::col_reduce_sum( - dsl::leaf(grad_output) * dsl::row_broadcast(gamma_) - * dsl::leaf(normalized_cache_)), + auto mgn_raw = dsl::compute_reduce( + engine, + dsl::col_reduce_sum(dsl::leaf(grad_output) * dsl::row_broadcast(gamma_) * + dsl::leaf(normalized_cache_)), F, B, p_.stable); NN_TRY_CHECK(mgn_raw); - auto mean_gn = dsl::compute(engine, - dsl::leaf(*mgn_raw) * dsl::rparam(inv_features), - mgn_raw->rows(), mgn_raw->cols(), p_.stable); + auto mean_gn = dsl::compute(engine, dsl::leaf(*mgn_raw) * dsl::rparam(inv_features), + mgn_raw->rows(), mgn_raw->cols(), p_.stable); NN_TRY_CHECK(mean_gn); // 3. grad_x = (gy - mean_g - normalized*mean_gn) * std_inv → (F,B) auto grad_x = dsl::compute(engine, - (dsl::leaf(grad_output) * dsl::row_broadcast(gamma_) - - dsl::col_broadcast(*mean_g) - - dsl::leaf(normalized_cache_) * dsl::col_broadcast(*mean_gn)) - * dsl::col_broadcast(std_cache_), - F, B, p_.stable); + (dsl::leaf(grad_output) * dsl::row_broadcast(gamma_) - + dsl::col_broadcast(*mean_g) - + dsl::leaf(normalized_cache_) * dsl::col_broadcast(*mean_gn)) * + dsl::col_broadcast(std_cache_), + F, B, p_.stable); NN_TRY_CHECK(grad_x); // 4. grad_gamma += row_reduce_sum(gy ⊙ normalized) → (F,1) // （∂L/∂γ_f = Σ_b gy·n：γ 是 out 的线性因子，导数不含 γ） - auto gg = dsl::compute_reduce(engine, - dsl::row_reduce_sum( - dsl::leaf(grad_output) - * dsl::leaf(normalized_cache_)), - F, B, p_.stable); + auto gg = dsl::compute_reduce( + engine, dsl::row_reduce_sum(dsl::leaf(grad_output) * dsl::leaf(normalized_cache_)), F, + B, p_.stable); NN_TRY_CHECK(gg); - auto grad_gamma_acc = dsl::compute_into(engine, - dsl::leaf(grad_gamma_) + dsl::leaf(*gg), grad_gamma_); + auto grad_gamma_acc = + dsl::compute_into(engine, dsl::leaf(grad_gamma_) + dsl::leaf(*gg), grad_gamma_); NN_TRY_CHECK(grad_gamma_acc); // 5. grad_beta += row_reduce_sum(grad_out) → (F,1) - auto gb = dsl::compute_reduce(engine, - dsl::row_reduce_sum(dsl::leaf(grad_output)), F, B, p_.stable); + auto gb = dsl::compute_reduce(engine, dsl::row_reduce_sum(dsl::leaf(grad_output)), F, B, + p_.stable); NN_TRY_CHECK(gb); - auto grad_beta_acc = dsl::compute_into(engine, - dsl::leaf(grad_beta_) + dsl::leaf(*gb), grad_beta_); + auto grad_beta_acc = + dsl::compute_into(engine, dsl::leaf(grad_beta_) + dsl::leaf(*gb), grad_beta_); NN_TRY_CHECK(grad_beta_acc); return grad_x; @@ -727,12 +671,12 @@ class RMSNorm final : public Layer std::size_t normalized_shape_; Scalar epsilon_; - Tensor gamma_; // (normalized_shape, 1) + Tensor gamma_; // (normalized_shape, 1) Tensor grad_gamma_; // backward 缓存 - Tensor normed_cache_; // (features, batch) - Tensor rms_inv_cache_; // (1, batch) + Tensor normed_cache_; // (features, batch) + Tensor rms_inv_cache_; // (1, batch) static constexpr Scalar EPSILON = 1e-5; @@ -740,29 +684,23 @@ class RMSNorm final : public Layer [[nodiscard]] const char* layer_name() const noexcept override { return "RMSNorm"; } explicit RMSNorm(std::size_t normalized_shape, Scalar epsilon = EPSILON) - : normalized_shape_(normalized_shape), epsilon_(epsilon) {} + : normalized_shape_(normalized_shape), epsilon_(epsilon) + { + } [[nodiscard]] Result init_impl(ComputeEngine& engine) override { // gamma 初始化为 1（无 beta）——M2 声明式：引擎填数 gamma_ = engine.create_tensor(normalized_shape_, 1, p_.param, InitSpec::constant(1)); - if (!gamma_.valid()) - NN_FAIL("RMSNorm: 参数初始化失败"); + if (!gamma_.valid()) NN_FAIL("RMSNorm: 参数初始化失败"); grad_gamma_ = engine.create_tensor(normalized_shape_, 1, p_.param, InitSpec::zero()); - if (!grad_gamma_.valid()) - NN_FAIL("RMSNorm: 梯度缓冲初始化失败"); + if (!grad_gamma_.valid()) NN_FAIL("RMSNorm: 梯度缓冲初始化失败"); return {}; } - [[nodiscard]] std::vector parameters() override - { - return {gamma_}; - } + [[nodiscard]] std::vector parameters() override { return {gamma_}; } - [[nodiscard]] std::vector param_gradients() override - { - return {grad_gamma_}; - } + [[nodiscard]] std::vector param_gradients() override { return {grad_gamma_}; } void clear_cache() override { @@ -786,44 +724,40 @@ class RMSNorm final : public Layer // 2. rms_inv= rsqrt(s_raw*(1/F) + ε) → (1,B) 原语 // 3. normed = x * rms_inv (col 广播) → (F,B) 融合逐元素 // 4. out = normed * gamma (row 广播) → (F,B) 融合逐元素 - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); - if (input.rows() != normalized_shape_) - NN_FAIL("rmsnorm forward: input shape mismatch"); + if (input.rows() != normalized_shape_) NN_FAIL("rmsnorm forward: input shape mismatch"); const Scalar inv_features = Scalar{1} / static_cast(normalized_shape_); const std::size_t F = normalized_shape_; const std::size_t B = input.cols(); // 1. s_raw = col_reduce_sum(x*x) → (1,B) - auto s_raw = dsl::compute_reduce(engine, - dsl::col_reduce_sum(dsl::leaf(input) * dsl::leaf(input)), F, B, p_.stable); + auto s_raw = dsl::compute_reduce( + engine, dsl::col_reduce_sum(dsl::leaf(input) * dsl::leaf(input)), F, B, p_.stable); NN_TRY_CHECK(s_raw); // 2. rms_inv = rsqrt(s_raw*(1/F) + ε) → (1,B) // （"乘 1/F → 加 ε → rsqrt"三步塌成单表达式、单次遍历；1/F、ε 由 // RParam 承载：值不进 expr_spec_key → 表达式结构保持 F/ε 无关） - auto rms_inv = dsl::compute(engine, - dsl::rsqrt(dsl::leaf(*s_raw) * dsl::rparam(inv_features) - + dsl::rparam(epsilon_)), + auto rms_inv = dsl::compute( + engine, + dsl::rsqrt(dsl::leaf(*s_raw) * dsl::rparam(inv_features) + dsl::rparam(epsilon_)), s_raw->rows(), s_raw->cols(), p_.stable); NN_TRY_CHECK(rms_inv); Tensor rms_inv_t = std::move(*rms_inv); - if (!checkpoint_mode_) - rms_inv_cache_ = rms_inv_t; + if (!checkpoint_mode_) rms_inv_cache_ = rms_inv_t; // 3. normed = x * rms_inv (col 广播) → (F,B) - auto normed = dsl::compute(engine, - dsl::leaf(input) * dsl::col_broadcast(rms_inv_t), F, B, p_.stable); + auto normed = + dsl::compute(engine, dsl::leaf(input) * dsl::col_broadcast(rms_inv_t), F, B, p_.stable); NN_TRY_CHECK(normed); Tensor normed_t = std::move(*normed); - if (!checkpoint_mode_) - normed_cache_ = normed_t; + if (!checkpoint_mode_) normed_cache_ = normed_t; // 4. out = normed * gamma (row 广播) → (F,B) - return dsl::compute(engine, - dsl::leaf(normed_t) * dsl::row_broadcast(gamma_), F, B, p_.stable); + return dsl::compute(engine, dsl::leaf(normed_t) * dsl::row_broadcast(gamma_), F, B, + p_.stable); } // ── backward ────────────────────────────────────────────────────────── @@ -833,8 +767,7 @@ class RMSNorm final : public Layer // m = m_raw * (1/F) // grad_x = (gy - m*normed) * rms_inv → (F,B) 融合逐元素 // grad_gamma += row_reduce_sum(gy ⊙ normed) → (F,1) 归约向量输出 - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { ComputeEngine& engine = engine_ref(); const Scalar inv_features = Scalar{1} / static_cast(normalized_shape_); @@ -842,35 +775,32 @@ class RMSNorm final : public Layer const std::size_t B = grad_output.cols(); // 1. m_raw = col_reduce_sum(gy ⊙ normed) → (1,B) - auto m_raw = dsl::compute_reduce(engine, - dsl::col_reduce_sum( - dsl::leaf(grad_output) * dsl::row_broadcast(gamma_) - * dsl::leaf(normed_cache_)), + auto m_raw = dsl::compute_reduce( + engine, + dsl::col_reduce_sum(dsl::leaf(grad_output) * dsl::row_broadcast(gamma_) * + dsl::leaf(normed_cache_)), F, B, p_.stable); NN_TRY_CHECK(m_raw); - auto m = dsl::compute(engine, - dsl::leaf(*m_raw) * dsl::rparam(inv_features), - m_raw->rows(), m_raw->cols(), p_.stable); + auto m = dsl::compute(engine, dsl::leaf(*m_raw) * dsl::rparam(inv_features), m_raw->rows(), + m_raw->cols(), p_.stable); NN_TRY_CHECK(m); // 2. grad_x = (gy - m*normed) * rms_inv → (F,B) auto grad_x = dsl::compute(engine, - (dsl::leaf(grad_output) * dsl::row_broadcast(gamma_) - - dsl::col_broadcast(*m) * dsl::leaf(normed_cache_)) - * dsl::col_broadcast(rms_inv_cache_), - F, B, p_.stable); + (dsl::leaf(grad_output) * dsl::row_broadcast(gamma_) - + dsl::col_broadcast(*m) * dsl::leaf(normed_cache_)) * + dsl::col_broadcast(rms_inv_cache_), + F, B, p_.stable); NN_TRY_CHECK(grad_x); // 3. grad_gamma += row_reduce_sum(gy ⊙ normed) → (F,1) // （∂L/∂γ_f = Σ_b gy·n：γ 是 out 的线性因子，导数不含 γ） - auto gg = dsl::compute_reduce(engine, - dsl::row_reduce_sum( - dsl::leaf(grad_output) - * dsl::leaf(normed_cache_)), - F, B, p_.stable); + auto gg = dsl::compute_reduce( + engine, dsl::row_reduce_sum(dsl::leaf(grad_output) * dsl::leaf(normed_cache_)), F, B, + p_.stable); NN_TRY_CHECK(gg); - auto grad_gamma_acc = dsl::compute_into(engine, - dsl::leaf(grad_gamma_) + dsl::leaf(*gg), grad_gamma_); + auto grad_gamma_acc = + dsl::compute_into(engine, dsl::leaf(grad_gamma_) + dsl::leaf(*gg), grad_gamma_); NN_TRY_CHECK(grad_gamma_acc); return grad_x; @@ -926,23 +856,23 @@ class BatchNorm final : public Layer private: std::size_t features_; Scalar epsilon_; - Scalar momentum_; // EMA 的新 batch 权重（默认 0.1） - bool training_ = true; // 训练/推理双态（set_training 切换） - bool stat_update_suppress_ = false; // forward_recompute 重算时抑制 EMA + Scalar momentum_; // EMA 的新 batch 权重（默认 0.1） + bool training_ = true; // 训练/推理双态（set_training 切换） + bool stat_update_suppress_ = false; // forward_recompute 重算时抑制 EMA // 可学习参数 - Tensor gamma_; // (features, 1) - Tensor beta_; // (features, 1) + Tensor gamma_; // (features, 1) + Tensor beta_; // (features, 1) Tensor grad_gamma_; Tensor grad_beta_; // 非可学习状态（extra_state → 序列化；初值 mean=0 / var=1） - Tensor running_mean_; // (features, 1) - Tensor running_var_; // (features, 1) + Tensor running_mean_; // (features, 1) + Tensor running_var_; // (features, 1) // backward 缓存 - Tensor normalized_cache_; // (features, batch) - Tensor std_cache_; // (features, 1) — 训练态 std_inv / 推理态 inv_std + Tensor normalized_cache_; // (features, batch) + Tensor std_cache_; // (features, 1) — 训练态 std_inv / 推理态 inv_std static constexpr Scalar EPSILON = 1e-5; static constexpr Scalar MOMENTUM = 0.1f; @@ -950,9 +880,10 @@ class BatchNorm final : public Layer public: [[nodiscard]] const char* layer_name() const noexcept override { return "BatchNorm"; } - explicit BatchNorm(std::size_t features, Scalar epsilon = EPSILON, - Scalar momentum = MOMENTUM) - : features_(features), epsilon_(epsilon), momentum_(momentum) {} + explicit BatchNorm(std::size_t features, Scalar epsilon = EPSILON, Scalar momentum = MOMENTUM) + : features_(features), epsilon_(epsilon), momentum_(momentum) + { + } [[nodiscard]] Result init_impl(ComputeEngine& engine) override { @@ -960,36 +891,28 @@ class BatchNorm final : public Layer // 推理态 = 恒等缩放，可直接训练早期评估）——M2 声明式：引擎填数。 gamma_ = engine.create_tensor(features_, 1, p_.param, InitSpec::constant(1)); beta_ = engine.create_tensor(features_, 1, p_.param, InitSpec::zero()); - if (!gamma_.valid() || !beta_.valid()) - NN_FAIL("BatchNorm: 参数初始化失败"); + if (!gamma_.valid() || !beta_.valid()) NN_FAIL("BatchNorm: 参数初始化失败"); grad_gamma_ = engine.create_tensor(features_, 1, p_.param, InitSpec::zero()); grad_beta_ = engine.create_tensor(features_, 1, p_.param, InitSpec::zero()); - if (!grad_gamma_.valid() || !grad_beta_.valid()) - NN_FAIL("BatchNorm: 梯度缓冲初始化失败"); + if (!grad_gamma_.valid() || !grad_beta_.valid()) NN_FAIL("BatchNorm: 梯度缓冲初始化失败"); running_mean_ = engine.create_tensor(features_, 1, p_.stable, InitSpec::zero()); - running_var_ = engine.create_tensor(features_, 1, p_.stable, - InitSpec::constant(1)); + running_var_ = engine.create_tensor(features_, 1, p_.stable, InitSpec::constant(1)); if (!running_mean_.valid() || !running_var_.valid()) NN_FAIL("BatchNorm: running 统计初始化失败"); return {}; } - [[nodiscard]] std::vector parameters() override - { - return {gamma_, beta_}; - } + [[nodiscard]] std::vector parameters() override { return {gamma_, beta_}; } [[nodiscard]] std::vector param_gradients() override - { - return {grad_gamma_, grad_beta_}; - } + { return {grad_gamma_, grad_beta_}; } // 非可学习状态：running 统计（不进 parameters → 优化器不更新） [[nodiscard]] std::vector extra_state() override { std::vector r; if (running_mean_.valid()) r.emplace_back(running_mean_); - if (running_var_.valid()) r.emplace_back(running_var_); + if (running_var_.valid()) r.emplace_back(running_var_); return r; } @@ -1023,12 +946,10 @@ class BatchNorm final : public Layer //（1/B、ε、momentum 都是 RParam：值不进 expr_spec_key） // 推理态：inv_std + 归一化 + 出参三步，无归约链（不同结构，scan 必须 // 单独覆盖——见 tools/scan_exprs.cpp 模型 pass 的推理态 run）。 - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); - if (input.rows() != features_) - NN_FAIL("batchnorm forward: input shape mismatch"); + if (input.rows() != features_) NN_FAIL("batchnorm forward: input shape mismatch"); const std::size_t F = features_; const std::size_t B = input.cols(); @@ -1036,103 +957,98 @@ class BatchNorm final : public Layer if (!training_) { // 1. inv_std = rsqrt(running_var + ε) → (F,1) - auto inv_std = dsl::compute(engine, - dsl::rsqrt(dsl::leaf(running_var_) + dsl::rparam(epsilon_)), - F, 1, p_.stable); + auto inv_std = + dsl::compute(engine, dsl::rsqrt(dsl::leaf(running_var_) + dsl::rparam(epsilon_)), F, + 1, p_.stable); NN_TRY_CHECK(inv_std); Tensor inv_std_t = std::move(*inv_std); - if (!checkpoint_mode_) - std_cache_ = inv_std_t; + if (!checkpoint_mode_) std_cache_ = inv_std_t; // 2. normalized = (x - running_mean) * inv_std → (F,B) auto normalized = dsl::compute(engine, - (dsl::leaf(input) - dsl::row_broadcast(running_mean_)) - * dsl::row_broadcast(inv_std_t), F, B, p_.stable); + (dsl::leaf(input) - dsl::row_broadcast(running_mean_)) * + dsl::row_broadcast(inv_std_t), + F, B, p_.stable); NN_TRY_CHECK(normalized); Tensor normalized_t = std::move(*normalized); - if (!checkpoint_mode_) - normalized_cache_ = normalized_t; + if (!checkpoint_mode_) normalized_cache_ = normalized_t; // 3. out = normalized*γ + β（与训练态同构，共享融合结构） return dsl::compute(engine, - dsl::leaf(normalized_t) * dsl::row_broadcast(gamma_) - + dsl::row_broadcast(beta_), - F, B, p_.stable); + dsl::leaf(normalized_t) * dsl::row_broadcast(gamma_) + + dsl::row_broadcast(beta_), + F, B, p_.stable); } // ── 训练态 ── const Scalar inv_batch = Scalar{1} / static_cast(B); // 1. mean_raw = row_reduce_sum(x) → (F,1) - auto mean_raw = dsl::compute_reduce(engine, - dsl::row_reduce_sum(dsl::leaf(input)), F, B, p_.stable); + auto mean_raw = + dsl::compute_reduce(engine, dsl::row_reduce_sum(dsl::leaf(input)), F, B, p_.stable); NN_TRY_CHECK(mean_raw); // 2. mean = mean_raw*(1/B) → (F,1) - auto mean = dsl::compute(engine, - dsl::leaf(*mean_raw) * dsl::rparam(inv_batch), - mean_raw->rows(), mean_raw->cols(), p_.stable); + auto mean = dsl::compute(engine, dsl::leaf(*mean_raw) * dsl::rparam(inv_batch), + mean_raw->rows(), mean_raw->cols(), p_.stable); NN_TRY_CHECK(mean); // 3. diff = x - mean (row 广播) → (F,B) - auto diff = dsl::compute(engine, - dsl::leaf(input) - dsl::row_broadcast(*mean), F, B, p_.stable); + auto diff = + dsl::compute(engine, dsl::leaf(input) - dsl::row_broadcast(*mean), F, B, p_.stable); NN_TRY_CHECK(diff); // 4. var_raw = row_reduce_sum(diff²) → (F,1) - auto var_raw = dsl::compute_reduce(engine, - dsl::row_reduce_sum(dsl::leaf(*diff) * dsl::leaf(*diff)), - F, B, p_.stable); + auto var_raw = dsl::compute_reduce( + engine, dsl::row_reduce_sum(dsl::leaf(*diff) * dsl::leaf(*diff)), F, B, p_.stable); NN_TRY_CHECK(var_raw); // 5. std_inv = rsqrt(var_raw*(1/B) + ε) → (F,1) - auto std_inv = dsl::compute(engine, - dsl::rsqrt(dsl::leaf(*var_raw) * dsl::rparam(inv_batch) - + dsl::rparam(epsilon_)), + auto std_inv = dsl::compute( + engine, + dsl::rsqrt(dsl::leaf(*var_raw) * dsl::rparam(inv_batch) + dsl::rparam(epsilon_)), var_raw->rows(), var_raw->cols(), p_.stable); NN_TRY_CHECK(std_inv); Tensor std_inv_t = std::move(*std_inv); - if (!checkpoint_mode_) - std_cache_ = std_inv_t; + if (!checkpoint_mode_) std_cache_ = std_inv_t; // 6. normalized = diff * std_inv (row 广播) → (F,B) - auto normalized = dsl::compute(engine, - dsl::leaf(*diff) * dsl::row_broadcast(std_inv_t), F, B, p_.stable); + auto normalized = + dsl::compute(engine, dsl::leaf(*diff) * dsl::row_broadcast(std_inv_t), F, B, p_.stable); NN_TRY_CHECK(normalized); Tensor normalized_t = std::move(*normalized); - if (!checkpoint_mode_) - normalized_cache_ = normalized_t; + if (!checkpoint_mode_) normalized_cache_ = normalized_t; // 7. running 统计 EMA（原地 compute_into；推理态无此步， // forward_recompute 重算时抑制——否则检查点场景会双倍更新） if (!stat_update_suppress_) { auto ema_m = dsl::compute_into(engine, - dsl::leaf(running_mean_) * dsl::rparam(1 - momentum_) - + dsl::leaf(*mean) * dsl::rparam(momentum_), - running_mean_); + dsl::leaf(running_mean_) * dsl::rparam(1 - momentum_) + + dsl::leaf(*mean) * dsl::rparam(momentum_), + running_mean_); NN_TRY_CHECK(ema_m); // 有偏 var（1/B 口径）直接以 var_raw*(1/B·m) 进 EMA，省一次分配 - auto ema_v = dsl::compute_into(engine, - dsl::leaf(running_var_) * dsl::rparam(1 - momentum_) - + dsl::leaf(*var_raw) * dsl::rparam(inv_batch * momentum_), - running_var_); + auto ema_v = + dsl::compute_into(engine, + dsl::leaf(running_var_) * dsl::rparam(1 - momentum_) + + dsl::leaf(*var_raw) * dsl::rparam(inv_batch * momentum_), + running_var_); NN_TRY_CHECK(ema_v); } // 8. out = normalized*γ + β → (F,B) return dsl::compute(engine, - dsl::leaf(normalized_t) * dsl::row_broadcast(gamma_) - + dsl::row_broadcast(beta_), - F, B, p_.stable); + dsl::leaf(normalized_t) * dsl::row_broadcast(gamma_) + + dsl::row_broadcast(beta_), + F, B, p_.stable); } // ── backward ────────────────────────────────────────────────────────── // 训练态：与 LayerNorm backward 同形（归约方向翻转为 batch、1/F → 1/B）； // 推理态：前向对 x 逐元素（无跨 b 耦合）→ grad_x 无均值修正项。 // γ/β 归约项两种模式同构（normalized_cache_ 语义一致）。 - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { ComputeEngine& engine = engine_ref(); if (!normalized_cache_.valid() || !std_cache_.valid()) @@ -1145,9 +1061,9 @@ class BatchNorm final : public Layer if (!training_) { auto grad_x = dsl::compute(engine, - dsl::leaf(grad_output) * dsl::row_broadcast(gamma_) - * dsl::row_broadcast(std_cache_), - F, B, p_.stable); + dsl::leaf(grad_output) * dsl::row_broadcast(gamma_) * + dsl::row_broadcast(std_cache_), + F, B, p_.stable); NN_TRY_CHECK(grad_x); NN_TRY(grad_gamma_acc, accumulate_gamma_(engine, grad_output)); NN_TRY(grad_beta_acc, accumulate_beta_(engine, grad_output)); @@ -1156,35 +1072,32 @@ class BatchNorm final : public Layer // ── 训练态 ── // 1. mean_g_raw = row_reduce_sum(gy) → (F,1) - auto mg_raw = dsl::compute_reduce(engine, - dsl::row_reduce_sum( - dsl::leaf(grad_output) * dsl::row_broadcast(gamma_)), - F, B, p_.stable); + auto mg_raw = dsl::compute_reduce( + engine, dsl::row_reduce_sum(dsl::leaf(grad_output) * dsl::row_broadcast(gamma_)), F, B, + p_.stable); NN_TRY_CHECK(mg_raw); - auto mean_g = dsl::compute(engine, - dsl::leaf(*mg_raw) * dsl::rparam(inv_batch), - mg_raw->rows(), mg_raw->cols(), p_.stable); + auto mean_g = dsl::compute(engine, dsl::leaf(*mg_raw) * dsl::rparam(inv_batch), + mg_raw->rows(), mg_raw->cols(), p_.stable); NN_TRY_CHECK(mean_g); // 2. mean_gn_raw = row_reduce_sum(gy ⊙ normalized) → (F,1) - auto mgn_raw = dsl::compute_reduce(engine, - dsl::row_reduce_sum( - dsl::leaf(grad_output) * dsl::row_broadcast(gamma_) - * dsl::leaf(normalized_cache_)), + auto mgn_raw = dsl::compute_reduce( + engine, + dsl::row_reduce_sum(dsl::leaf(grad_output) * dsl::row_broadcast(gamma_) * + dsl::leaf(normalized_cache_)), F, B, p_.stable); NN_TRY_CHECK(mgn_raw); - auto mean_gn = dsl::compute(engine, - dsl::leaf(*mgn_raw) * dsl::rparam(inv_batch), - mgn_raw->rows(), mgn_raw->cols(), p_.stable); + auto mean_gn = dsl::compute(engine, dsl::leaf(*mgn_raw) * dsl::rparam(inv_batch), + mgn_raw->rows(), mgn_raw->cols(), p_.stable); NN_TRY_CHECK(mean_gn); // 3. grad_x = (gy - mean_g - normalized*mean_gn) * std_inv → (F,B) auto grad_x = dsl::compute(engine, - (dsl::leaf(grad_output) * dsl::row_broadcast(gamma_) - - dsl::row_broadcast(*mean_g) - - dsl::leaf(normalized_cache_) * dsl::row_broadcast(*mean_gn)) - * dsl::row_broadcast(std_cache_), - F, B, p_.stable); + (dsl::leaf(grad_output) * dsl::row_broadcast(gamma_) - + dsl::row_broadcast(*mean_g) - + dsl::leaf(normalized_cache_) * dsl::row_broadcast(*mean_gn)) * + dsl::row_broadcast(std_cache_), + F, B, p_.stable); NN_TRY_CHECK(grad_x); // 4. grad_gamma += row_reduce_sum(grad_out ⊙ normalized) → (F,1) @@ -1200,7 +1113,7 @@ class BatchNorm final : public Layer [[nodiscard]] Result forward_recompute(const Tensor& saved_input) override { const bool prev_ckpt = checkpoint_mode_; - const bool prev_sup = stat_update_suppress_; + const bool prev_sup = stat_update_suppress_; checkpoint_mode_ = false; stat_update_suppress_ = true; auto r = forward(saved_input); @@ -1211,42 +1124,34 @@ class BatchNorm final : public Layer private: // grad_gamma += row_reduce_sum(grad_out ⊙ normalized)（训练/推理态共用） - [[nodiscard]] Result accumulate_gamma_( - ComputeEngine& engine, const Tensor& grad_output) + [[nodiscard]] Result accumulate_gamma_(ComputeEngine& engine, const Tensor& grad_output) { const std::size_t F = features_; const std::size_t B = grad_output.cols(); - auto gg = dsl::compute_reduce(engine, - dsl::row_reduce_sum( - dsl::leaf(grad_output) * dsl::leaf(normalized_cache_)), - F, B, p_.stable); + auto gg = dsl::compute_reduce( + engine, dsl::row_reduce_sum(dsl::leaf(grad_output) * dsl::leaf(normalized_cache_)), F, + B, p_.stable); NN_TRY_CHECK(gg); - return dsl::compute_into(engine, - dsl::leaf(grad_gamma_) + dsl::leaf(*gg), grad_gamma_); + return dsl::compute_into(engine, dsl::leaf(grad_gamma_) + dsl::leaf(*gg), grad_gamma_); } // grad_beta += row_reduce_sum(grad_out)（训练/推理态共用） - [[nodiscard]] Result accumulate_beta_( - ComputeEngine& engine, const Tensor& grad_output) + [[nodiscard]] Result accumulate_beta_(ComputeEngine& engine, const Tensor& grad_output) { const std::size_t F = features_; const std::size_t B = grad_output.cols(); - auto gb = dsl::compute_reduce(engine, - dsl::row_reduce_sum(dsl::leaf(grad_output)), F, B, p_.stable); + auto gb = dsl::compute_reduce(engine, dsl::row_reduce_sum(dsl::leaf(grad_output)), F, B, + p_.stable); NN_TRY_CHECK(gb); - return dsl::compute_into(engine, - dsl::leaf(grad_beta_) + dsl::leaf(*gb), grad_beta_); + return dsl::compute_into(engine, dsl::leaf(grad_beta_) + dsl::leaf(*gb), grad_beta_); } }; // ── 归一化层工厂：按 NormType 创建 LayerNorm / RMSNorm / BatchNorm ──────── -[[nodiscard]] inline std::unique_ptr make_norm_layer( - std::size_t d_model, NormType norm_type) +[[nodiscard]] inline std::unique_ptr make_norm_layer(std::size_t d_model, NormType norm_type) { - if (norm_type == NormType::RMSNorm) - return std::make_unique(d_model); - if (norm_type == NormType::BatchNorm) - return std::make_unique(d_model); + if (norm_type == NormType::RMSNorm) return std::make_unique(d_model); + if (norm_type == NormType::BatchNorm) return std::make_unique(d_model); return std::make_unique(d_model); } @@ -1255,12 +1160,11 @@ class BatchNorm final : public Layer // 用**运行期开关**选线性层类型，成员类型统一为 `std::unique_ptr`。 // BitLinear 与 Linear 的参数形状/顺序完全一致（(out,in) 权重 + (out,1) 偏置）， // 因此序列化与优化器侧零改动。 -[[nodiscard]] inline std::unique_ptr make_linear_layer( - std::size_t in_features, std::size_t out_features, - WeightQuant quant = WeightQuant::None) +[[nodiscard]] inline std::unique_ptr make_linear_layer(std::size_t in_features, + std::size_t out_features, + WeightQuant quant = WeightQuant::None) { - if (quant == WeightQuant::T1_58) - return std::make_unique(in_features, out_features); + if (quant == WeightQuant::T1_58) return std::make_unique(in_features, out_features); return std::make_unique(in_features, out_features); } } // namespace nn diff --git a/include/neuralnet.cpp/compute_layer_rapt.hpp b/include/neuralnet.cpp/compute_layer_rapt.hpp index 67a18fca..384666d5 100644 --- a/include/neuralnet.cpp/compute_layer_rapt.hpp +++ b/include/neuralnet.cpp/compute_layer_rapt.hpp @@ -66,12 +66,12 @@ class ReLULinearAttention final : public Layer Tensor Kp_cache_; // (BH*d_k, seq) ReLU(RoPE(K)) Tensor V_re_cache_; // (BH*d_k, seq) V（不 ReLU、不 RoPE） std::size_t batch_cache_ = 0; - std::size_t seq_cache_ = 0; + std::size_t seq_cache_ = 0; // z-scan 缓存：V_ones / e_0 用于计算 z = Σk'（RLA-2 分母） Tensor V_ones_cache_; // (BH*dk, seq) 全1矩阵 Tensor e_0_cache_; // (BH*dk, seq) 每头首行=1、其余=0（单位向量） - std::size_t ones_BH_ = 0; + std::size_t ones_BH_ = 0; std::size_t ones_seq_ = 0; // RMSNorm 缓存（RLA-2：投影后 RMSNorm → RoPE → ReLU） @@ -86,8 +86,7 @@ class ReLULinearAttention final : public Layer // 由 doc_ids_ 构建文档边界向量（size = batch*seq）：boundary[b*seq+t]=1 // 表示位置 t 是文档起点（t==0 或与前一位置文档不同）。 - [[nodiscard]] std::vector build_boundary_( - std::size_t batch, std::size_t seq) const + [[nodiscard]] std::vector build_boundary_(std::size_t batch, std::size_t seq) const { std::vector boundary; if (!has_doc_ids_) return boundary; @@ -113,29 +112,26 @@ class ReLULinearAttention final : public Layer // 确保 V_ones / e_0 缓存与当前 BH·dk × seq 尺寸匹配。 // 返回 Result：分配/upload 失败（常见为显存不足）时传播真实错误， // 避免后续 scan/outer_col 对空张量报 "import: invalid tensor" 掩盖根因。 - [[nodiscard]] Result ensure_ones_( - ComputeEngine& engine, std::size_t BH, std::size_t dk, std::size_t seq) + [[nodiscard]] Result ensure_ones_(ComputeEngine& engine, std::size_t BH, std::size_t dk, + std::size_t seq) { // 两个缓存都有效且尺寸匹配才命中（避免半初始化：V_ones 有效但 e_0 缺失）。 - if (ones_BH_ == BH && ones_seq_ == seq && - V_ones_cache_.valid() && e_0_cache_.valid()) + if (ones_BH_ == BH && ones_seq_ == seq && V_ones_cache_.valid() && e_0_cache_.valid()) return {}; const std::size_t rows = BH * dk; // V_ones：全1（声明式常数初始化，M2 InitSpec） - V_ones_cache_ = engine.create_tensor(rows, seq, Precision::F32, - InitSpec::constant(Scalar{1})); - if (!V_ones_cache_.valid()) - NN_FAIL("ensure_ones_: V_ones 分配失败"); + V_ones_cache_ = + engine.create_tensor(rows, seq, Precision::F32, InitSpec::constant(Scalar{1})); + if (!V_ones_cache_.valid()) NN_FAIL("ensure_ones_: V_ones 分配失败"); // e_0：每头首行=1，其余=0（用于 suffix(scale·q) 的外积构造） std::vector e0(rows * seq, Scalar{0}); // 宿主桥（17 §3 D11） for (std::size_t bh = 0; bh < BH; ++bh) for (std::size_t t = 0; t < seq; ++t) e0[(bh * dk) * seq + t] = Scalar{1}; - auto e0t = detail::upload_span(engine, rows, seq, Precision::F32, - std::span(e0)); + auto e0t = detail::upload_span(engine, rows, seq, Precision::F32, std::span(e0)); NN_TRY_CHECK(e0t); e_0_cache_ = std::move(*e0t); - ones_BH_ = BH; + ones_BH_ = BH; ones_seq_ = seq; return {}; } @@ -143,9 +139,9 @@ class ReLULinearAttention final : public Layer // ── 逐头 RMSNorm 前向（RLA-2：投影后 RMSNorm → RoPE → ReLU） ──── // 对 (BH*dk, seq) 张量的每个 dk 块独立归一化：y = x / sqrt(mean(x²) + eps) // 无学习参数（固定增益=1），符合极简原则。 - [[nodiscard]] Result rms_norm_forward_( - ComputeEngine& engine, const Tensor& input, - std::size_t BH, std::size_t dk, std::size_t seq, Tensor* rms_inv_out) + [[nodiscard]] Result rms_norm_forward_(ComputeEngine& engine, const Tensor& input, + std::size_t BH, std::size_t dk, std::size_t seq, + Tensor* rms_inv_out) { const Scalar inv_dk = Scalar{1} / static_cast(dk); const Scalar eps = Scalar{1e-5}; @@ -187,11 +183,11 @@ class ReLULinearAttention final : public Layer // 必须分两步：归约融合 shader 尚不支持"归约后仍有逐元素后处理" // 的形态（GPU 侧会静默错值，见 GpuEngine::eval_expr_reduce 的显式 // 拒绝）；1/dk、eps 由 RParam 承载 → 值不进 expr_spec_key。 - auto s = dsl::compute_reduce(engine, - dsl::col_reduce_sum(dsl::leaf(*x) * dsl::leaf(*x)), dk, seq, p_.compute); + auto s = dsl::compute_reduce(engine, dsl::col_reduce_sum(dsl::leaf(*x) * dsl::leaf(*x)), + dk, seq, p_.compute); NN_TRY_CHECK(s); - auto ri = dsl::compute(engine, - dsl::rsqrt(dsl::leaf(*s) * dsl::rparam(inv_dk) + dsl::rparam(eps)), + auto ri = dsl::compute( + engine, dsl::rsqrt(dsl::leaf(*s) * dsl::rparam(inv_dk) + dsl::rparam(eps)), s->rows(), s->cols(), p_.compute); NN_TRY_CHECK(ri); if (rms_inv_out) @@ -206,8 +202,8 @@ class ReLULinearAttention final : public Layer NN_TRY(ins_ri, engine.insert_rows(rms_out, bh, *ri)); } } - auto n = dsl::compute(engine, - dsl::leaf(*x) * dsl::col_broadcast(*ri), dk, seq, p_.compute); + auto n = + dsl::compute(engine, dsl::leaf(*x) * dsl::col_broadcast(*ri), dk, seq, p_.compute); NN_TRY_CHECK(n); NN_TRY(ins, engine.insert_rows(output, bh * dk, *n)); } @@ -218,10 +214,9 @@ class ReLULinearAttention final : public Layer // ── 逐头 RMSNorm 反向 ─────────────────────────────────────────── // dL/dx = (1/rms) · (g - y · (g·y)/dk) // 其中 y = normed（缓存），rms = 1/rms_inv（缓存）。 - [[nodiscard]] Result rms_norm_backward_( - ComputeEngine& engine, const Tensor& grad, - const Tensor& normed, const Tensor& rms_inv, - std::size_t BH, std::size_t dk, std::size_t seq) + [[nodiscard]] Result rms_norm_backward_(ComputeEngine& engine, const Tensor& grad, + const Tensor& normed, const Tensor& rms_inv, + std::size_t BH, std::size_t dk, std::size_t seq) { const Scalar inv_dk = Scalar{1} / static_cast(dk); // RV-02：分配失败原因直接上抛（同前向：精度与内部 DSL 一致，P2-3） @@ -234,18 +229,17 @@ class ReLULinearAttention final : public Layer NN_TRY(ri, engine.slice_rows(rms_inv, bh, 1)); // m = (1/dk) · col_reduce_sum(gy * y) → (1, seq) // 归约与后处理分两步（原因同上：归约融合 shader 不支持归约后后处理） - auto m_raw = dsl::compute_reduce(engine, - dsl::col_reduce_sum(dsl::leaf(*gy) * dsl::leaf(*y)), dk, seq, p_.compute); + auto m_raw = dsl::compute_reduce( + engine, dsl::col_reduce_sum(dsl::leaf(*gy) * dsl::leaf(*y)), dk, seq, p_.compute); NN_TRY_CHECK(m_raw); - auto m = dsl::compute(engine, - dsl::leaf(*m_raw) * dsl::rparam(inv_dk), - m_raw->rows(), m_raw->cols(), p_.compute); + auto m = dsl::compute(engine, dsl::leaf(*m_raw) * dsl::rparam(inv_dk), m_raw->rows(), + m_raw->cols(), p_.compute); NN_TRY_CHECK(m); // grad_x = (gy − m·y) · rms_inv：三式合一，单次遍历、无中间张量 auto gx = dsl::compute(engine, - (dsl::leaf(*gy) - dsl::col_broadcast(*m) * dsl::leaf(*y)) - * dsl::col_broadcast(*ri), - dk, seq, p_.compute); + (dsl::leaf(*gy) - dsl::col_broadcast(*m) * dsl::leaf(*y)) * + dsl::col_broadcast(*ri), + dk, seq, p_.compute); NN_TRY_CHECK(gx); NN_TRY(ins, engine.insert_rows(output, bh * dk, *gx)); } @@ -253,15 +247,11 @@ class ReLULinearAttention final : public Layer } public: - ReLULinearAttention(std::size_t d_model, std::size_t num_heads, - std::size_t seq_len = 0, - bool causal = true, - PosEncodingType pos_enc = PosEncodingType::RoPE) - : d_model_(d_model), num_heads_(num_heads), - d_k_(d_model / num_heads), - seq_len_(seq_len), causal_(causal), - w_q_(d_model, d_model), w_k_(d_model, d_model), - w_v_(d_model, d_model), w_o_(d_model, d_model) + ReLULinearAttention(std::size_t d_model, std::size_t num_heads, std::size_t seq_len = 0, + bool causal = true, PosEncodingType pos_enc = PosEncodingType::RoPE) + : d_model_(d_model), num_heads_(num_heads), d_k_(d_model / num_heads), seq_len_(seq_len), + causal_(causal), w_q_(d_model, d_model), w_k_(d_model, d_model), w_v_(d_model, d_model), + w_o_(d_model, d_model) { // RV-01：整除不成立 → d_k_ 错值（除零/越界索引）在后续热路径静默发生 NN_ASSERT_ALWAYS(d_model % num_heads == 0, @@ -340,26 +330,30 @@ class ReLULinearAttention final : public Layer Q_rms_inv_cache_ = Tensor{}; K_rms_inv_cache_ = Tensor{}; batch_cache_ = 0; - seq_cache_ = 0; - w_q_.clear_cache(); w_k_.clear_cache(); - w_v_.clear_cache(); w_o_.clear_cache(); + seq_cache_ = 0; + w_q_.clear_cache(); + w_k_.clear_cache(); + w_v_.clear_cache(); + w_o_.clear_cache(); } // 文档感知：记录每位置文档 id（batch-major），forward/backward 据此 // 在文档边界重置运行态（每个 token 只聚合本文档内前缀）。 void set_doc_ids(std::span ids) override { - if (ids.empty()) { doc_ids_.clear(); has_doc_ids_ = false; return; } + if (ids.empty()) + { + doc_ids_.clear(); + has_doc_ids_ = false; + return; + } doc_ids_.assign(ids.begin(), ids.end()); has_doc_ids_ = true; } // 绝对位置偏移（滑动窗生成用）：转发给位置编码策略（RoPE 生效，其余 no-op）， // 使重计算式生成的位置从真实起点算起而非每次从 0 重置。 - void set_position_offset(std::size_t off) - { - pos_->set_position_offset(off); - } + void set_position_offset(std::size_t off) { pos_->set_position_offset(off); } std::vector activation_cache() override { @@ -373,10 +367,14 @@ class ReLULinearAttention final : public Layer if (K_normed_cache_.valid()) r.emplace_back(K_normed_cache_); if (Q_rms_inv_cache_.valid()) r.emplace_back(Q_rms_inv_cache_); if (K_rms_inv_cache_.valid()) r.emplace_back(K_rms_inv_cache_); - auto wq = w_q_.activation_cache(); r.insert(r.end(), wq.begin(), wq.end()); - auto wk = w_k_.activation_cache(); r.insert(r.end(), wk.begin(), wk.end()); - auto wv = w_v_.activation_cache(); r.insert(r.end(), wv.begin(), wv.end()); - auto wo = w_o_.activation_cache(); r.insert(r.end(), wo.begin(), wo.end()); + auto wq = w_q_.activation_cache(); + r.insert(r.end(), wq.begin(), wq.end()); + auto wk = w_k_.activation_cache(); + r.insert(r.end(), wk.begin(), wk.end()); + auto wv = w_v_.activation_cache(); + r.insert(r.end(), wv.begin(), wv.end()); + auto wo = w_o_.activation_cache(); + r.insert(r.end(), wo.begin(), wo.end()); return r; } @@ -388,12 +386,10 @@ class ReLULinearAttention final : public Layer // // 实现：两次 scan_prefix_outer（主扫描 + z-scan），z-scan 用 V=ones 获得 z。 // ══════════════════════════════════════════════════════════════════════ - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); - if (input.rows() != d_model_) - NN_FAIL("ReLULinearAttention forward: input shape mismatch"); + if (input.rows() != d_model_) NN_FAIL("ReLULinearAttention forward: input shape mismatch"); const std::size_t total = input.cols(); const std::size_t seq = (seq_len_ > 0) ? seq_len_ : total; const std::size_t batch = (seq_len_ > 0) ? (total / seq_len_) : 1; @@ -445,11 +441,11 @@ class ReLULinearAttention final : public Layer NN_TRY(qr, pos_->apply_qk(engine, Q, seq, false)); NN_TRY(kr, pos_->apply_qk(engine, K, seq, false)); } - auto Qp = dsl::compute(engine, - dsl::max(dsl::leaf(Q), Scalar{0}), Q.rows(), Q.cols(), p_.compute); + auto Qp = + dsl::compute(engine, dsl::max(dsl::leaf(Q), Scalar{0}), Q.rows(), Q.cols(), p_.compute); NN_TRY_CHECK(Qp); - auto Kp = dsl::compute(engine, - dsl::max(dsl::leaf(K), Scalar{0}), K.rows(), K.cols(), p_.compute); + auto Kp = + dsl::compute(engine, dsl::max(dsl::leaf(K), Scalar{0}), K.rows(), K.cols(), p_.compute); NN_TRY_CHECK(Kp); // O(L·d_k²) 运行态前缀和扫描（消除 O(L²) 得分矩阵物化）。 @@ -463,8 +459,7 @@ class ReLULinearAttention final : public Layer std::vector bm(batch * seq, Scalar{0}); // 宿主桥（17 §3 D11） for (std::size_t i = 0; i < boundary.size(); ++i) bm[i] = static_cast(boundary[i]); - auto bt = detail::upload_span(engine, 1, batch * seq, Precision::F32, - std::span(bm)); + auto bt = detail::upload_span(engine, 1, batch * seq, Precision::F32, std::span(bm)); NN_TRY_CHECK(bt); boundary_t = std::move(*bt); has_bnd = true; @@ -478,9 +473,8 @@ class ReLULinearAttention final : public Layer const Tensor& dummy = *dummy_r; // 主扫描：读出 [0) B·P = num = B·q'（RLA-2 分子） - auto Sc = engine.scan_prefix_outer(*Kp, V, *Qp, V, dummy, dummy, false, - d_k_, num_heads_, causal_, - boundary_t, has_bnd); + auto Sc = engine.scan_prefix_outer(*Kp, V, *Qp, V, dummy, dummy, false, d_k_, num_heads_, + causal_, boundary_t, has_bnd); NN_TRY_CHECK(Sc); NN_TRY(BP_r, engine.slice_rows(*Sc, 0, BHdk)); @@ -488,17 +482,14 @@ class ReLULinearAttention final : public Layer // [0) B·P = q'·z（标量，头内逐行重复） // [2) B^T·R = dk·z（向量，backward 用） NN_TRY(ones_r, ensure_ones_(engine, BH, d_k_, seq)); - auto Z_sc = engine.scan_prefix_outer(*Kp, V_ones_cache_, *Qp, V_ones_cache_, - dummy, dummy, false, - d_k_, num_heads_, causal_, - boundary_t, has_bnd); + auto Z_sc = engine.scan_prefix_outer(*Kp, V_ones_cache_, *Qp, V_ones_cache_, dummy, dummy, + false, d_k_, num_heads_, causal_, boundary_t, has_bnd); NN_TRY_CHECK(Z_sc); NN_TRY(u_r, engine.slice_rows(*Z_sc, 0, BHdk)); // den = q'·z + ε；out = num / den（逐元素链融合为单 kernel：num/(u+ε)） - auto div_r = dsl::compute(engine, - dsl::leaf(*BP_r) / (dsl::leaf(*u_r) + Scalar{1e-4}), - (*BP_r).rows(), (*BP_r).cols(), p_.compute); + auto div_r = dsl::compute(engine, dsl::leaf(*BP_r) / (dsl::leaf(*u_r) + Scalar{1e-4}), + (*BP_r).rows(), (*BP_r).cols(), p_.compute); NN_TRY_CHECK(div_r); Tensor out_t = std::move(*div_r); @@ -521,7 +512,7 @@ class ReLULinearAttention final : public Layer V_re_cache_ = std::move(V); } batch_cache_ = batch; - seq_cache_ = seq; + seq_cache_ = seq; return w_o_.forward(concat); } @@ -542,19 +533,18 @@ class ReLULinearAttention final : public Layer // dB = outer(gnum, q), S_B = suffix(dB) // gV = S_B·k, gK = S_B^T·v + suffix(scale·q) // ══════════════════════════════════════════════════════════════════════ - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { ComputeEngine& engine = engine_ref(); // 缓存前置校验：checkpoint 模式（尚未 forward_recompute）或 offload // （尚未 import）下缓存为空，backward 必须立刻返回明确错误——否则空 // 张量会流进 matmul 报出无关错误，甚至静默算出垃圾梯度。 - if (seq_cache_ == 0 || batch_cache_ == 0 || - !Qp_cache_.valid() || !Kp_cache_.valid() || !V_re_cache_.valid() || - !Q_normed_cache_.valid() || !K_normed_cache_.valid() || + if (seq_cache_ == 0 || batch_cache_ == 0 || !Qp_cache_.valid() || !Kp_cache_.valid() || + !V_re_cache_.valid() || !Q_normed_cache_.valid() || !K_normed_cache_.valid() || !Q_rms_inv_cache_.valid() || !K_rms_inv_cache_.valid()) { - NN_FAIL("ReLULinearAttention::backward: forward 缓存缺失" "（checkpoint 模式需先 forward_recompute）"); + NN_FAIL("ReLULinearAttention::backward: forward 缓存缺失" + "（checkpoint 模式需先 forward_recompute）"); } const std::size_t seq = seq_cache_; const std::size_t batch = batch_cache_; @@ -582,8 +572,7 @@ class ReLULinearAttention final : public Layer std::vector bm(batch * seq, Scalar{0}); // 宿主桥（17 §3 D11） for (std::size_t i = 0; i < boundary.size(); ++i) bm[i] = static_cast(boundary[i]); - auto bt = detail::upload_span(engine, 1, batch * seq, Precision::F32, - std::span(bm)); + auto bt = detail::upload_span(engine, 1, batch * seq, Precision::F32, std::span(bm)); NN_TRY_CHECK(bt); boundary_t = std::move(*bt); has_bnd = true; @@ -598,25 +587,22 @@ class ReLULinearAttention final : public Layer // ── z-scan：获取 z 和 q·z（backward 需要 z 向量） ──────────── NN_TRY(ones_r, ensure_ones_(engine, BH, d_k_, seq)); - auto Z_sc = engine.scan_prefix_outer(Kp_cache_, V_ones_cache_, Qp_cache_, V_ones_cache_, - dummy, dummy, false, - d_k_, num_heads_, causal_, - boundary_t, has_bnd); + auto Z_sc = + engine.scan_prefix_outer(Kp_cache_, V_ones_cache_, Qp_cache_, V_ones_cache_, dummy, + dummy, false, d_k_, num_heads_, causal_, boundary_t, has_bnd); NN_TRY_CHECK(Z_sc); auto u_r = engine.slice_rows(*Z_sc, 0, BHdk); // q·z（标量重复 dk 次） NN_TRY_CHECK(u_r); auto z2_r = engine.slice_rows(*Z_sc, 2 * BHdk, BHdk); // dk·z（向量） NN_TRY_CHECK(z2_r); - auto z_inv_r = dsl::compute(engine, - dsl::leaf(*z2_r) * dsl::rparam(Scalar{1} / static_cast(d_k_)), + auto z_inv_r = dsl::compute( + engine, dsl::leaf(*z2_r) * dsl::rparam(Scalar{1} / static_cast(d_k_)), z2_r->rows(), z2_r->cols(), p_.compute); NN_TRY_CHECK(z_inv_r); // z = [2)/dk // ── pass 1：主扫描，读出 B^T·g 和 r = g·(B·q) ──────────── - auto Sc = engine.scan_prefix_outer(Kp_cache_, V_re_cache_, Qp_cache_, gcr, - dummy, dummy, false, - d_k_, num_heads_, causal_, - boundary_t, has_bnd); + auto Sc = engine.scan_prefix_outer(Kp_cache_, V_re_cache_, Qp_cache_, gcr, dummy, dummy, + false, d_k_, num_heads_, causal_, boundary_t, has_bnd); NN_TRY_CHECK(Sc); auto BTR_r = engine.slice_rows(*Sc, 2 * BHdk, BHdk); // B^T·g NN_TRY_CHECK(BTR_r); @@ -628,21 +614,20 @@ class ReLULinearAttention final : public Layer // scale = -r/(u+ε)² （分母修正系数，用于 gQ/gK） // gQt = B^T·g/(u+ε) + scale·z （gQ 完整表达式） const std::size_t cRows = (*u_r).rows(), cCols = (*u_r).cols(); - auto gnum_r = dsl::compute(engine, - dsl::leaf(gcr) / (dsl::leaf(*u_r) + Scalar{1e-4}), - cRows, cCols, p_.compute); + auto gnum_r = dsl::compute(engine, dsl::leaf(gcr) / (dsl::leaf(*u_r) + Scalar{1e-4}), cRows, + cCols, p_.compute); NN_TRY_CHECK(gnum_r); auto scale_r = dsl::compute(engine, - dsl::neg(dsl::leaf(*r_r)) - / ((dsl::leaf(*u_r) + Scalar{1e-4}) * (dsl::leaf(*u_r) + Scalar{1e-4})), - cRows, cCols, p_.compute); + dsl::neg(dsl::leaf(*r_r)) / ((dsl::leaf(*u_r) + Scalar{1e-4}) * + (dsl::leaf(*u_r) + Scalar{1e-4})), + cRows, cCols, p_.compute); NN_TRY_CHECK(scale_r); // -r/den²（标量重复 dk 次） // ── gQ = B^T·gnum + scale·z（全链融合为单 kernel） ───────────── auto gQt_r = dsl::compute(engine, - dsl::leaf(*BTR_r) / (dsl::leaf(*u_r) + Scalar{1e-4}) - + dsl::leaf(*scale_r) * dsl::leaf(*z_inv_r), - cRows, cCols, p_.compute); + dsl::leaf(*BTR_r) / (dsl::leaf(*u_r) + Scalar{1e-4}) + + dsl::leaf(*scale_r) * dsl::leaf(*z_inv_r), + cRows, cCols, p_.compute); NN_TRY_CHECK(gQt_r); Tensor gQt = std::move(*gQt_r); @@ -654,9 +639,8 @@ class ReLULinearAttention final : public Layer if (causal_) { // 因果：S_B = suffix(dB)，gV = S_B·k，gK_B = S_B^T·v - auto SBc_r = engine.scan_suffix_outer(*dB_r, Kp_cache_, V_re_cache_, - d_k_, num_heads_, true, - boundary_t, has_bnd); + auto SBc_r = engine.scan_suffix_outer(*dB_r, Kp_cache_, V_re_cache_, d_k_, num_heads_, + true, boundary_t, has_bnd); NN_TRY_CHECK(SBc_r); NN_TRY(gv_r, engine.slice_rows(*SBc_r, 0, BHdk)); gVt = std::move(*gv_r); @@ -664,14 +648,12 @@ class ReLULinearAttention final : public Layer // suffix(scale·q)：D = outer(q, e_0, scale)，suffix(D)·e_0 NN_TRY(D_zq_r, engine.outer_col(Qp_cache_, e_0_cache_, *scale_r, d_k_, true)); - auto SZ_r = engine.scan_suffix_outer(*D_zq_r, e_0_cache_, e_0_cache_, - d_k_, num_heads_, true, - boundary_t, has_bnd); + auto SZ_r = engine.scan_suffix_outer(*D_zq_r, e_0_cache_, e_0_cache_, d_k_, num_heads_, + true, boundary_t, has_bnd); NN_TRY_CHECK(SZ_r); NN_TRY(suffix_sq_r, engine.slice_rows(*SZ_r, 0, BHdk)); - auto gk_r = dsl::compute(engine, - dsl::leaf(*gK_B_r) + dsl::leaf(*suffix_sq_r), - gK_B_r->rows(), gK_B_r->cols(), p_.compute); + auto gk_r = dsl::compute(engine, dsl::leaf(*gK_B_r) + dsl::leaf(*suffix_sq_r), + gK_B_r->rows(), gK_B_r->cols(), p_.compute); NN_TRY_CHECK(gk_r); gKt = std::move(*gk_r); } @@ -679,22 +661,19 @@ class ReLULinearAttention final : public Layer { // 双向：A/B 为全集常数 → dB 在所有位置广播相同值 // 归约步用 dsl::compute_reduce - auto dB_sum_r = dsl::compute_reduce(engine, - dsl::row_reduce_sum(dsl::leaf(*dB_r)), - dB_r->rows(), dB_r->cols(), p_.compute); + auto dB_sum_r = dsl::compute_reduce(engine, dsl::row_reduce_sum(dsl::leaf(*dB_r)), + dB_r->rows(), dB_r->cols(), p_.compute); NN_TRY_CHECK(dB_sum_r); // Bb = row_broadcast(dB_sum)：单条表达式一次 dispatch 完成， // 不物化中间张量。 // 注：IR 规定"输出 = 最后一条指令的 dst"，故**单视图表达式不合法** // （指令表为空会被 validate_expr_spec 拒绝）。这里与一个**运行时 0** // （RParam，编译期无法被常量折叠掉）相加，使表达式合法且语义不变。 - auto Bb_r = dsl::compute(engine, - dsl::row_broadcast(*dB_sum_r) + dsl::rparam(Scalar{0}), - BHdk * d_k_, seq, p_.compute); + auto Bb_r = dsl::compute(engine, dsl::row_broadcast(*dB_sum_r) + dsl::rparam(Scalar{0}), + BHdk * d_k_, seq, p_.compute); NN_TRY_CHECK(Bb_r); - auto SBc_r = engine.scan_suffix_outer(*Bb_r, Kp_cache_, V_re_cache_, - d_k_, num_heads_, false, - dummy, false); + auto SBc_r = engine.scan_suffix_outer(*Bb_r, Kp_cache_, V_re_cache_, d_k_, num_heads_, + false, dummy, false); NN_TRY_CHECK(SBc_r); NN_TRY(gv_r, engine.slice_rows(*SBc_r, 0, BHdk)); gVt = std::move(*gv_r); @@ -702,28 +681,25 @@ class ReLULinearAttention final : public Layer // 双向 gK 的常数项：Σ_t scale_t · q_t（逐元素链与行归约融合为单次 // dispatch，不物化 scale_q 中间张量） - auto tsq_r = dsl::compute_reduce(engine, - dsl::row_reduce_sum(dsl::leaf(*scale_r) * dsl::leaf(Qp_cache_)), - BHdk, seq, p_.compute); + auto tsq_r = dsl::compute_reduce( + engine, dsl::row_reduce_sum(dsl::leaf(*scale_r) * dsl::leaf(Qp_cache_)), BHdk, seq, + p_.compute); NN_TRY_CHECK(tsq_r); // gk = gK_B + row_broadcast(tsq)：单条表达式一次 dispatch 完成 // （不物化中间张量） - auto gk_r = dsl::compute(engine, - dsl::leaf(*gK_B_r) + dsl::row_broadcast(*tsq_r), - gK_B_r->rows(), gK_B_r->cols(), p_.compute); + auto gk_r = dsl::compute(engine, dsl::leaf(*gK_B_r) + dsl::row_broadcast(*tsq_r), + gK_B_r->rows(), gK_B_r->cols(), p_.compute); NN_TRY_CHECK(gk_r); gKt = std::move(*gk_r); } // ── ReLU 反向：(x>0) ? g : 0 ─────────────────────────────── - auto gq_relu = dsl::compute(engine, - dsl::select(dsl::leaf(Qp_cache_) > Scalar{0}, - dsl::leaf(gQt), Scalar{0}), + auto gq_relu = dsl::compute( + engine, dsl::select(dsl::leaf(Qp_cache_) > Scalar{0}, dsl::leaf(gQt), Scalar{0}), gQt.rows(), gQt.cols(), p_.compute); NN_TRY_CHECK(gq_relu); - auto gk_relu = dsl::compute(engine, - dsl::select(dsl::leaf(Kp_cache_) > Scalar{0}, - dsl::leaf(gKt), Scalar{0}), + auto gk_relu = dsl::compute( + engine, dsl::select(dsl::leaf(Kp_cache_) > Scalar{0}, dsl::leaf(gKt), Scalar{0}), gKt.rows(), gKt.cols(), p_.compute); NN_TRY_CHECK(gk_relu); @@ -736,12 +712,12 @@ class ReLULinearAttention final : public Layer // ── RLA-2 RMSNorm 反向（RoPE 反向之后、rearrange 之前） ────── { const std::size_t BHrms = batch * num_heads_; - auto gq_rn = rms_norm_backward_(engine, *gq_relu, Q_normed_cache_, - Q_rms_inv_cache_, BHrms, d_k_, seq); + auto gq_rn = rms_norm_backward_(engine, *gq_relu, Q_normed_cache_, Q_rms_inv_cache_, + BHrms, d_k_, seq); NN_TRY_CHECK(gq_rn); gq_relu = std::move(*gq_rn); - auto gk_rn = rms_norm_backward_(engine, *gk_relu, K_normed_cache_, - K_rms_inv_cache_, BHrms, d_k_, seq); + auto gk_rn = rms_norm_backward_(engine, *gk_relu, K_normed_cache_, K_rms_inv_cache_, + BHrms, d_k_, seq); NN_TRY_CHECK(gk_rn); gk_relu = std::move(*gk_rn); } @@ -768,8 +744,8 @@ class ReLULinearAttention final : public Layer NN_TRY(giv, w_v_.backward(gv_r)); // grad_input = gq + gk + gv：三路累加**原地**融合为单趟（目标传递， // 不额外分配）；求和按 (gq + gk) + gv 的固定顺序结合，跨 run 确定 - auto acc = dsl::compute_into(engine, - dsl::leaf(*giq) + dsl::leaf(*gik) + dsl::leaf(*giv), *giq); + auto acc = + dsl::compute_into(engine, dsl::leaf(*giq) + dsl::leaf(*gik) + dsl::leaf(*giv), *giq); NN_TRY_CHECK(acc); return giq; } @@ -779,9 +755,8 @@ class ReLULinearAttention final : public Layer // B += k' ⊗ v, z += k' // num = B·q', den = q'·z + ε, out = num / den // 每步 O(d_k²)，与序列长度无关。 - [[nodiscard]] Result forward_step( - ComputeEngine& engine, const Tensor& input, - Tensor& B_state, Tensor& z_state, std::size_t pos) + [[nodiscard]] Result forward_step(ComputeEngine& engine, const Tensor& input, + Tensor& B_state, Tensor& z_state, std::size_t pos) { auto q_res = w_q_.forward(input); // (d_model, 1) = (H*dk, 1) NN_TRY_CHECK(q_res); @@ -800,33 +775,35 @@ class ReLULinearAttention final : public Layer NN_TRY(qr, pos_->apply_qk_step(engine, Q, pos, false)); NN_TRY(kr, pos_->apply_qk_step(engine, K, pos, false)); } - auto Qp = dsl::compute(engine, - dsl::max(dsl::leaf(Q), Scalar{0}), Q.rows(), Q.cols(), p_.compute); + auto Qp = + dsl::compute(engine, dsl::max(dsl::leaf(Q), Scalar{0}), Q.rows(), Q.cols(), p_.compute); NN_TRY_CHECK(Qp); - auto Kp = dsl::compute(engine, - dsl::max(dsl::leaf(K), Scalar{0}), K.rows(), K.cols(), p_.compute); + auto Kp = + dsl::compute(engine, dsl::max(dsl::leaf(K), Scalar{0}), K.rows(), K.cols(), p_.compute); NN_TRY_CHECK(Kp); const std::size_t H = num_heads_; const std::size_t dk = d_k_; // 更新运行态：B += k' ⊗ v，z += k'（原地累加，dsl::compute_into） // dsl::matmul(batch)：纯 {0,1} 结构（attention backward 已登记同 key） - auto B_add_r = dsl::compute(engine, // v·k'^T → (H*dk, dk) - dsl::matmul(V, *Kp, false, true, H), - V.rows(), Kp->rows() / H, p_.compute); + auto B_add_r = + dsl::compute(engine, // v·k'^T → (H*dk, dk) + dsl::matmul(V, *Kp, false, true, H), V.rows(), Kp->rows() / H, p_.compute); NN_TRY_CHECK(B_add_r); - { auto r = dsl::compute_into(engine, - dsl::leaf(B_state) + dsl::leaf(*B_add_r), B_state); - NN_TRY_CHECK(r); } - { auto r = dsl::compute_into(engine, // z += k'（(H*dk, 1)） - dsl::leaf(z_state) + dsl::leaf(*Kp), z_state); - NN_TRY_CHECK(r); } + { + auto r = dsl::compute_into(engine, dsl::leaf(B_state) + dsl::leaf(*B_add_r), B_state); + NN_TRY_CHECK(r); + } + { + auto r = dsl::compute_into(engine, // z += k'（(H*dk, 1)） + dsl::leaf(z_state) + dsl::leaf(*Kp), z_state); + NN_TRY_CHECK(r); + } // num = B·q'：batched_matmul(B_state, Qp, H) → (H*dk, 1) // dsl::matmul(batch)：纯 {0,0} 结构（scan 手工登记同 key） - auto num_r = dsl::compute(engine, - dsl::matmul(B_state, *Qp, false, false, H), - B_state.rows(), Qp->cols(), p_.compute); + auto num_r = dsl::compute(engine, dsl::matmul(B_state, *Qp, false, false, H), + B_state.rows(), Qp->cols(), p_.compute); NN_TRY_CHECK(num_r); // den = q'·z：逐头标量点积（通过宿主桥在 CPU 上计算，17 §3 D11； @@ -844,16 +821,14 @@ class ReLULinearAttention final : public Layer for (std::size_t j = 0; j < dk; ++j) den_full[h * dk + j] = dot; } - auto den_t = detail::upload_span(engine, H * dk, 1, Precision::F32, - std::span(den_full)); + auto den_t = detail::upload_span(engine, H * dk, 1, Precision::F32, std::span(den_full)); NN_TRY_CHECK(den_t); // out = num / (den + ε)：与 forward 的除法（:491）**逐 token 同构**—— // 同一个 AOT 键由 scan_exprs 的 ReLULinearAttention dry-run 覆盖， // GPU 闭合世界可命中。 - auto out_r = dsl::compute(engine, - dsl::leaf(*num_r) / (dsl::leaf(*den_t) + Scalar{1e-4}), - (*num_r).rows(), (*num_r).cols(), p_.compute); + auto out_r = dsl::compute(engine, dsl::leaf(*num_r) / (dsl::leaf(*den_t) + Scalar{1e-4}), + (*num_r).rows(), (*num_r).cols(), p_.compute); NN_TRY_CHECK(out_r); return w_o_.forward(*out_r); } @@ -872,16 +847,13 @@ class RAPTBlock final : public Layer ActivationOffloader offloader_; // L1-offload（与 GPTBlock 共用实现） public: - RAPTBlock(std::size_t d_model, std::size_t num_heads, std::size_t d_ff, - std::size_t seq_len, PosEncodingType pos_enc, - ActivationType activation = ActivationType::GeLU, - NormType norm_type = NormType::LayerNorm, - bool causal = true, + RAPTBlock(std::size_t d_model, std::size_t num_heads, std::size_t d_ff, std::size_t seq_len, + PosEncodingType pos_enc, ActivationType activation = ActivationType::GeLU, + NormType norm_type = NormType::LayerNorm, bool causal = true, PrecisionProfile precision = PrecisionProfile{}) : norm1_(make_norm_layer(d_model, norm_type)), attn_(d_model, num_heads, seq_len, causal, pos_enc), - norm2_(make_norm_layer(d_model, norm_type)), - ff_(d_model, d_ff, activation) + norm2_(make_norm_layer(d_model, norm_type)), ff_(d_model, d_ff, activation) { // D7：将精度配置注入所有子层（§9.2）—— override 内已逐子层下传 set_precision_profile(precision); @@ -913,17 +885,13 @@ class RAPTBlock final : public Layer std::vector parameters() override { - return collect_refs(norm1_->parameters(), - attn_.parameters(), - norm2_->parameters(), + return collect_refs(norm1_->parameters(), attn_.parameters(), norm2_->parameters(), ff_.parameters()); } std::vector param_gradients() override { - return collect_refs(norm1_->param_gradients(), - attn_.param_gradients(), - norm2_->param_gradients(), - ff_.param_gradients()); + return collect_refs(norm1_->param_gradients(), attn_.param_gradients(), + norm2_->param_gradients(), ff_.param_gradients()); } void set_checkpoint_mode(bool enabled) override @@ -937,8 +905,10 @@ class RAPTBlock final : public Layer void clear_cache() override { - norm1_->clear_cache(); attn_.clear_cache(); - norm2_->clear_cache(); ff_.clear_cache(); + norm1_->clear_cache(); + attn_.clear_cache(); + norm2_->clear_cache(); + ff_.clear_cache(); } // RAPTBlock 与 GPTBlock 一样是"重计算单元"。 @@ -947,8 +917,7 @@ class RAPTBlock final : public Layer // 重计算：必须走**虚函数** set_checkpoint_mode 关闭本块+子层的 checkpoint // 模式（基类默认实现只改本块的标志位，子层仍处于 checkpoint 模式 → // forward 不会重建子层缓存 → backward 要么报错、要么误用上一 step 的陈旧缓存）。 - [[nodiscard]] Result forward_recompute( - const Tensor& saved_input) override + [[nodiscard]] Result forward_recompute(const Tensor& saved_input) override { set_checkpoint_mode(false); auto r = forward(saved_input); @@ -961,17 +930,11 @@ class RAPTBlock final : public Layer void set_offload_enabled(bool enabled) { offloader_.set_enabled(enabled); } [[nodiscard]] bool offload_enabled() const noexcept { return offloader_.enabled(); } [[nodiscard]] Result export_activations(ComputeEngine& engine) - { - return offloader_.export_activations(engine, activation_cache()); - } + { return offloader_.export_activations(engine, activation_cache()); } [[nodiscard]] Result import_activations(ComputeEngine& engine) - { - return offloader_.import_activations(engine); - } + { return offloader_.import_activations(engine); } [[nodiscard]] std::size_t offload_slab_bytes() const noexcept - { - return offloader_.slab_bytes(); - } + { return offloader_.slab_bytes(); } // 新一轮 forward 入口：上一轮导出未被 backward 消费（推理/评估）时作废 void discard_offloaded_activations() noexcept { offloader_.discard(); } @@ -979,73 +942,60 @@ class RAPTBlock final : public Layer std::vector activation_cache() override { std::vector r; - auto a = attn_.activation_cache(); r.insert(r.end(), a.begin(), a.end()); - auto n1 = norm1_->activation_cache(); r.insert(r.end(), n1.begin(), n1.end()); - auto n2 = norm2_->activation_cache(); r.insert(r.end(), n2.begin(), n2.end()); - auto f = ff_.activation_cache(); r.insert(r.end(), f.begin(), f.end()); + auto a = attn_.activation_cache(); + r.insert(r.end(), a.begin(), a.end()); + auto n1 = norm1_->activation_cache(); + r.insert(r.end(), n1.begin(), n1.end()); + auto n2 = norm2_->activation_cache(); + r.insert(r.end(), n2.begin(), n2.end()); + auto f = ff_.activation_cache(); + r.insert(r.end(), f.begin(), f.end()); return r; } // 文档感知：转发给内部 RLA-2 注意力（文档边界处重置运行态） - void set_doc_ids(std::span ids) override - { - attn_.set_doc_ids(ids); - } + void set_doc_ids(std::span ids) override { attn_.set_doc_ids(ids); } // 绝对位置偏移：转发给内部 RLA-2 注意力（滑动窗生成时 RoPE 用绝对位置） - void set_position_offset(std::size_t off) - { - attn_.set_position_offset(off); - } + void set_position_offset(std::size_t off) { attn_.set_position_offset(off); } - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); // 两个残差分支（与 GPTBlock/TransformerEncoderLayer 同骨架）。 // P2-3：残差相加的精度实参改为 `p_.compute`（此前硬编码 F32）—— // 与 GPTBlock / TransformerEncoderLayer 同口径；默认 profile 的 // compute = F32 → 逐位不变，`--f16` 下残差走 f16 存储。 - auto r1 = prenorm_residual_forward_(engine, input, *norm1_, attn_, - p_.compute); + auto r1 = prenorm_residual_forward_(engine, input, *norm1_, attn_, p_.compute); NN_TRY_CHECK(r1); - return prenorm_residual_forward_(engine, *r1, *norm2_, ff_, - p_.compute); + return prenorm_residual_forward_(engine, *r1, *norm2_, ff_, p_.compute); } - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { ComputeEngine& engine = engine_ref(); // activation offload：从 host 恢复激活再反向（替代重计算） - if (offloader_.offloaded()) - { - NN_TRY(im, offloader_.import_activations(engine)); - } + if (offloader_.offloaded()) { NN_TRY(im, offloader_.import_activations(engine)); } // 两处残差分流 + 子层反向（与 GPTBlock / TransformerEncoderLayer 同一骨架）。 // P2-3：精度实参改为 `p_.compute`（同前向；默认 profile 逐位不变）。 - return prenorm_residual_backward_(engine, grad_output, - *norm1_, attn_, *norm2_, ff_, + return prenorm_residual_backward_(engine, grad_output, *norm1_, attn_, *norm2_, ff_, p_.compute); } // 增量推理：单 token → norm1 → RLA-2 运行态注意力 → 残差 → norm2 → FFN → 残差 // B_state/z_state: RLA-2 运行态 KV cache（见 ReLULinearAttention::forward_step） - [[nodiscard]] Result forward_step( - ComputeEngine& engine, const Tensor& input, - Tensor& B_state, Tensor& z_state, std::size_t pos) + [[nodiscard]] Result forward_step(ComputeEngine& engine, const Tensor& input, + Tensor& B_state, Tensor& z_state, std::size_t pos) { NN_TRY(n1, norm1_->forward(input)); NN_TRY(a, attn_.forward_step(engine, *n1, B_state, z_state, pos)); - auto r1 = dsl::compute(engine, - dsl::leaf(input) + dsl::leaf(*a), - input.rows(), input.cols(), p_.compute); + auto r1 = dsl::compute(engine, dsl::leaf(input) + dsl::leaf(*a), input.rows(), input.cols(), + p_.compute); NN_TRY_CHECK(r1); NN_TRY(n2, norm2_->forward(*r1)); NN_TRY(f, ff_.forward(*n2)); - return dsl::compute(engine, - dsl::leaf(*r1) + dsl::leaf(*f), - r1->rows(), r1->cols(), p_.compute); + return dsl::compute(engine, dsl::leaf(*r1) + dsl::leaf(*f), r1->rows(), r1->cols(), + p_.compute); } }; @@ -1092,13 +1042,10 @@ class RAPTModel final : public Layer std::size_t num_heads, std::size_t d_ff, std::size_t num_layers, PosEncodingType pos_enc = PosEncodingType::RoPE, ActivationType activation = ActivationType::GeLU, - NormType norm_type = NormType::LayerNorm, - bool causal = true, + NormType norm_type = NormType::LayerNorm, bool causal = true, PrecisionProfile precision = PrecisionProfile{}) - : vocab_size_(vocab_size), d_model_(d_model), seq_len_(seq_len), - num_heads_(num_heads), - ln_f_(make_norm_layer(d_model, norm_type)), - lm_head_(d_model, vocab_size) + : vocab_size_(vocab_size), d_model_(d_model), seq_len_(seq_len), num_heads_(num_heads), + ln_f_(make_norm_layer(d_model, norm_type)), lm_head_(d_model, vocab_size) { // D7：将精度配置注入自身和所有子层（§9.2） set_precision_profile(precision); @@ -1113,8 +1060,8 @@ class RAPTModel final : public Layer 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, pos_enc, - activation, norm_type, causal, precision); + blocks_.emplace_back(d_model, num_heads, d_ff, seq_len, pos_enc, activation, norm_type, + causal, precision); } pos_encoder_->set_precision_profile(precision); } @@ -1127,7 +1074,8 @@ class RAPTModel final : public Layer Layer::set_precision_profile(profile); if (ln_f_) ln_f_->set_precision_profile(profile); lm_head_.set_precision_profile(profile); - for (auto& b : blocks_) b.set_precision_profile(profile); + for (auto& b : blocks_) + b.set_precision_profile(profile); if (pos_encoder_) pos_encoder_->set_precision_profile(profile); } @@ -1137,49 +1085,39 @@ class RAPTModel final : public Layer constexpr Scalar emb_init_std = 0.02; token_emb_ = engine.create_tensor(vocab_size_, d_model_, p_.param, InitSpec::normal(0, emb_init_std, kInitSeed)); - if (!token_emb_.valid()) - NN_FAIL("RAPTModel: token_emb 初始化失败"); - grad_token_emb_ = engine.create_tensor(vocab_size_, d_model_, p_.param, - InitSpec::zero()); - if (!grad_token_emb_.valid()) - NN_FAIL("RAPTModel: token_emb 梯度缓冲初始化失败"); - if (pos_encoder_) - { - NN_TRY(r, pos_encoder_->init(engine)); - } + if (!token_emb_.valid()) NN_FAIL("RAPTModel: token_emb 初始化失败"); + grad_token_emb_ = engine.create_tensor(vocab_size_, d_model_, p_.param, InitSpec::zero()); + if (!grad_token_emb_.valid()) NN_FAIL("RAPTModel: token_emb 梯度缓冲初始化失败"); + if (pos_encoder_) { NN_TRY(r, pos_encoder_->init(engine)); } for (auto& b : blocks_) { NN_TRY(r, b.init(engine)); } - if (ln_f_) + if (ln_f_) { NN_TRY(r, ln_f_->init(engine)); } { - NN_TRY(r, ln_f_->init(engine)); + NN_TRY(r, lm_head_.init(engine)); } - { NN_TRY(r, lm_head_.init(engine)); } return {}; } std::vector parameters() override { - return collect_refs(token_emb_, - pos_encoder_->parameters(), + return collect_refs(token_emb_, pos_encoder_->parameters(), collect_block_refs_(blocks_, &RAPTBlock::parameters), - ln_f_->parameters(), - lm_head_.parameters()); + ln_f_->parameters(), lm_head_.parameters()); } std::vector param_gradients() override { - return collect_refs(grad_token_emb_, - pos_encoder_->param_gradients(), + return collect_refs(grad_token_emb_, pos_encoder_->param_gradients(), collect_block_refs_(blocks_, &RAPTBlock::param_gradients), - ln_f_->param_gradients(), - lm_head_.param_gradients()); + ln_f_->param_gradients(), lm_head_.param_gradients()); } void set_checkpoint_mode(bool enabled) override { Layer::set_checkpoint_mode(enabled); - for (auto& b : blocks_) b.set_checkpoint_mode(enabled); + for (auto& b : blocks_) + b.set_checkpoint_mode(enabled); ln_f_->set_checkpoint_mode(enabled); lm_head_.set_checkpoint_mode(enabled); } @@ -1191,7 +1129,8 @@ class RAPTModel final : public Layer // 同样保留），清理会让梯度写不回词嵌入表。 void clear_cache() override { - for (auto& b : blocks_) b.clear_cache(); + for (auto& b : blocks_) + b.clear_cache(); ln_f_->clear_cache(); lm_head_.clear_cache(); checkpoint_inputs_.clear(); @@ -1200,14 +1139,19 @@ class RAPTModel final : public Layer // 文档感知：记录 doc_ids，forward 时下发给各块（文档边界处重置 RLA-2 运行态） void set_doc_ids(std::span ids) override { - if (ids.empty()) { doc_ids_.clear(); return; } + if (ids.empty()) + { + doc_ids_.clear(); + return; + } doc_ids_.assign(ids.begin(), ids.end()); } // 绝对位置偏移：下发给各块（滑动窗生成时 RoPE 用绝对位置，非 0..seq-1 重置） void set_position_offset(std::size_t off) { - for (auto& b : blocks_) b.set_position_offset(off); + for (auto& b : blocks_) + b.set_position_offset(off); } // ── batch 录制粒度（防 TDR）：每 N 个块提交一次，0 = 不启用 ── @@ -1222,7 +1166,8 @@ class RAPTModel final : public Layer void set_activation_offload(bool enabled) override { activation_offload_ = enabled; - for (auto& b : blocks_) b.set_offload_enabled(enabled); + for (auto& b : blocks_) + b.set_offload_enabled(enabled); } [[nodiscard]] bool activation_offload() const noexcept { return activation_offload_; } @@ -1230,19 +1175,20 @@ class RAPTModel final : public Layer [[nodiscard]] std::size_t offload_ram_bytes() override { std::size_t total = 0; - for (auto& b : blocks_) total += b.offload_slab_bytes(); + for (auto& b : blocks_) + total += b.offload_slab_bytes(); return total; } - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); // 新一轮 forward：上一轮 offload 的激活若未被 backward 消费（推理/评估 // 路径），其 slab 数据已作废——入口统一丢弃，保证 export 状态机从 // Idle 开始。 if (activation_offload_) - for (auto& b : blocks_) b.discard_offloaded_activations(); + for (auto& b : blocks_) + b.discard_offloaded_activations(); const std::size_t seq = input.rows(); const std::size_t batch = input.cols(); batch_size_ = batch; @@ -1286,8 +1232,7 @@ class RAPTModel final : public Layer NN_TRY(ex, b.export_activations(engine)); } // 按间隔 flush，将大录制拆成多个小提交（防 TDR） - if (flush_interval_ > 0 && (bi + 1) % flush_interval_ == 0 - && bi + 1 < blocks_.size()) + if (flush_interval_ > 0 && (bi + 1) % flush_interval_ == 0 && bi + 1 < blocks_.size()) { NN_TRY(fr, engine.flush_batch()); } @@ -1296,8 +1241,7 @@ class RAPTModel final : public Layer return lm_head_.forward(*ln); } - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { ComputeEngine& engine = engine_ref(); const std::size_t seq = seq_len_; @@ -1323,11 +1267,11 @@ class RAPTModel final : public Layer } auto br = blocks_[idx].backward(grad_x); if (!br) - NN_FAIL("RAPTModel::backward: block " + std::to_string(idx) + " failed: " + br.error().message); + NN_FAIL("RAPTModel::backward: block " + std::to_string(idx) + + " failed: " + br.error().message); grad_x = std::move(*br); // 重算/恢复出来的激活用后即释放，避免跨块累积（否则抵消省显存收益） - if (checkpoint_every_ > 0 || activation_offload_) - blocks_[idx].clear_cache(); + if (checkpoint_every_ > 0 || activation_offload_) blocks_[idx].clear_cache(); if (flush_interval_ > 0 && (bi + 1) % flush_interval_ == 0 && bi + 1 < n) { NN_TRY(fr, engine.flush_batch()); @@ -1339,10 +1283,8 @@ class RAPTModel final : public Layer NN_TRY(pr, pos_encoder_->backward(engine, *grad_T, batch, seq)); NN_TRY(sr, engine.scatter_add_rows(grad_token_emb_, stored_tokens_tensor_, *grad_T)); - Tensor grad_input = engine.create_tensor(seq, batch, Precision::F32, - InitSpec::zero()); - if (!grad_input.valid()) - NN_FAIL("RAPT token_emb backward: 梯度张量分配失败"); + Tensor grad_input = engine.create_tensor(seq, batch, Precision::F32, InitSpec::zero()); + if (!grad_input.valid()) NN_FAIL("RAPT token_emb backward: 梯度张量分配失败"); return grad_input; } @@ -1356,27 +1298,28 @@ class RAPTModel final : public Layer std::size_t eos_token_id = static_cast(-1), std::size_t min_new_tokens = 0) { - if (prompt.empty()) - NN_FAIL("RAPT generate: empty prompt"); + if (prompt.empty()) NN_FAIL("RAPT generate: empty prompt"); const std::size_t dk = d_model_ / num_heads_; // RLA-2 每块两个运行态：B_state (d_model, d_k) + z_state (d_model, 1) std::vector statesB, statesZ; for (std::size_t i = 0; i < blocks_.size(); ++i) { statesB.emplace_back(engine.create_tensor(d_model_, dk)); - { NN_TRY(r, engine.zero(statesB.back())); } + { + NN_TRY(r, engine.zero(statesB.back())); + } statesZ.emplace_back(engine.create_tensor(d_model_, 1)); - { NN_TRY(r, engine.zero(statesZ.back())); } + { + NN_TRY(r, engine.zero(statesZ.back())); + } } // 处理单个 token：embed → 各块 forward_step → LN → LM head → 返回 logits 列 - auto step_one = [&](std::size_t tok, std::size_t pos) - -> Result> + auto step_one = [&](std::size_t tok, std::size_t pos) -> Result> { std::vector idx_v(1); // 宿主桥（17 §3 D11） idx_v[0] = static_cast(tok); - auto idx_t = detail::upload_span(engine, 1, 1, Precision::F32, - std::span(idx_v)); + auto idx_t = detail::upload_span(engine, 1, 1, Precision::F32, std::span(idx_v)); NN_TRY_CHECK(idx_t); auto emb = engine.gather_rows(token_emb_, *idx_t); // (1, d_model) NN_TRY_CHECK(emb); diff --git a/include/neuralnet.cpp/compute_layer_softmax.hpp b/include/neuralnet.cpp/compute_layer_softmax.hpp index cad1e357..8c62aed0 100644 --- a/include/neuralnet.cpp/compute_layer_softmax.hpp +++ b/include/neuralnet.cpp/compute_layer_softmax.hpp @@ -50,8 +50,7 @@ class Softmax final : public Layer return r; } - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); // 行 softmax（数值稳定）：out = exp(x - row_max) / Σ_c exp(x - row_max) @@ -59,30 +58,28 @@ class Softmax final : public Layer // Tensor（shifted/exp_shift/row_max/row_sum 的物化）由融合 kernel 消解， // 仅 input 与 output 落显存。 // 表达式文本只写在本 Layer；AOT 收集由 scan_exprs dry-run 本方法完成。 - auto out = dsl::compute(engine, - dsl::exp(dsl::leaf(input) - dsl::row_reduce_max(input)) - / dsl::row_reduce_sum( - dsl::exp(dsl::leaf(input) - dsl::row_reduce_max(input))), + auto out = dsl::compute( + engine, + dsl::exp(dsl::leaf(input) - dsl::row_reduce_max(input)) / + dsl::row_reduce_sum(dsl::exp(dsl::leaf(input) - dsl::row_reduce_max(input))), p_.stable); // 输出仍是整网格（归约视图按行广播）→ 由表达式推导 NN_TRY_CHECK(out); - if (checkpoint_mode_) - return out; + if (checkpoint_mode_) return out; // 单缓冲：把结果移入 output_cache_（唯一持有者），返回共享同一 buffer output_cache_ = std::move(*out); return output_cache_; } - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { ComputeEngine& engine = engine_ref(); // grad_x = out ⊙ (grad_output - row_dot(out ⊙ grad_output)) // 单表达式融合：row_dot 为归约指令，消除 ep/gmd 等全尺寸中间 Tensor。 - auto out = dsl::compute(engine, - dsl::leaf(output_cache_) - * (dsl::leaf(grad_output) - - dsl::row_reduce_sum(dsl::leaf(output_cache_) - * dsl::leaf(grad_output))), + auto out = dsl::compute( + engine, + dsl::leaf(output_cache_) * + (dsl::leaf(grad_output) - + dsl::row_reduce_sum(dsl::leaf(output_cache_) * dsl::leaf(grad_output))), p_.stable); // 输出仍是整网格（归约指令按行广播）→ 由表达式推导 return out; } diff --git a/include/neuralnet.cpp/compute_layer_transformer.hpp b/include/neuralnet.cpp/compute_layer_transformer.hpp index 078cd93a..cba19acf 100644 --- a/include/neuralnet.cpp/compute_layer_transformer.hpp +++ b/include/neuralnet.cpp/compute_layer_transformer.hpp @@ -34,14 +34,14 @@ class PositionalEncoding final : public Layer [[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; + const std::size_t batch = (tile_size_ > 0) ? total_len / tile_size_ : 1; std::vector enc(d_model_ * total_len); // 宿主桥（17 §3 D11） const std::size_t half = d_model_ / 2; std::vector freqs(half); for (std::size_t i = 0; i < half; ++i) - freqs[i] = Scalar{1} / std::pow(Scalar{10000}, - static_cast(2 * i) / static_cast(d_model_)); + freqs[i] = Scalar{1} / std::pow(Scalar{10000}, static_cast(2 * i) / + static_cast(d_model_)); for (std::size_t b = 0; b < batch; ++b) { @@ -52,13 +52,14 @@ class PositionalEncoding final : public Layer for (std::size_t i = 0; i < half; ++i) { const Scalar angle = pos_d * freqs[i]; - enc[(2 * i) * total_len + col_off + pos] = std::sin(angle); + enc[(2 * i) * total_len + col_off + pos] = std::sin(angle); enc[(2 * i + 1) * total_len + col_off + pos] = std::cos(angle); } if (d_model_ % 2 == 1) { - const Scalar freq_last = Scalar{1} / std::pow(Scalar{10000}, - static_cast(2 * half) / static_cast(d_model_)); + const Scalar freq_last = + Scalar{1} / std::pow(Scalar{10000}, static_cast(2 * half) / + static_cast(d_model_)); enc[(d_model_ - 1) * total_len + col_off + pos] = std::sin(pos_d * freq_last); } } @@ -70,35 +71,26 @@ class PositionalEncoding final : public Layer } public: - PositionalEncoding(std::size_t d_model, std::size_t max_len = 5000, - std::size_t tile_size = 0) + 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) { } - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); - if (input.rows() != d_model_) - NN_FAIL("PE forward: d_model mismatch"); + if (input.rows() != d_model_) NN_FAIL("PE forward: d_model mismatch"); const std::size_t total_len = input.cols(); if (tile_size_ > 0 && total_len % tile_size_ != 0) NN_FAIL("PE forward: total_len not divisible by tile_size"); - if (tile_size_ == 0 && total_len > max_len_) - NN_FAIL("PE forward: seq_len exceeds max_len"); - if (cached_total_ != total_len) - { - NN_TRY(r, rebuild_encoding(engine, total_len)); - } + if (tile_size_ == 0 && total_len > max_len_) NN_FAIL("PE forward: seq_len exceeds max_len"); + if (cached_total_ != total_len) { NN_TRY(r, rebuild_encoding(engine, total_len)); } - return dsl::compute(engine, - dsl::leaf(input) + dsl::leaf(encoding_cache_), - input.rows(), input.cols(), p_.compute); + return dsl::compute(engine, dsl::leaf(input) + dsl::leaf(encoding_cache_), input.rows(), + input.cols(), p_.compute); } - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { return grad_output; // 位置编码不可学习，梯度直接穿透 } @@ -124,15 +116,15 @@ class TransformerEncoderLayer final : public Layer Tensor residual2_cache_; public: - [[nodiscard]] const char* layer_name() const noexcept override { return "TransformerEncoderLayer"; } + [[nodiscard]] const char* layer_name() const noexcept override + { return "TransformerEncoderLayer"; } - TransformerEncoderLayer(std::size_t d_model, std::size_t num_heads, - std::size_t d_ff, std::size_t seq_len = 0, - NormType norm_type = NormType::LayerNorm) - : self_attn_(d_model, num_heads, seq_len), - norm1_(make_norm_layer(d_model, norm_type)), - ff_(d_model, d_ff), - norm2_(make_norm_layer(d_model, norm_type)) {} + TransformerEncoderLayer(std::size_t d_model, std::size_t num_heads, std::size_t d_ff, + std::size_t seq_len = 0, NormType norm_type = NormType::LayerNorm) + : self_attn_(d_model, num_heads, seq_len), norm1_(make_norm_layer(d_model, norm_type)), + ff_(d_model, d_ff), norm2_(make_norm_layer(d_model, norm_type)) + { + } [[nodiscard]] Result init_impl(ComputeEngine& engine) override { @@ -155,18 +147,14 @@ class TransformerEncoderLayer final : public Layer std::vector parameters() override { - return collect_refs(self_attn_.parameters(), - norm1_->parameters(), - ff_.parameters(), + return collect_refs(self_attn_.parameters(), norm1_->parameters(), ff_.parameters(), norm2_->parameters()); } std::vector param_gradients() override { - return collect_refs(self_attn_.param_gradients(), - norm1_->param_gradients(), - ff_.param_gradients(), - norm2_->param_gradients()); + return collect_refs(self_attn_.param_gradients(), norm1_->param_gradients(), + ff_.param_gradients(), norm2_->param_gradients()); } // 梯度检查点：把模式传播给内部注意力/归一化/FFN @@ -181,8 +169,7 @@ class TransformerEncoderLayer final : public Layer [[nodiscard]] bool recompute_supported() const override { return true; } - [[nodiscard]] Result forward_recompute( - const Tensor& saved_input) override + [[nodiscard]] Result forward_recompute(const Tensor& saved_input) override { // 临时关闭本层及其子层的 checkpoint 模式，使 forward 重建缓存 set_checkpoint_mode(false); @@ -211,26 +198,22 @@ class TransformerEncoderLayer final : public Layer norm2_->set_training(training); } - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); // 残差分支 1（Pre-Norm + 自注意力），与 GPTBlock/RAPTBlock 同骨架 NN_TRY(r2, prenorm_residual_forward_(engine, input, *norm1_, self_attn_, p_.compute)); Tensor res2 = std::move(*r2); - if (!checkpoint_mode_) - residual2_cache_ = res2; + if (!checkpoint_mode_) residual2_cache_ = res2; // 残差分支 2（Pre-Norm + FFN） return prenorm_residual_forward_(engine, res2, *norm2_, ff_, p_.compute); } - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { ComputeEngine& engine = engine_ref(); // 两处残差分流 + 子层反向（与 GPTBlock / RAPTBlock 同一骨架） - return prenorm_residual_backward_(engine, grad_output, - *norm1_, self_attn_, *norm2_, ff_, + return prenorm_residual_backward_(engine, grad_output, *norm1_, self_attn_, *norm2_, ff_, p_.compute); } }; @@ -270,11 +253,9 @@ class TransformerEncoder final : public Layer public: [[nodiscard]] const char* layer_name() const noexcept override { return "TransformerEncoder"; } - 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, - NormType norm_type = NormType::LayerNorm, - bool final_norm = false) + 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, + NormType norm_type = NormType::LayerNorm, bool final_norm = false) : d_model_(d_model), num_patches_(num_patches), inv_num_patches_(Scalar{1} / static_cast(num_patches)), // PE 启用 tiling：每个样本独立使用 (d_model, num_patches) 的编码 @@ -283,8 +264,7 @@ class TransformerEncoder final : public Layer // 所有 EncoderLayer 传入 seq_len=num_patches，启用 MHA 批量化路径 for (std::size_t i = 0; i < num_layers; ++i) layers_.emplace_back(d_model, num_heads, d_ff, num_patches, norm_type); - if (final_norm) - ln_f_ = make_norm_layer(d_model, norm_type); + if (final_norm) ln_f_ = make_norm_layer(d_model, norm_type); } [[nodiscard]] Result init_impl(ComputeEngine& engine) override @@ -299,13 +279,9 @@ class TransformerEncoder final : public Layer } // 预创建 ones_row_ (1, num_patches) 全1，用于 backward 广播 // （M2 声明式：引擎填数，原 from_matrix 口径 = F32） - if (ln_f_) - { - NN_TRY(rn, ln_f_->init(engine)); - } + if (ln_f_) { NN_TRY(rn, ln_f_->init(engine)); } ones_row_ = engine.create_tensor(1, num_patches_, Precision::F32, InitSpec::constant(1)); - if (!ones_row_.valid()) - NN_FAIL("TransformerEncoder: ones_row_ 初始化失败"); + if (!ones_row_.valid()) NN_FAIL("TransformerEncoder: ones_row_ 初始化失败"); return {}; } @@ -313,7 +289,8 @@ class TransformerEncoder final : public Layer void set_precision_profile(const PrecisionProfile& profile) override { Layer::set_precision_profile(profile); - for (auto& l : layers_) l.set_precision_profile(profile); + for (auto& l : layers_) + l.set_precision_profile(profile); if (ln_f_) ln_f_->set_precision_profile(profile); pos_encoding_.set_precision_profile(profile); } @@ -323,14 +300,16 @@ class TransformerEncoder final : public Layer void set_training(bool training) override { Layer::set_training(training); - for (auto& l : layers_) l.set_training(training); + for (auto& l : layers_) + l.set_training(training); if (ln_f_) ln_f_->set_training(training); } // 缓存释放同 GPTModel/RAPTModel：逐层下传（含 final norm） void clear_cache() override { - for (auto& l : layers_) l.clear_cache(); + for (auto& l : layers_) + l.clear_cache(); if (ln_f_) ln_f_->clear_cache(); } @@ -356,13 +335,11 @@ class TransformerEncoder final : public Layer return refs; } - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); // input: (d_model, batch * num_patches) - if (input.rows() != d_model_) - NN_FAIL("TransformerEncoder: row count mismatch"); + if (input.rows() != d_model_) NN_FAIL("TransformerEncoder: row count mismatch"); if (input.cols() % num_patches_ != 0) NN_FAIL("TransformerEncoder: cols not divisible by num_patches"); batch_size_ = input.cols() / num_patches_; @@ -394,19 +371,18 @@ class TransformerEncoder final : public Layer // 再 rearrange_3d inverse 回 (d_model, batch)。 NN_TRY(re, engine.rearrange_3d(x, d_model_, batch_size_, num_patches_, false)); // 归约步走 dsl::compute_reduce（输出与输入同形） - auto row_sum = dsl::compute_reduce(engine, - dsl::row_reduce_sum(dsl::leaf(*re)), re->rows(), re->cols(), p_.compute); + auto row_sum = dsl::compute_reduce(engine, dsl::row_reduce_sum(dsl::leaf(*re)), re->rows(), + re->cols(), p_.compute); NN_TRY_CHECK(row_sum); NN_TRY(rs_re, engine.rearrange_3d(*row_sum, d_model_, batch_size_, 1, true)); // 原地缩放走 dsl::compute_into（inv_num_patches_ 走 rparam，值不进 key） - auto r = dsl::compute_into(engine, - dsl::leaf(*rs_re) * dsl::rparam(inv_num_patches_), *rs_re); + auto r = + dsl::compute_into(engine, dsl::leaf(*rs_re) * dsl::rparam(inv_num_patches_), *rs_re); NN_TRY_CHECK(r); return *rs_re; } - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { ComputeEngine& engine = engine_ref(); // grad_output: (d_model, batch) — 来自下游 Linear.backward @@ -419,15 +395,13 @@ class TransformerEncoder final : public Layer // rearrange_3d(result, d_model, batch, num_patches, true) → (d_model, batch*num_patches) NN_TRY(g_re, engine.rearrange_3d(grad_output, d_model_, batch_size_, 1, false)); // 原地缩放走 dsl::compute_into（inv_num_patches_ 走 rparam，值不进 key） - auto r = dsl::compute_into(engine, - dsl::leaf(*g_re) * dsl::rparam(inv_num_patches_), *g_re); + auto r = dsl::compute_into(engine, dsl::leaf(*g_re) * dsl::rparam(inv_num_patches_), *g_re); NN_TRY_CHECK(r); // (*g_re): (batch*d_model, 1) × ones_row_ (1, num_patches) → (batch*d_model, num_patches) // 纯 matmul 走 DSL 直写（结构经 scan 的 TransformerEncoderLayer dry-run 登记； // 注意该调用在 pooling 路径，scan 的 enc dry-run 覆盖 forward/backward） - auto unpooled = dsl::compute(engine, - dsl::matmul(*g_re, ones_row_, false, false), - g_re->rows(), ones_row_.cols(), p_.compute); + auto unpooled = dsl::compute(engine, dsl::matmul(*g_re, ones_row_, false, false), + g_re->rows(), ones_row_.cols(), p_.compute); NN_TRY_CHECK(unpooled); NN_TRY(grad, engine.rearrange_3d(*unpooled, d_model_, batch_size_, num_patches_, true)); Tensor grad_x = std::move(*grad); @@ -442,8 +416,7 @@ class TransformerEncoder final : public Layer for (auto it = layers_.rbegin(); it != layers_.rend(); ++it) { - const std::size_t li = - static_cast(layers_.rend() - it) - 1; + const std::size_t li = static_cast(layers_.rend() - it) - 1; NN_PROF_SCOPE(nn::prof::interned("bwd.enc", li)); NN_TRY(br, it->backward(grad_x)); grad_x = std::move(*br); @@ -482,14 +455,10 @@ class PatchEmbedding final : public Layer public: [[nodiscard]] const char* layer_name() const noexcept override { return "PatchEmbedding"; } - 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_(patch_dim_, d_model) + 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_(patch_dim_, d_model) { // RV-01：不整除 → num_patches_ 少算、后续按错误形状索引 → 无条件终止 NN_ASSERT_ALWAYS(img_size % patch_size == 0, @@ -497,9 +466,7 @@ class PatchEmbedding final : public Layer } [[nodiscard]] Result init_impl(ComputeEngine& engine) override - { - return projection_.init(engine); - } + { return projection_.init(engine); } // ── D7：精度配置下传（§9.2）：投影层是复合层的唯一子层 ─────────────── void set_precision_profile(const PrecisionProfile& profile) override @@ -509,16 +476,13 @@ class PatchEmbedding final : public Layer } [[nodiscard]] std::size_t num_patches() const noexcept { return num_patches_; } - [[nodiscard]] std::size_t d_model() const noexcept { return d_model_; } + [[nodiscard]] std::size_t d_model() const noexcept { return d_model_; } - std::vector parameters() override - { return projection_.parameters(); } + std::vector parameters() override { return projection_.parameters(); } - std::vector param_gradients() override - { return projection_.param_gradients(); } + std::vector param_gradients() override { return projection_.param_gradients(); } - [[nodiscard]] Result forward( - const Tensor& input) override + [[nodiscard]] Result forward(const Tensor& input) override { ComputeEngine& engine = engine_ref(); input_cache_ = input; @@ -546,23 +510,22 @@ class PatchEmbedding final : public Layer for (std::size_t pc = 0; pc < patch_size_; ++pc) { const std::size_t flat = pr * patch_size_ + pc; - const std::size_t pix = (gr + pr) * img_size_ + (gc + pc); + const std::size_t pix = (gr + pr) * img_size_ + (gc + pc); all_patches[flat * (batch * num_patches_) + col_idx] = (*in_v)[pix * in_cols + b]; } } } - auto ap_t = detail::upload_span(engine, patch_dim_, batch * num_patches_, - Precision::F32, std::span(all_patches)); + auto ap_t = detail::upload_span(engine, patch_dim_, batch * num_patches_, Precision::F32, + std::span(all_patches)); NN_TRY_CHECK(ap_t); // Step 2: 投影 → (d_model, batch * num_patches) — 已是 batch-major，无需重排 return projection_.forward(*ap_t); } - [[nodiscard]] Result backward( - const Tensor& grad_output) override + [[nodiscard]] Result backward(const Tensor& grad_output) override { ComputeEngine& engine = engine_ref(); // grad_output: (d_model, batch * num_patches) — batch-major @@ -588,15 +551,14 @@ class PatchEmbedding final : public Layer for (std::size_t pc = 0; pc < patch_size_; ++pc) { const std::size_t flat = pr * patch_size_ + pc; - const std::size_t pix = (gr + pr) * img_size_ + (gc + pc); - const std::size_t gi = pix * batch + b; - grad_input[gi] = grad_input[gi] - + (*gp_v)[flat * gp_cols + col_idx]; + const std::size_t pix = (gr + pr) * img_size_ + (gc + pc); + const std::size_t gi = pix * batch + b; + grad_input[gi] = grad_input[gi] + (*gp_v)[flat * gp_cols + col_idx]; } } } - return detail::upload_span(engine, img_size_ * img_size_, batch, - Precision::F32, std::span(grad_input)); + return detail::upload_span(engine, img_size_ * img_size_, batch, Precision::F32, + std::span(grad_input)); } }; diff --git a/include/neuralnet.cpp/compute_loss.hpp b/include/neuralnet.cpp/compute_loss.hpp index d686cb4f..03f46656 100644 --- a/include/neuralnet.cpp/compute_loss.hpp +++ b/include/neuralnet.cpp/compute_loss.hpp @@ -63,8 +63,8 @@ class Loss [[nodiscard]] const PrecisionProfile& precision_profile() const noexcept { return p_; } // forward 计算损失标量，并缓存 backward 所需中间结果 - [[nodiscard]] virtual Result forward( - ComputeEngine& engine, const Tensor& pred, const Tensor& target) = 0; + [[nodiscard]] virtual Result forward(ComputeEngine& engine, const Tensor& pred, + const Tensor& target) = 0; // backward 返回对 pred 的梯度 [[nodiscard]] virtual Result backward() = 0; @@ -86,32 +86,29 @@ class MSELoss final : public Loss public: MSELoss() = default; - [[nodiscard]] Result forward( - ComputeEngine& engine, const Tensor& pred, const Tensor& target) override + [[nodiscard]] Result forward(ComputeEngine& engine, const Tensor& pred, + const Tensor& target) override { if (pred.rows() != target.rows() || pred.cols() != target.cols()) NN_FAIL("mse loss: shape mismatch"); - if (pred.size() == 0) - NN_FAIL("mse loss: empty input"); + if (pred.size() == 0) NN_FAIL("mse loss: empty input"); const Scalar total = static_cast(pred.size()); const Scalar scale = Scalar{2} / total; // diff = pred - target（逐元素融合） - auto diff = dsl::compute(engine, - dsl::leaf(pred) - dsl::leaf(target), - pred.rows(), pred.cols(), p_.stable); + auto diff = dsl::compute(engine, dsl::leaf(pred) - dsl::leaf(target), pred.rows(), + pred.cols(), p_.stable); NN_TRY_CHECK(diff); // Σ diff²：逐元素链与列归约融合为**单次** dispatch（GPU 上 1 个融合 // kernel，不物化 diff_sq 中间张量；CPU 该路径更慢，但 GPU 是主战场）。 // 必须在下面"就地缩放 diff"之前完成（缩放与 diff 共享缓冲）。 - auto col_sum = dsl::compute_reduce(engine, - dsl::col_reduce_sum(dsl::leaf(*diff) * dsl::leaf(*diff)), - diff->rows(), diff->cols(), p_.stable); + auto col_sum = + dsl::compute_reduce(engine, dsl::col_reduce_sum(dsl::leaf(*diff) * dsl::leaf(*diff)), + diff->rows(), diff->cols(), p_.stable); NN_TRY_CHECK(col_sum); // 归约步用 dsl::compute_reduce（col_sum (1,B) → 行归约 (1,1)） - auto total_t = dsl::compute_reduce(engine, - dsl::row_reduce_sum(dsl::leaf(*col_sum)), - col_sum->rows(), col_sum->cols(), p_.stable); + auto total_t = dsl::compute_reduce(engine, dsl::row_reduce_sum(dsl::leaf(*col_sum)), + col_sum->rows(), col_sum->cols(), p_.stable); NN_TRY_CHECK(total_t); auto total_v = detail::download_vector(engine, *total_t); // 宿主桥（17 §3 D11） NN_TRY_CHECK(total_v); @@ -121,18 +118,14 @@ class MSELoss final : public Loss // 结构跨形状共享 AOT shader）。**就地**缩放在 diff 自己的缓冲上完成 // （零额外分配；guard 已用完 diff²，且 grad_input_ 与 diff 共享缓冲 // ——但此时 diff² 已取用，故不产生"平方被缩放值"的污染）。 - auto grad = dsl::compute_into(engine, - dsl::leaf(*diff) * dsl::rparam(scale), *diff); + auto grad = dsl::compute_into(engine, dsl::leaf(*diff) * dsl::rparam(scale), *diff); NN_TRY_CHECK(grad); grad_input_ = std::move(*diff); return loss; } - [[nodiscard]] Result backward() override - { - return grad_input_; - } + [[nodiscard]] Result backward() override { return grad_input_; } }; // ══════════════════════════════════════════════════════════════════════════ @@ -170,41 +163,37 @@ class CrossEntropyLoss final : public Loss // 直接 log(softmax) 在极负 logits/大词表下 softmax→0 → log→-inf， // 与 target=0 相乘得 0*(-inf)=NaN；稳定形式中 shifted 与 log(col_sum) // 均有限，可避免该 NaN（稠密/软标签路径，稀疏 kernel 已用稳定形式）。 - [[nodiscard]] Result softmax_cols_( - ComputeEngine& engine, const Tensor& logits) const + [[nodiscard]] Result softmax_cols_(ComputeEngine& engine, + const Tensor& logits) const { // col_max 用 dsl::compute_reduce（归约向量 (1,cols)，供下面 col_broadcast） - auto col_max = dsl::compute_reduce(engine, - dsl::col_reduce_max(dsl::leaf(logits)), - logits.rows(), logits.cols(), p_.stable); + auto col_max = dsl::compute_reduce(engine, dsl::col_reduce_max(dsl::leaf(logits)), + logits.rows(), logits.cols(), p_.stable); NN_TRY_CHECK(col_max); // denom = Σ_r exp(logits[r][c] - col_max[c]) 用 IR 表达式 // （列归约 + ColBroadcast 视图 + exp，单 kernel，不物化 exp 张量） - auto col_sum = dsl::compute_reduce(engine, - dsl::col_reduce_sum( - dsl::exp(dsl::leaf(logits) - dsl::col_broadcast(*col_max))), + auto col_sum = dsl::compute_reduce( + engine, dsl::col_reduce_sum(dsl::exp(dsl::leaf(logits) - dsl::col_broadcast(*col_max))), logits.rows(), logits.cols(), p_.stable); NN_TRY_CHECK(col_sum); // shifted / softmax / log_softmax 全部用**列广播表达式**：GPU 上各为 1 个 // 融合 kernel，不产生 (classes,batch) 的中间拷贝（列广播是表达式内的一次 // 读取，不额外分配、不额外 dispatch）。 - auto shifted = dsl::compute(engine, - dsl::leaf(logits) - dsl::col_broadcast(*col_max), - logits.rows(), logits.cols(), p_.stable); + auto shifted = dsl::compute(engine, dsl::leaf(logits) - dsl::col_broadcast(*col_max), + logits.rows(), logits.cols(), p_.stable); NN_TRY_CHECK(shifted); - auto softmax = dsl::compute(engine, - dsl::exp(dsl::leaf(*shifted)) / dsl::col_broadcast(*col_sum), - shifted->rows(), shifted->cols(), p_.stable); + auto softmax = + dsl::compute(engine, dsl::exp(dsl::leaf(*shifted)) / dsl::col_broadcast(*col_sum), + shifted->rows(), shifted->cols(), p_.stable); NN_TRY_CHECK(softmax); // 稳定 log_softmax = shifted - log(col_sum) // col_sum ≥ 1（因 max 元素 shifted=0 → exp=1），故 log(col_sum) 有限 - auto log_col_sum = dsl::compute(engine, - dsl::log(dsl::leaf(*col_sum)), - col_sum->rows(), col_sum->cols(), p_.stable); + auto log_col_sum = dsl::compute(engine, dsl::log(dsl::leaf(*col_sum)), col_sum->rows(), + col_sum->cols(), p_.stable); NN_TRY_CHECK(log_col_sum); - auto log_softmax = dsl::compute(engine, - dsl::leaf(*shifted) - dsl::col_broadcast(*log_col_sum), - shifted->rows(), shifted->cols(), p_.stable); + auto log_softmax = + dsl::compute(engine, dsl::leaf(*shifted) - dsl::col_broadcast(*log_col_sum), + shifted->rows(), shifted->cols(), p_.stable); NN_TRY_CHECK(log_softmax); return DenseSoftmax{/*softmax=*/std::move(*softmax), /*log_softmax=*/std::move(*log_softmax)}; @@ -213,16 +202,15 @@ class CrossEntropyLoss final : public Loss public: CrossEntropyLoss() = default; - [[nodiscard]] Result forward( - ComputeEngine& engine, const Tensor& logits, const Tensor& target) override + [[nodiscard]] Result forward(ComputeEngine& engine, const Tensor& logits, + const Tensor& target) override { const std::size_t classes = logits.rows(); - const std::size_t batch = logits.cols(); + const std::size_t batch = logits.cols(); if (target.rows() != classes || target.cols() != batch) NN_FAIL("cross_entropy loss: shape mismatch"); - if (classes == 0 || batch == 0) - NN_FAIL("cross_entropy loss: empty input"); + if (classes == 0 || batch == 0) NN_FAIL("cross_entropy loss: empty input"); // 1. softmax / log_softmax = softmax_cols(logits)（稳定形式） auto sm = softmax_cols_(engine, logits); NN_TRY_CHECK(sm); @@ -234,9 +222,9 @@ class CrossEntropyLoss final : public Loss // 1/batch 由 RParam 承载（值不进 expr_spec_key）→ 与整个逐元素链 // 融合为单 kernel。 auto grad = dsl::compute(engine, - (dsl::leaf(sm->softmax) - dsl::leaf(target)) - * dsl::rparam(Scalar{1} / static_cast(batch)), - sm->softmax.rows(), sm->softmax.cols(), p_.stable); + (dsl::leaf(sm->softmax) - dsl::leaf(target)) * + dsl::rparam(Scalar{1} / static_cast(batch)), + sm->softmax.rows(), sm->softmax.cols(), p_.stable); NN_TRY_CHECK(grad); grad_input_ = std::move(*grad); @@ -251,14 +239,13 @@ class CrossEntropyLoss final : public Loss // 「loss 归约输出按 P 落盘」的真实后果；14 号台账的「4.0137 冻结」 // 即该溢出的另一种硬件表现）。归约内部本就恒 f32 累加，输出改 F32 // 即得正确值，且调用方随后本就 cast 到 F32 回读。 - auto col_s = dsl::compute_reduce(engine, - dsl::col_reduce_sum(dsl::leaf(target) * dsl::leaf(sm->log_softmax)), + auto col_s = dsl::compute_reduce( + engine, dsl::col_reduce_sum(dsl::leaf(target) * dsl::leaf(sm->log_softmax)), target.rows(), target.cols(), p_.stable); NN_TRY_CHECK(col_s); // 归约步用 dsl::compute_reduce（(1,B) → (1,1)） - auto total_t = dsl::compute_reduce(engine, - dsl::row_reduce_sum(dsl::leaf(*col_s)), - col_s->rows(), col_s->cols(), Precision::F32); + auto total_t = dsl::compute_reduce(engine, dsl::row_reduce_sum(dsl::leaf(*col_s)), + col_s->rows(), col_s->cols(), Precision::F32); NN_TRY_CHECK(total_t); // 6. loss = -total / batch — 下载标量 auto total_v = detail::download_vector(engine, *total_t); // 宿主桥（17 §3 D11） @@ -266,10 +253,7 @@ class CrossEntropyLoss final : public Loss return -(*total_v)[0] / static_cast(batch); } - [[nodiscard]] Result backward() override - { - return grad_input_; - } + [[nodiscard]] Result backward() override { return grad_input_; } // ── 稀疏标签版 forward：接受整数标签而非 one-hot 矩阵 ────────── // @@ -295,45 +279,38 @@ class CrossEntropyLoss final : public Loss // num_valid_out 回传有效 token 数；loss = -sum / num_valid。 // 训练热循环用本接口 + engine.submit_scalar_readback 异步取 loss， // 避免每 step to_matrix 触发 end_batch + wait_in_flight 全流水线 drain。 - [[nodiscard]] Result forward_sparse_sum( - ComputeEngine& engine, const Tensor& logits, - std::span labels, - std::span loss_mask, - std::size_t vocab_size, - std::size_t& num_valid_out, - Tensor* grad_reuse = nullptr) + [[nodiscard]] Result forward_sparse_sum(ComputeEngine& engine, const Tensor& logits, + std::span labels, + std::span loss_mask, + std::size_t vocab_size, + std::size_t& num_valid_out, + Tensor* grad_reuse = nullptr) { const std::size_t classes = logits.rows(); - const std::size_t total = logits.cols(); + const std::size_t total = logits.cols(); if (vocab_size == 0) vocab_size = classes; - if (labels.size() != total) - NN_FAIL("sparse CE: labels size mismatch"); + if (labels.size() != total) NN_FAIL("sparse CE: labels size mismatch"); if (!loss_mask.empty() && loss_mask.size() != total) NN_FAIL("sparse CE: mask size mismatch"); - if (classes == 0 || total == 0) - NN_FAIL("sparse CE: empty input"); + if (classes == 0 || total == 0) NN_FAIL("sparse CE: empty input"); // ── 融合路径（不物化全 softmax；失败直接透传错误，不降级） ── - return fused_forward_sparse_(engine, logits, labels, loss_mask, vocab_size, - num_valid_out, grad_reuse); + return fused_forward_sparse_(engine, logits, labels, loss_mask, vocab_size, num_valid_out, + grad_reuse); } // 同步版稀疏 CE（测试/评估等非热路径）：内部下载标量（宿主桥，17 §3 D11）。 - [[nodiscard]] Result forward_sparse( - ComputeEngine& engine, const Tensor& logits, - std::span labels, - std::span loss_mask = {}, - std::size_t vocab_size = 0) + [[nodiscard]] Result forward_sparse(ComputeEngine& engine, const Tensor& logits, + std::span labels, + std::span loss_mask = {}, + std::size_t vocab_size = 0) { std::size_t num_valid = 0; - auto sum_t = forward_sparse_sum(engine, logits, labels, loss_mask, - vocab_size, num_valid); + auto sum_t = forward_sparse_sum(engine, logits, labels, loss_mask, vocab_size, num_valid); NN_TRY_CHECK(sum_t); auto sum_v = detail::download_vector(engine, *sum_t); NN_TRY_CHECK(sum_v); - return (num_valid > 0) - ? -(*sum_v)[0] / static_cast(num_valid) - : Scalar{0}; + return (num_valid > 0) ? -(*sum_v)[0] / static_cast(num_valid) : Scalar{0}; } // ── 融合路径实现（online-softmax 单趟 fold + 两条 IR 链）──────────── @@ -345,13 +322,12 @@ class CrossEntropyLoss final : public Loss // grad = (exp(logits - cb(lse)) - select(Row==cb(labels),1,0)) * cb(mask) * inv // exp(x−s) ≡ exp(x−m)/l（同一数学值；online 归约改变浮点结合序，容差口径） // 派发序：lse → loss_sum → grad（grad_reuse 原地覆写 logits，前两者须先读完） - [[nodiscard]] Result fused_forward_sparse_( - ComputeEngine& engine, const Tensor& logits, - std::span labels, - std::span loss_mask, - std::size_t vocab_size, - std::size_t& num_valid_out, - Tensor* grad_reuse = nullptr) + [[nodiscard]] Result fused_forward_sparse_(ComputeEngine& engine, const Tensor& logits, + std::span labels, + std::span loss_mask, + std::size_t vocab_size, + std::size_t& num_valid_out, + Tensor* grad_reuse = nullptr) { const std::size_t classes = logits.rows(); const std::size_t total = logits.cols(); @@ -361,10 +337,8 @@ class CrossEntropyLoss final : public Loss // 宿主桥（17 §3 D11）：L2 不出现 Matrix，标量缓冲直传引擎。 std::vector labels_v(total); for (std::size_t i = 0; i < total; ++i) - labels_v[i] = static_cast( - (labels[i] < vocab_size) ? labels[i] : 0); - auto labels_t = detail::upload_span(engine, 1, total, Precision::F32, - std::span(labels_v)); + labels_v[i] = static_cast((labels[i] < vocab_size) ? labels[i] : 0); + auto labels_t = detail::upload_span(engine, 1, total, Precision::F32, std::span(labels_v)); NN_TRY_CHECK(labels_t); // 2. 有效 mask（总是构造，(1,total) 0/1；含 mask 缺失与 label 越界修正）： // valid = (loss_mask 空 || mask[i]>=0.5) && labels[i] < vocab_size @@ -374,15 +348,14 @@ class CrossEntropyLoss final : public Loss const bool masked = !loss_mask.empty() && loss_mask[i] < Scalar{0.5}; mask_v[i] = (!masked && labels[i] < vocab_size) ? Scalar{1} : Scalar{0}; } - auto mask_t = detail::upload_span(engine, 1, total, Precision::F32, - std::span(mask_v)); + auto mask_t = detail::upload_span(engine, 1, total, Precision::F32, std::span(mask_v)); NN_TRY_CHECK(mask_t); // 3. num_valid（与 mask 判定一致） std::size_t num_valid = 0; for (std::size_t i = 0; i < total; ++i) if (mask_v[i] >= Scalar{0.5}) ++num_valid; - const Scalar inv_num_valid = (num_valid > 0) - ? Scalar{1} / static_cast(num_valid) : Scalar{0}; + const Scalar inv_num_valid = + (num_valid > 0) ? Scalar{1} / static_cast(num_valid) : Scalar{0}; // 4. lse = logsumexp（online 双状态单趟 fold）：m/l 一并归约、单值 // s = m + log(l) 输出——省掉 col_max 与 denom 两趟 logits 扫描。 @@ -394,8 +367,8 @@ class CrossEntropyLoss final : public Loss Tensor lse_rs; { NN_PROF_SCOPE("ce.fold"); - auto lse = engine.eval_expr(lse_spec, std::span{&logits, 1}, - total, 1, p_.stable); + auto lse = engine.eval_expr(lse_spec, std::span{&logits, 1}, total, 1, + p_.stable); NN_TRY_CHECK(lse); // (total,1) → (1,total)：reshape 零拷贝（GPU 元数据重贴；CPU 小拷贝）， // 供下方 loss/grad 两条链按列广播（ColBroadcast 读 b[c]）取用 @@ -419,10 +392,10 @@ class CrossEntropyLoss final : public Loss { NN_PROF_SCOPE("ce.lossum"); auto tr = dsl::compute_reduce(engine, - dsl::row_reduce_sum( - (dsl::row_gather(logits, *labels_t) - dsl::col_broadcast(lse_rs)) - * dsl::col_broadcast(*mask_t)), - 1, total, Precision::F32); + dsl::row_reduce_sum((dsl::row_gather(logits, *labels_t) - + dsl::col_broadcast(lse_rs)) * + dsl::col_broadcast(*mask_t)), + 1, total, Precision::F32); NN_TRY_CHECK(tr); total_t = std::move(*tr); } @@ -431,11 +404,9 @@ class CrossEntropyLoss final : public Loss // 1/num_valid 由 RParam 承载（运行时值、不进 expr_spec_key）→ 与整个 // 逐元素链融合为单 kernel auto grad_expr = - (dsl::exp(dsl::leaf(logits) - dsl::col_broadcast(lse_rs)) - - dsl::select(dsl::row() == dsl::col_broadcast(*labels_t), - Scalar{1}, Scalar{0})) - * dsl::col_broadcast(*mask_t) - * dsl::rparam(inv_num_valid); + (dsl::exp(dsl::leaf(logits) - dsl::col_broadcast(lse_rs)) - + dsl::select(dsl::row() == dsl::col_broadcast(*labels_t), Scalar{1}, Scalar{0})) * + dsl::col_broadcast(*mask_t) * dsl::rparam(inv_num_valid); // 显存优化：grad_reuse 与 logits 同形同精度时把梯度**原地**写进该 // 缓冲（典型用法：传 &logits）——省一份 (vocab×total) 显存，训练峰值 @@ -445,9 +416,8 @@ class CrossEntropyLoss final : public Loss // 中，且 loss_sum 已于本派发之前算完（步骤 5）。 { NN_PROF_SCOPE("ce.grad"); - if (grad_reuse != nullptr && grad_reuse->valid() - && grad_reuse->rows() == classes && grad_reuse->cols() == total - && grad_reuse->precision() == logits.precision()) + if (grad_reuse != nullptr && grad_reuse->valid() && grad_reuse->rows() == classes && + grad_reuse->cols() == total && grad_reuse->precision() == logits.precision()) { auto w = dsl::compute_into(engine, grad_expr, *grad_reuse); NN_TRY_CHECK(w); @@ -470,4 +440,3 @@ class CrossEntropyLoss final : public Loss }; } // namespace nn - diff --git a/include/neuralnet.cpp/compute_optimizer.hpp b/include/neuralnet.cpp/compute_optimizer.hpp index dba1289a..27082980 100644 --- a/include/neuralnet.cpp/compute_optimizer.hpp +++ b/include/neuralnet.cpp/compute_optimizer.hpp @@ -62,8 +62,7 @@ class Optimizer // 校验 params/grads 数量一致（各子类 step() 开头调用） [[nodiscard]] Result validate_sizes_() const { - if (params_.size() != grads_.size()) - NN_FAIL("Optimizer: params/grads size mismatch"); + if (params_.size() != grads_.size()) NN_FAIL("Optimizer: params/grads size mismatch"); return {}; } @@ -75,8 +74,7 @@ class Optimizer buffers.reserve(params_.size()); for (auto& p : params_) { - auto buf = engine_.create_tensor(p.get().rows(), p.get().cols(), - p_.optimizer); + auto buf = engine_.create_tensor(p.get().rows(), p.get().cols(), p_.optimizer); auto r = engine_.zero(buf); NN_TRY_CHECK(r); buffers.push_back(std::move(buf)); @@ -100,8 +98,8 @@ class Optimizer return {}; } - [[nodiscard]] static ComputeEngine& bind_engine_( - ComputeEngine& engine, std::shared_ptr& owner) + [[nodiscard]] static ComputeEngine& bind_engine_(ComputeEngine& engine, + std::shared_ptr& owner) { owner = engine.weak_from_this().lock(); if (!owner) @@ -111,15 +109,12 @@ class Optimizer } public: - Optimizer(ComputeEngine& engine, - std::vector params, - std::vector grads, + Optimizer(ComputeEngine& engine, std::vector params, std::vector grads, PrecisionProfile precision = PrecisionProfile{}) - : engine_owner_(), - engine_(bind_engine_(engine, engine_owner_)), - params_(std::move(params)), - grads_(std::move(grads)), - p_(precision) {} + : engine_owner_(), engine_(bind_engine_(engine, engine_owner_)), params_(std::move(params)), + grads_(std::move(grads)), p_(precision) + { + } virtual ~Optimizer() = default; @@ -137,8 +132,7 @@ class Optimizer // 在 step() 之前调用： backward() → clip_grad_norm(max_norm) → step(). [[nodiscard]] Result clip_grad_norm(Scalar max_norm) { - if (max_norm <= 0 || grads_.empty()) - return {}; + if (max_norm <= 0 || grads_.empty()) return {}; // 计算全局梯度平方和：每个梯度张量归约到 (1,1) 标量后在设备端累加， // 最后仅下载一次（而非每张量一次 PCIe 往返）。 @@ -147,18 +141,16 @@ class Optimizer { auto& g = g_ref.get(); // g² = g * g（逐元素乘法） - auto g_sq_r = dsl::compute(engine_, dsl::leaf(g) * dsl::leaf(g), - g.rows(), g.cols()); + auto g_sq_r = dsl::compute(engine_, dsl::leaf(g) * dsl::leaf(g), g.rows(), g.cols()); NN_TRY_CHECK(g_sq_r); // 按行求和 → (rows, 1)（dsl::compute_reduce） - auto row_sums_r = dsl::compute_reduce(engine_, - dsl::row_reduce_sum(dsl::leaf(*g_sq_r)), - g_sq_r->rows(), g_sq_r->cols()); + auto row_sums_r = dsl::compute_reduce(engine_, dsl::row_reduce_sum(dsl::leaf(*g_sq_r)), + g_sq_r->rows(), g_sq_r->cols()); NN_TRY_CHECK(row_sums_r); // 按列求和 → (1, 1) 标量张量（dsl::compute_reduce） - auto col_sum_r = dsl::compute_reduce(engine_, - dsl::col_reduce_sum(dsl::leaf(*row_sums_r)), - row_sums_r->rows(), row_sums_r->cols()); + auto col_sum_r = + dsl::compute_reduce(engine_, dsl::col_reduce_sum(dsl::leaf(*row_sums_r)), + row_sums_r->rows(), row_sums_r->cols()); NN_TRY_CHECK(col_sum_r); if (acc.valid()) { @@ -180,15 +172,14 @@ class Optimizer const Scalar total_sq = total_buf[0]; Scalar norm = std::sqrt(total_sq); - if (norm <= max_norm) - return {}; + if (norm <= max_norm) return {}; // 等比例缩放所有梯度（目标传递：原地、单 dispatch） Scalar scale = max_norm / norm; for (auto& g_ref : grads_) { - auto r = dsl::compute_into(engine_, - dsl::leaf(g_ref.get()) * dsl::rparam(scale), g_ref.get()); + auto r = dsl::compute_into(engine_, dsl::leaf(g_ref.get()) * dsl::rparam(scale), + g_ref.get()); NN_TRY_CHECK(r); } return {}; @@ -227,8 +218,7 @@ class Optimizer [[nodiscard]] virtual Result load_state(std::span in) { if (!in.empty()) - NN_FAIL("该优化器没有状态，但文件里的状态张量数 = " + - std::to_string(in.size())); + NN_FAIL("该优化器没有状态，但文件里的状态张量数 = " + std::to_string(in.size())); return {}; } [[nodiscard]] virtual std::uint64_t state_step() const noexcept { return 0; } @@ -236,16 +226,14 @@ class Optimizer protected: // 状态张量回写（`load_state` 的实现件）：形状必须一致；精度不同经引擎 cast。 - [[nodiscard]] Result restore_state_tensor_( - Tensor& dst, const Tensor& src, const char* what) + [[nodiscard]] Result restore_state_tensor_(Tensor& dst, const Tensor& src, + const char* what) { if (dst.rows() != src.rows() || dst.cols() != src.cols()) NN_FAIL(std::string("优化器状态形状不匹配（") + what + "）：期望 " + - std::to_string(dst.rows()) + "x" + std::to_string(dst.cols()) + - "，文件里是 " + std::to_string(src.rows()) + "x" + - std::to_string(src.cols())); - if (dst.precision() == src.precision()) - return engine_.copy_into(dst, src); + std::to_string(dst.rows()) + "x" + std::to_string(dst.cols()) + "，文件里是 " + + std::to_string(src.rows()) + "x" + std::to_string(src.cols())); + if (dst.precision() == src.precision()) return engine_.copy_into(dst, src); return engine_.cast_into(src, dst); } }; @@ -261,12 +249,11 @@ class SGD : public Optimizer Scalar lr_; public: - SGD(ComputeEngine& engine, - std::vector params, - std::vector grads, - Scalar lr, - PrecisionProfile precision = PrecisionProfile{}) - : Optimizer(engine, std::move(params), std::move(grads), precision), lr_(lr) {} + SGD(ComputeEngine& engine, std::vector params, std::vector grads, + Scalar lr, PrecisionProfile precision = PrecisionProfile{}) + : Optimizer(engine, std::move(params), std::move(grads), precision), lr_(lr) + { + } void set_lr(Scalar lr) override { lr_ = lr; } @@ -276,8 +263,8 @@ class SGD : public Optimizer for (std::size_t i = 0; i < params_.size(); ++i) { // p -= lr * g（目标传递：原地、单 dispatch） - auto r = dsl::compute_into(engine_, - dsl::leaf(params_[i]) + dsl::leaf(grads_[i]) * dsl::rparam(-lr_), + auto r = dsl::compute_into( + engine_, dsl::leaf(params_[i]) + dsl::leaf(grads_[i]) * dsl::rparam(-lr_), params_[i]); NN_TRY_CHECK(r); } @@ -298,19 +285,14 @@ class SGDWithMomentum : public Optimizer std::vector velocities_; public: - SGDWithMomentum(ComputeEngine& engine, - std::vector params, - std::vector grads, - Scalar lr, Scalar beta = 0.9, - PrecisionProfile precision = PrecisionProfile{}) - : Optimizer(engine, std::move(params), std::move(grads), precision), - lr_(lr), beta_(beta) + SGDWithMomentum(ComputeEngine& engine, std::vector params, + std::vector grads, Scalar lr, Scalar beta = 0.9, + PrecisionProfile precision = PrecisionProfile{}) + : Optimizer(engine, std::move(params), std::move(grads), precision), lr_(lr), beta_(beta) { auto v_r = create_zero_buffers_(); - if (!v_r) - record_init_error_(v_r.error()); // 不 abort：由 step() 经 Result 上抛 - else - velocities_ = std::move(*v_r); + if (!v_r) record_init_error_(v_r.error()); // 不 abort：由 step() 经 Result 上抛 + else velocities_ = std::move(*v_r); } void set_lr(Scalar lr) override { lr_ = lr; } @@ -325,13 +307,13 @@ class SGDWithMomentum : public Optimizer { // v = β*v + (1-β)*g（两趟原语融合为单 dispatch 的原地目标传递） auto r = dsl::compute_into(engine_, - dsl::leaf(velocities_[i]) * dsl::rparam(beta_) - + dsl::leaf(grads_[i]) * dsl::rparam(one_minus_beta), - velocities_[i]); + dsl::leaf(velocities_[i]) * dsl::rparam(beta_) + + dsl::leaf(grads_[i]) * dsl::rparam(one_minus_beta), + velocities_[i]); NN_TRY_CHECK(r); // p -= lr * v（原地、单 dispatch） - r = dsl::compute_into(engine_, - dsl::leaf(params_[i]) + dsl::leaf(velocities_[i]) * dsl::rparam(-lr_), + r = dsl::compute_into( + engine_, dsl::leaf(params_[i]) + dsl::leaf(velocities_[i]) * dsl::rparam(-lr_), params_[i]); NN_TRY_CHECK(r); } @@ -342,16 +324,16 @@ class SGDWithMomentum : public Optimizer [[nodiscard]] Result save_state(std::vector& out) const override { NN_TRY(r, check_ready_()); - for (const auto& v : velocities_) out.push_back(v); + for (const auto& v : velocities_) + out.push_back(v); return {}; } [[nodiscard]] Result load_state(std::span in) override { NN_TRY(r, check_ready_()); if (in.size() != velocities_.size()) - NN_FAIL("SGD + Momentum 状态张量数不匹配：期望 " + - std::to_string(velocities_.size()) + "（每参数 1 个），实际 " + - std::to_string(in.size())); + NN_FAIL("SGD + Momentum 状态张量数不匹配：期望 " + std::to_string(velocities_.size()) + + "（每参数 1 个），实际 " + std::to_string(in.size())); for (std::size_t i = 0; i < in.size(); ++i) NN_TRY_CHECK(restore_state_tensor_(velocities_[i], in[i], "momentum")); return {}; @@ -386,9 +368,7 @@ class Adam : public Optimizer // （docs/internal/24 §2）；实践上配 `--max-norm` 梯度裁剪即可约束 |g|。 // · S 取 2 的幂 → 与 g 的相乘在 f16 下**无损**（指数平移），不引入额外舍入。 [[nodiscard]] Scalar optimizer_state_scale_() const noexcept - { - return precision_bytes(p_.optimizer) >= 4 ? Scalar{1} : Scalar{4096}; - } + { return precision_bytes(p_.optimizer) >= 4 ? Scalar{1} : Scalar{4096}; } // Adam 核心更新（protected，AdamW 复用） // inv_bc1/inv_bc2 由 step() 提前计算（每步仅一次 pow），失败时不推进 t。 @@ -400,8 +380,7 @@ class Adam : public Optimizer // K1 m = β1*m + (1-β1)*g （m_ owned，直接重赋值） // K2 v = β2*v + (1-β2)*g² （v_ owned，直接重赋值） // K3 p -= lr * (inv_bc1*m) / (sqrt(inv_bc2*v)+eps)（p 为模型张量，目标传递原地更新） - [[nodiscard]] Result adam_update_( - std::size_t i, Scalar inv_bc1, Scalar inv_bc2) + [[nodiscard]] Result adam_update_(std::size_t i, Scalar inv_bc1, Scalar inv_bc2) { const Scalar one_minus_beta1 = Scalar{1} - beta1_; const Scalar one_minus_beta2 = Scalar{1} - beta2_; @@ -426,17 +405,18 @@ class Adam : public Optimizer // K1: m = β1*m + (1-β1)*g（单 kernel 融合；状态精度 = p_.optimizer） auto m_new = dsl::compute(engine_, - dsl::leaf(m_[i]) * dsl::rparam(beta1_) + - dsl::leaf(g) * dsl::rparam(one_minus_beta1 * s_scale), - rows, cols, p_.optimizer); + dsl::leaf(m_[i]) * dsl::rparam(beta1_) + + dsl::leaf(g) * dsl::rparam(one_minus_beta1 * s_scale), + rows, cols, p_.optimizer); NN_TRY_CHECK(m_new); m_[i] = std::move(*m_new); // K2: v = β2*v + (1-β2)*g²（单 kernel 融合；状态精度 = p_.optimizer） - auto v_new = dsl::compute(engine_, - dsl::leaf(v_[i]) * dsl::rparam(beta2_) + - dsl::leaf(g) * dsl::leaf(g) * dsl::rparam(one_minus_beta2 * s2), - rows, cols, p_.optimizer); + auto v_new = + dsl::compute(engine_, + dsl::leaf(v_[i]) * dsl::rparam(beta2_) + + dsl::leaf(g) * dsl::leaf(g) * dsl::rparam(one_minus_beta2 * s2), + rows, cols, p_.optimizer); NN_TRY_CHECK(v_new); v_[i] = std::move(*v_new); @@ -448,19 +428,17 @@ class Adam : public Optimizer // T1_58（latent 恒 f32）：照搬 p_.param 会拿非存储精度去建张量 // （CPU 走 f32 兜底、GPU 在 check_precision_supported 处硬报错 = // 两端行为不一致）。见 docs/internal/21-quantized-weights.md §4.3.2。 - const Precision upd_prec = (precision_bytes(p_.param) != 0) - ? p_.param - : params_[i].get().precision(); - auto delta = dsl::compute(engine_, - -dsl::rparam(lr_) - * ((dsl::leaf(m_[i]) * dsl::rparam(inv_bc1)) - / (dsl::sqrt(dsl::leaf(v_[i]) * dsl::rparam(inv_bc2)) - + dsl::rparam(eps_scaled))), + const Precision upd_prec = + (precision_bytes(p_.param) != 0) ? p_.param : params_[i].get().precision(); + auto delta = dsl::compute( + engine_, + -dsl::rparam(lr_) * + ((dsl::leaf(m_[i]) * dsl::rparam(inv_bc1)) / + (dsl::sqrt(dsl::leaf(v_[i]) * dsl::rparam(inv_bc2)) + dsl::rparam(eps_scaled))), rows, cols, upd_prec); NN_TRY_CHECK(delta); // p += delta（目标传递：原地、单 dispatch，不额外分配） - auto r = dsl::compute_into(engine_, - dsl::leaf(params_[i]) + dsl::leaf(*delta), params_[i]); + auto r = dsl::compute_into(engine_, dsl::leaf(params_[i]) + dsl::leaf(*delta), params_[i]); NN_TRY_CHECK(r); return {}; } @@ -493,19 +471,12 @@ class Adam : public Optimizer } public: - Adam(ComputeEngine& engine, - std::vector params, - std::vector grads, - Scalar lr, - Scalar beta1 = 0.9, - Scalar beta2 = 0.999, - Scalar eps = 1e-8, + Adam(ComputeEngine& engine, std::vector params, std::vector grads, + Scalar lr, Scalar beta1 = 0.9, Scalar beta2 = 0.999, Scalar eps = 1e-8, PrecisionProfile precision = PrecisionProfile{}) - : Optimizer(engine, std::move(params), std::move(grads), precision), - lr_(lr), beta1_(beta1), beta2_(beta2), eps_(eps), t_(0) - { - init_moments_(); - } + : Optimizer(engine, std::move(params), std::move(grads), precision), lr_(lr), beta1_(beta1), + beta2_(beta2), eps_(eps), t_(0) + { init_moments_(); } void set_lr(Scalar lr) override { lr_ = lr; } @@ -559,13 +530,9 @@ class Adam : public Optimizer return {}; } [[nodiscard]] std::uint64_t state_step() const noexcept override - { - return static_cast(t_); - } + { return static_cast(t_); } void set_state_step(std::uint64_t step) noexcept override - { - t_ = static_cast(step); - } + { t_ = static_cast(step); } }; // ══════════════════════════════════════════════════════════════════════════ @@ -587,18 +554,13 @@ class AdamW : public Adam Scalar wd_; // 权重衰减系数 public: - AdamW(ComputeEngine& engine, - std::vector params, - std::vector grads, - Scalar lr, - Scalar beta1 = 0.9, - Scalar beta2 = 0.999, - Scalar eps = 1e-8, - Scalar weight_decay = 0.01, - PrecisionProfile precision = PrecisionProfile{}) - : Adam(engine, std::move(params), std::move(grads), - lr, beta1, beta2, eps, precision), - wd_(weight_decay) {} + AdamW(ComputeEngine& engine, std::vector params, std::vector grads, + Scalar lr, Scalar beta1 = 0.9, Scalar beta2 = 0.999, Scalar eps = 1e-8, + Scalar weight_decay = 0.01, PrecisionProfile precision = PrecisionProfile{}) + : Adam(engine, std::move(params), std::move(grads), lr, beta1, beta2, eps, precision), + wd_(weight_decay) + { + } // set_lr 复用 Adam::set_lr（lr_ 为 protected，无需 override） @@ -623,8 +585,9 @@ class AdamW : public Adam // 权重衰减解耦：p = (1 - lr*wd) * p（目标传递：原地、单 dispatch） if (wd_ != Scalar{0}) { - NN_TRY(dr, dsl::compute_into(engine_, - dsl::leaf(params_[i]) * dsl::rparam(decay_factor), params_[i])); + NN_TRY(dr, + dsl::compute_into(engine_, dsl::leaf(params_[i]) * dsl::rparam(decay_factor), + params_[i])); } // Adam 更新 @@ -722,34 +685,36 @@ class AdamWQuant : public Optimizer const Scalar one_m_b2 = Scalar{1} - beta2_; // ── 尺度前进（上界；只依赖 g，不需要新状态的 f32 副本）────────────── - NN_TRY(gabs, dsl::compute_reduce(engine_, - dsl::row_reduce_max(dsl::abs(dsl::leaf(g))), rows, cols, Precision::F32)); - NN_TRY(gmax, dsl::compute_reduce(engine_, - dsl::row_reduce_max(dsl::leaf(g) * dsl::leaf(g)), rows, cols, Precision::F32)); + NN_TRY(gabs, dsl::compute_reduce(engine_, dsl::row_reduce_max(dsl::abs(dsl::leaf(g))), rows, + cols, Precision::F32)); + NN_TRY(gmax, dsl::compute_reduce(engine_, dsl::row_reduce_max(dsl::leaf(g) * dsl::leaf(g)), + rows, cols, Precision::F32)); auto sm_new = dsl::compute(engine_, - dsl::max(dsl::leaf(sm_[i]) * dsl::rparam(beta1_) - + dsl::leaf(*gabs) * dsl::rparam(one_m_b1 * kSafety), - kScaleFloor), - rows, 1, Precision::F32); + dsl::max(dsl::leaf(sm_[i]) * dsl::rparam(beta1_) + + dsl::leaf(*gabs) * dsl::rparam(one_m_b1 * kSafety), + kScaleFloor), + rows, 1, Precision::F32); NN_TRY_CHECK(sm_new); auto sv_new = dsl::compute(engine_, - dsl::max(dsl::leaf(sv_[i]) * dsl::rparam(beta2_) - + dsl::leaf(*gmax) * dsl::rparam(one_m_b2 * kSafety), - kScaleFloor), - rows, 1, Precision::F32); + dsl::max(dsl::leaf(sv_[i]) * dsl::rparam(beta2_) + + dsl::leaf(*gmax) * dsl::rparam(one_m_b2 * kSafety), + kScaleFloor), + rows, 1, Precision::F32); NN_TRY_CHECK(sv_new); // ── 状态写回（去量化 → 新状态 → ÷新尺度 → f16 RNE 写回，单 dispatch）── - auto mw = dsl::compute_into(engine_, - (dsl::leaf(m16_[i]) * dsl::row_broadcast(sm_[i]) * dsl::rparam(beta1_) - + dsl::leaf(g) * dsl::rparam(one_m_b1)) - / dsl::row_broadcast(*sm_new), + auto mw = dsl::compute_into( + engine_, + (dsl::leaf(m16_[i]) * dsl::row_broadcast(sm_[i]) * dsl::rparam(beta1_) + + dsl::leaf(g) * dsl::rparam(one_m_b1)) / + dsl::row_broadcast(*sm_new), m16_[i]); NN_TRY_CHECK(mw); - auto vw = dsl::compute_into(engine_, - (dsl::leaf(v16_[i]) * dsl::row_broadcast(sv_[i]) * dsl::rparam(beta2_) - + dsl::leaf(g) * dsl::leaf(g) * dsl::rparam(one_m_b2)) - / dsl::row_broadcast(*sv_new), + auto vw = dsl::compute_into( + engine_, + (dsl::leaf(v16_[i]) * dsl::row_broadcast(sv_[i]) * dsl::rparam(beta2_) + + dsl::leaf(g) * dsl::leaf(g) * dsl::rparam(one_m_b2)) / + dsl::row_broadcast(*sv_new), v16_[i]); NN_TRY_CHECK(vw); @@ -758,32 +723,26 @@ class AdamWQuant : public Optimizer sv_[i] = std::move(*sv_new); // ── 参数更新（AdamW：weight decay 解耦到 p；偏置修正走 rparam）────── - auto upd = dsl::compute_into(engine_, - dsl::leaf(params_[i]) - - dsl::rparam(lr_) - * (dsl::rparam(inv_bc1) - * (dsl::leaf(m16_[i]) * dsl::row_broadcast(sm_[i])) - / (dsl::sqrt(dsl::rparam(inv_bc2) - * (dsl::leaf(v16_[i]) * dsl::row_broadcast(sv_[i]))) - + dsl::rparam(eps_)) - + dsl::rparam(wd_) * dsl::leaf(params_[i])), + auto upd = dsl::compute_into( + engine_, + dsl::leaf(params_[i]) - + dsl::rparam(lr_) * + (dsl::rparam(inv_bc1) * (dsl::leaf(m16_[i]) * dsl::row_broadcast(sm_[i])) / + (dsl::sqrt(dsl::rparam(inv_bc2) * + (dsl::leaf(v16_[i]) * dsl::row_broadcast(sv_[i]))) + + dsl::rparam(eps_)) + + dsl::rparam(wd_) * dsl::leaf(params_[i])), params_[i]); NN_TRY_CHECK(upd); return {}; } public: - AdamWQuant(ComputeEngine& engine, - std::vector params, - std::vector grads, - Scalar lr, - Scalar beta1 = 0.9, - Scalar beta2 = 0.999, - Scalar eps = 1e-8, - Scalar weight_decay = 0.01, - PrecisionProfile precision = PrecisionProfile{}) - : Optimizer(engine, std::move(params), std::move(grads), precision), - lr_(lr), beta1_(beta1), beta2_(beta2), eps_(eps), wd_(weight_decay) + AdamWQuant(ComputeEngine& engine, std::vector params, std::vector grads, + Scalar lr, Scalar beta1 = 0.9, Scalar beta2 = 0.999, Scalar eps = 1e-8, + Scalar weight_decay = 0.01, PrecisionProfile precision = PrecisionProfile{}) + : Optimizer(engine, std::move(params), std::move(grads), precision), lr_(lr), beta1_(beta1), + beta2_(beta2), eps_(eps), wd_(weight_decay) { auto r = init_state_(); if (!r) record_init_error_(r.error()); @@ -825,8 +784,7 @@ class AdamWQuant : public Optimizer { NN_TRY(r, check_ready_()); if (in.size() != 4 * m16_.size()) - NN_FAIL("AdamW（量化状态）状态张量数不匹配：期望 " + - std::to_string(4 * m16_.size()) + + NN_FAIL("AdamW（量化状态）状态张量数不匹配：期望 " + std::to_string(4 * m16_.size()) + "（每参数 m16/v16/sm/sv），实际 " + std::to_string(in.size())); for (std::size_t i = 0; i < m16_.size(); ++i) { @@ -838,13 +796,9 @@ class AdamWQuant : public Optimizer return {}; } [[nodiscard]] std::uint64_t state_step() const noexcept override - { - return static_cast(t_); - } + { return static_cast(t_); } void set_state_step(std::uint64_t step) noexcept override - { - t_ = static_cast(step); - } + { t_ = static_cast(step); } // ── 诊断 / 测试视图 ─────────────────────────────────────────────────── // 状态实际字节数（用于与 f32 Adam 的状态账对拍：后者 = Σ 2·rows·cols·4） @@ -871,9 +825,9 @@ class AdamWQuant : public Optimizer // // 参考：Keller Jordan et al., "Muon: An optimizer for hidden layers in neural networks" // ══════════════════════════════════════════════════════════════════════════ -[[nodiscard]] inline Result newton_schulz_orthogonalize( - ComputeEngine& engine, const Tensor& G, std::size_t steps = 5, Scalar eps = 1e-7f, - Precision prec = Precision::F32) +[[nodiscard]] inline Result +newton_schulz_orthogonalize(ComputeEngine& engine, const Tensor& G, std::size_t steps = 5, + Scalar eps = 1e-7f, Precision prec = Precision::F32) { // 调优后的 quintic 多项式系数（使 φ^N(x) → 1 for x ∈ [0,1]） constexpr Scalar a = 3.4445f; @@ -883,13 +837,12 @@ class AdamWQuant : public Optimizer // 计算 Frobenius 范数的平方：||G||_F² = Σ g_ij² // 逐元素链与行归约**融合为单 dispatch**（不物化 norm_sq 中间张量）， // 再对 (rows,1) 小向量做列归约 → (1,1) - auto row_sum_norm = dsl::compute_reduce(engine, - dsl::row_reduce_sum(dsl::leaf(G) * dsl::leaf(G)), G.rows(), G.cols(), prec); + auto row_sum_norm = dsl::compute_reduce( + engine, dsl::row_reduce_sum(dsl::leaf(G) * dsl::leaf(G)), G.rows(), G.cols(), prec); NN_TRY_CHECK(row_sum_norm); // 二级归约（(rows,1) → (1,1)）用 dsl::compute_reduce - auto total_norm_sq = dsl::compute_reduce(engine, - dsl::col_reduce_sum(dsl::leaf(*row_sum_norm)), - row_sum_norm->rows(), row_sum_norm->cols(), prec); + auto total_norm_sq = dsl::compute_reduce(engine, dsl::col_reduce_sum(dsl::leaf(*row_sum_norm)), + row_sum_norm->rows(), row_sum_norm->cols(), prec); NN_TRY_CHECK(total_norm_sq); // 从 (1,1) Tensor 提取标量值 // 使用宿主桥下载后取值（一次 PCIe 下载，可接受；17 §3 D11） @@ -901,8 +854,8 @@ class AdamWQuant : public Optimizer // X = G * inv_norm_scalar（归一化）：单表达式（GPU 上 1 个融合 kernel + // 1 次分配） - auto X = dsl::compute(engine, - dsl::leaf(G) * dsl::rparam(inv_norm_scalar), G.rows(), G.cols(), prec); + auto X = + dsl::compute(engine, dsl::leaf(G) * dsl::rparam(inv_norm_scalar), G.rows(), G.cols(), prec); NN_TRY_CHECK(X); // 选更小一侧构造母矩阵，避免显存爆炸（Muon 显存 > AdamW 的根因）： // - 短宽/方阵（m ≤ n）：行正交化，母矩阵 A = X·X^T（m×m，m 为短边） @@ -924,21 +877,21 @@ class AdamWQuant : public Optimizer { // 纯 matmul 经 DSL 直写（scan 的 optimizer dry-run 自动登记结构） auto A = dsl::compute(engine, // A = X·X^T - dsl::matmul(*X, *X, false, true), m, m, prec); + dsl::matmul(*X, *X, false, true), m, m, prec); NN_TRY_CHECK(A); auto A_sq = dsl::compute(engine, // A² - dsl::matmul(*A, *A, false, false), m, m, prec); + dsl::matmul(*A, *A, false, false), m, m, prec); NN_TRY_CHECK(A_sq); // A = b·A + c·A²：一次原地目标传递（GPU 上 1 个 kernel） - auto accA = dsl::compute_into(engine, - dsl::leaf(*A) * dsl::rparam(b) + dsl::leaf(*A_sq) * dsl::rparam(c), *A); + auto accA = dsl::compute_into( + engine, dsl::leaf(*A) * dsl::rparam(b) + dsl::leaf(*A_sq) * dsl::rparam(c), *A); NN_TRY_CHECK(accA); auto BX = dsl::compute(engine, // B·X - dsl::matmul(*A, *X, false, false), m, n, prec); + dsl::matmul(*A, *X, false, false), m, n, prec); NN_TRY_CHECK(BX); // BX += a·X（原地目标传递：单 dispatch，不额外分配） - auto accBX = dsl::compute_into(engine, - dsl::leaf(*BX) + dsl::leaf(*X) * dsl::rparam(a), *BX); + auto accBX = + dsl::compute_into(engine, dsl::leaf(*BX) + dsl::leaf(*X) * dsl::rparam(a), *BX); NN_TRY_CHECK(accBX); X = std::move(*BX); } @@ -952,21 +905,21 @@ class AdamWQuant : public Optimizer { // 纯 matmul 经 DSL 直写（同上，scan optimizer dry-run 登记） auto Gr = dsl::compute(engine, // G = X^T·X（n×n） - dsl::matmul(*X, *X, true, false), n, n, prec); + dsl::matmul(*X, *X, true, false), n, n, prec); NN_TRY_CHECK(Gr); auto Gr_sq = dsl::compute(engine, // G² - dsl::matmul(*Gr, *Gr, false, false), n, n, prec); + dsl::matmul(*Gr, *Gr, false, false), n, n, prec); NN_TRY_CHECK(Gr_sq); // G = b·G + c·G²（同上：一次原地目标传递，单 kernel） - auto accG = dsl::compute_into(engine, - dsl::leaf(*Gr) * dsl::rparam(b) + dsl::leaf(*Gr_sq) * dsl::rparam(c), *Gr); + auto accG = dsl::compute_into( + engine, dsl::leaf(*Gr) * dsl::rparam(b) + dsl::leaf(*Gr_sq) * dsl::rparam(c), *Gr); NN_TRY_CHECK(accG); auto XM = dsl::compute(engine, // X·G - dsl::matmul(*X, *Gr, false, false), m, n, prec); + dsl::matmul(*X, *Gr, false, false), m, n, prec); NN_TRY_CHECK(XM); // XM += a·X（原地目标传递：单 dispatch） - auto accXM = dsl::compute_into(engine, - dsl::leaf(*XM) + dsl::leaf(*X) * dsl::rparam(a), *XM); + auto accXM = + dsl::compute_into(engine, dsl::leaf(*XM) + dsl::leaf(*X) * dsl::rparam(a), *XM); NN_TRY_CHECK(accXM); X = std::move(*XM); } @@ -1003,24 +956,15 @@ class Muon : public Optimizer std::vector velocities_; public: - Muon(ComputeEngine& engine, - std::vector params, - std::vector grads, - Scalar lr, - Scalar momentum = 0.95f, - bool nesterov = true, - std::size_t ns_steps = 5, - Scalar ns_eps = 1e-7f, - PrecisionProfile precision = PrecisionProfile{}) - : Optimizer(engine, std::move(params), std::move(grads), precision), - lr_(lr), momentum_(momentum), nesterov_(nesterov), - ns_steps_(ns_steps), ns_eps_(ns_eps) + Muon(ComputeEngine& engine, std::vector params, std::vector grads, + Scalar lr, Scalar momentum = 0.95f, bool nesterov = true, std::size_t ns_steps = 5, + Scalar ns_eps = 1e-7f, PrecisionProfile precision = PrecisionProfile{}) + : Optimizer(engine, std::move(params), std::move(grads), precision), lr_(lr), + momentum_(momentum), nesterov_(nesterov), ns_steps_(ns_steps), ns_eps_(ns_eps) { auto v_r = create_zero_buffers_(); - if (!v_r) - record_init_error_(v_r.error()); // 不 abort：由 step() 经 Result 上抛 - else - velocities_ = std::move(*v_r); + if (!v_r) record_init_error_(v_r.error()); // 不 abort：由 step() 经 Result 上抛 + else velocities_ = std::move(*v_r); } void set_lr(Scalar lr) override { lr_ = lr; } @@ -1034,8 +978,8 @@ class Muon : public Optimizer const Tensor& g = grads_[i]; // 1. SGD-Momentum: v = μ*v + g（两趟原语融合为单 dispatch） - auto r = dsl::compute_into(engine_, - dsl::leaf(velocities_[i]) * dsl::rparam(momentum_) + dsl::leaf(g), + auto r = dsl::compute_into( + engine_, dsl::leaf(velocities_[i]) * dsl::rparam(momentum_) + dsl::leaf(g), velocities_[i]); NN_TRY_CHECK(r); // 确定用于正交化的更新方向（Nesterov 时需临时缓冲，否则直接用 v） @@ -1043,8 +987,8 @@ class Muon : public Optimizer if (nesterov_) { // Nesterov: update = g + μ*v（单表达式：一次 dispatch + 一次分配） - auto buf = dsl::compute(engine_, - dsl::leaf(g) + dsl::leaf(velocities_[i]) * dsl::rparam(momentum_), + auto buf = dsl::compute( + engine_, dsl::leaf(g) + dsl::leaf(velocities_[i]) * dsl::rparam(momentum_), g.rows(), g.cols(), p_.optimizer); NN_TRY_CHECK(buf); nesterov_buf = std::move(*buf); @@ -1054,8 +998,8 @@ class Muon : public Optimizer // 2. Newton-Schulz 正交化（仅对 ≥2D 参数，即 rows > 1 且 cols > 1） if (params_[i].get().rows() > 1 && params_[i].get().cols() > 1) { - auto ortho_update = newton_schulz_orthogonalize( - engine_, update, ns_steps_, ns_eps_, p_.optimizer); + auto ortho_update = + newton_schulz_orthogonalize(engine_, update, ns_steps_, ns_eps_, p_.optimizer); NN_TRY_CHECK(ortho_update); // 3. 参数更新: p -= lr * 0.2 * sqrt(max(m,n)) * NS(update) // NorMuon 论文 / KellerJordan 参考实现的形状缩放：NS 输出谱范数为 1， @@ -1063,20 +1007,19 @@ class Muon : public Optimizer const std::size_t m = params_[i].get().rows(); const std::size_t n = params_[i].get().cols(); const std::size_t big = m > n ? m : n; - const Scalar muon_scale = - Scalar{0.2} * std::sqrt(static_cast(big)); + const Scalar muon_scale = Scalar{0.2} * std::sqrt(static_cast(big)); r = dsl::compute_into(engine_, - dsl::leaf(params_[i]) - + dsl::leaf(*ortho_update) * dsl::rparam(-lr_ * muon_scale), - params_[i]); + dsl::leaf(params_[i]) + + dsl::leaf(*ortho_update) * dsl::rparam(-lr_ * muon_scale), + params_[i]); NN_TRY_CHECK(r); } else { // 非 2D 参数（bias 等）：标准 SGD 更新（目标传递：原地、单 dispatch） r = dsl::compute_into(engine_, - dsl::leaf(params_[i]) + dsl::leaf(update) * dsl::rparam(-lr_), - params_[i]); + dsl::leaf(params_[i]) + dsl::leaf(update) * dsl::rparam(-lr_), + params_[i]); NN_TRY_CHECK(r); } } @@ -1087,7 +1030,8 @@ class Muon : public Optimizer [[nodiscard]] Result save_state(std::vector& out) const override { NN_TRY(r, check_ready_()); - for (const auto& v : velocities_) out.push_back(v); + for (const auto& v : velocities_) + out.push_back(v); return {}; } [[nodiscard]] Result load_state(std::span in) override @@ -1123,27 +1067,27 @@ struct OptimizerMeta { std::string_view name; std::string_view display; - bool beta; - std::uint8_t state_bufs; - bool half_state; + bool beta; + std::uint8_t state_bufs; + bool half_state; }; inline constexpr OptimizerMeta kOptimizerTable[] = { - {"sgd", "SGD", false, 0, false}, - {"sgd_momentum", "SGD + Momentum", false, 1, false}, - {"adam", "Adam", true, 2, false}, - {"adamw", "AdamW", true, 2, false}, - {"adamwquant", "AdamW（量化状态）", true, 2, true }, - {"muon", "Muon", false, 1, false}, + {"sgd", "SGD", false, 0, false}, + {"sgd_momentum", "SGD + Momentum", false, 1, false}, + {"adam", "Adam", true, 2, false}, + {"adamw", "AdamW", true, 2, false}, + {"adamwquant", "AdamW（量化状态）", true, 2, true}, + {"muon", "Muon", false, 1, false}, }; [[nodiscard]] constexpr std::span optimizer_table() noexcept { return kOptimizerTable; } // 按 CLI 名查元数据；未知名 → nullptr（调用方负责报错，不静默降级） -[[nodiscard]] inline const OptimizerMeta *find_optimizer(std::string_view name) noexcept +[[nodiscard]] inline const OptimizerMeta* find_optimizer(std::string_view name) noexcept { - for (const auto &m : kOptimizerTable) + for (const auto& m : kOptimizerTable) if (m.name == name) return &m; return nullptr; } @@ -1153,7 +1097,7 @@ inline constexpr OptimizerMeta kOptimizerTable[] = { [[nodiscard]] inline std::string optimizer_names(std::string_view sep) { std::string s; - for (const auto &m : kOptimizerTable) + for (const auto& m : kOptimizerTable) { if (!s.empty()) s += sep; s += m.name; @@ -1168,36 +1112,29 @@ inline constexpr OptimizerMeta kOptimizerTable[] = { // 走 in-place（存储精度不可变）。默认全 F32。 // beta1/beta2：仅 adam/adamw 使用（sgd 家族与 muon 忽略）。默认 0.9/0.999； // LLM 预训练惯例常取 beta2=0.95（如 GPT-3 / LLaMA 配方）。 -[[nodiscard]] inline std::unique_ptr create_optimizer( - std::string_view name, - ComputeEngine& engine, - std::vector params, - std::vector grads, - Scalar lr, - Scalar weight_decay = 0, - PrecisionProfile precision = PrecisionProfile{}, - Scalar beta1 = 0.9, - Scalar beta2 = 0.999) +[[nodiscard]] inline std::unique_ptr +create_optimizer(std::string_view name, ComputeEngine& engine, std::vector params, + std::vector grads, Scalar lr, Scalar weight_decay = 0, + PrecisionProfile precision = PrecisionProfile{}, Scalar beta1 = 0.9, + Scalar beta2 = 0.999) { if (name == "sgd") - return std::make_unique(engine, std::move(params), std::move(grads), lr, - precision); + return std::make_unique(engine, std::move(params), std::move(grads), lr, precision); if (name == "sgd_momentum") - return std::make_unique(engine, std::move(params), - std::move(grads), lr, 0.9, precision); + return std::make_unique(engine, std::move(params), std::move(grads), lr, + 0.9, precision); if (name == "adam") - return std::make_unique(engine, std::move(params), std::move(grads), lr, - beta1, beta2, 1e-8, precision); + return std::make_unique(engine, std::move(params), std::move(grads), lr, beta1, beta2, + 1e-8, precision); if (name == "adamw") - return std::make_unique(engine, std::move(params), std::move(grads), lr, - beta1, beta2, 1e-8, weight_decay, precision); + return std::make_unique(engine, std::move(params), std::move(grads), lr, beta1, + beta2, 1e-8, weight_decay, precision); if (name == "adamwquant") - return std::make_unique(engine, std::move(params), std::move(grads), - lr, beta1, beta2, 1e-8, weight_decay, - precision); + return std::make_unique(engine, std::move(params), std::move(grads), lr, beta1, + beta2, 1e-8, weight_decay, precision); if (name == "muon") - return std::make_unique(engine, std::move(params), std::move(grads), lr, - 0.95f, true, 5, 1e-7f, precision); + return std::make_unique(engine, std::move(params), std::move(grads), lr, 0.95f, true, + 5, 1e-7f, precision); return nullptr; // 未知名称，调用方应处理 } diff --git a/include/neuralnet.cpp/compute_position_encoding.hpp b/include/neuralnet.cpp/compute_position_encoding.hpp index 30d0df4c..ed04c6cf 100644 --- a/include/neuralnet.cpp/compute_position_encoding.hpp +++ b/include/neuralnet.cpp/compute_position_encoding.hpp @@ -68,13 +68,13 @@ class RotaryEmbedding const Scalar pd = static_cast(pos); for (std::size_t j = 0; j < half; ++j) { - const Scalar theta = pd / std::pow(Scalar{10000}, - static_cast(2 * j) / static_cast(d_k_)); + const Scalar theta = pd / std::pow(Scalar{10000}, static_cast(2 * j) / + static_cast(d_k_)); const Scalar cv = std::cos(theta); const Scalar sv = std::sin(theta); - c[j * cols + col] = cv; + c[j * cols + col] = cv; c[(half + j) * cols + col] = cv; - s[j * cols + col] = sv; + s[j * cols + col] = sv; s[(half + j) * cols + col] = sv; } } @@ -100,10 +100,7 @@ class RotaryEmbedding public: RotaryEmbedding() = default; - explicit RotaryEmbedding(std::size_t d_k) - : d_k_(d_k) - { - } + explicit RotaryEmbedding(std::size_t d_k) : d_k_(d_k) {} [[nodiscard]] Result init(ComputeEngine& /*engine*/) { return {}; } @@ -126,35 +123,31 @@ class RotaryEmbedding } // 全表应用：q 为 (batch*H*d_k, seq)，cos/sin 为 (d_k, seq) 短表 - [[nodiscard]] Result apply( - ComputeEngine& engine, const Tensor& q, - std::size_t seq, bool backward) + [[nodiscard]] Result apply(ComputeEngine& engine, const Tensor& q, std::size_t seq, + bool backward) { if (d_k_ == 0 || d_k_ % 2 != 0) NN_FAIL("RotaryEmbedding::apply: d_k must be positive and even"); - if (seq != seq_cached_) - { - NN_TRY(r, rebuild(engine, seq)); - } + if (seq != seq_cached_) { NN_TRY(r, rebuild(engine, seq)); } const std::uint32_t dk = static_cast(d_k_); // 内联 RoPE 表达式（LLaMA half-swap；backward 为旋转正交，逆 = 反角）： // forward: out = q·cos + rotate_half(q)·sin // backward: out = q·cos − rotate_half(q)·sin if (backward) return dsl::compute(engine, - dsl::leaf(q) * dsl::row_mod(cos_cache_, dk) - - dsl::rotate_half(q, dk) * dsl::row_mod(sin_cache_, dk), - q.rows(), q.cols(), p_.compute); + dsl::leaf(q) * dsl::row_mod(cos_cache_, dk) - + dsl::rotate_half(q, dk) * dsl::row_mod(sin_cache_, dk), + q.rows(), q.cols(), p_.compute); return dsl::compute(engine, - dsl::leaf(q) * dsl::row_mod(cos_cache_, dk) - + dsl::rotate_half(q, dk) * dsl::row_mod(sin_cache_, dk), - q.rows(), q.cols(), p_.compute); + dsl::leaf(q) * dsl::row_mod(cos_cache_, dk) + + dsl::rotate_half(q, dk) * dsl::row_mod(sin_cache_, dk), + q.rows(), q.cols(), p_.compute); } // 增量推理：q 为 (H*d_k, 1)，位置 = pos（cur_len） // 直接生成 (d_k,1) 位置表，避免对全表切片取列。 - [[nodiscard]] Result apply_step( - ComputeEngine& engine, const Tensor& q, std::size_t pos, bool backward) + [[nodiscard]] Result apply_step(ComputeEngine& engine, const Tensor& q, std::size_t pos, + bool backward) { if (d_k_ == 0 || d_k_ % 2 != 0) NN_FAIL("RotaryEmbedding::apply_step: d_k must be positive and even"); @@ -165,13 +158,13 @@ class RotaryEmbedding const std::uint32_t dk = static_cast(d_k_); if (backward) return dsl::compute(engine, - dsl::leaf(q) * dsl::row_mod(*cr, dk) - - dsl::rotate_half(q, dk) * dsl::row_mod(*sr, dk), - q.rows(), q.cols(), p_.compute); + dsl::leaf(q) * dsl::row_mod(*cr, dk) - + dsl::rotate_half(q, dk) * dsl::row_mod(*sr, dk), + q.rows(), q.cols(), p_.compute); return dsl::compute(engine, - dsl::leaf(q) * dsl::row_mod(*cr, dk) - + dsl::rotate_half(q, dk) * dsl::row_mod(*sr, dk), - q.rows(), q.cols(), p_.compute); + dsl::leaf(q) * dsl::row_mod(*cr, dk) + + dsl::rotate_half(q, dk) * dsl::row_mod(*sr, dk), + q.rows(), q.cols(), p_.compute); } }; @@ -208,20 +201,18 @@ class PositionEncoder // ── 嵌入侧注入（Learned / Sinusoidal 覆写；其余 no-op）────────────── // apply：token_emb_T 为 (d_model, batch*seq)，返回 x = token_emb_T (+ pos_emb) - [[nodiscard]] virtual Result apply( - ComputeEngine& /*engine*/, const Tensor& token_emb_T, - std::size_t /*batch*/, std::size_t /*seq*/) + [[nodiscard]] virtual Result apply(ComputeEngine& /*engine*/, const Tensor& token_emb_T, + std::size_t /*batch*/, std::size_t /*seq*/) { return token_emb_T; } // 增量前向：x 为 (d_model, 1)，返回 x + pos_emb[pos] - [[nodiscard]] virtual Result apply_step( - ComputeEngine& /*engine*/, const Tensor& x, std::size_t /*pos*/) + [[nodiscard]] virtual Result apply_step(ComputeEngine& /*engine*/, const Tensor& x, + std::size_t /*pos*/) { return x; } // 反向：累计位置梯度到 grad_pos_emb_（默认 no-op） - [[nodiscard]] virtual Result backward( - ComputeEngine& /*engine*/, const Tensor& /*grad_T*/, - std::size_t /*batch*/, std::size_t /*seq*/) + [[nodiscard]] virtual Result backward(ComputeEngine& /*engine*/, const Tensor& /*grad_T*/, + std::size_t /*batch*/, std::size_t /*seq*/) { return {}; } [[nodiscard]] virtual std::vector parameters() { return {}; } @@ -229,13 +220,13 @@ class PositionEncoder // ── 注意力 Q/K 侧（RoPE 覆写；其余 no-op）───────────────────────── // 全表：qk 为 (batch*H*d_k, seq)，列 = position；**原地替换** - [[nodiscard]] virtual Result apply_qk( - ComputeEngine& /*engine*/, Tensor& /*qk*/, std::size_t /*seq*/, bool /*backward*/) + [[nodiscard]] virtual Result apply_qk(ComputeEngine& /*engine*/, Tensor& /*qk*/, + std::size_t /*seq*/, bool /*backward*/) { return {}; } // 增量：qk 为 (H*d_k, 1)，位置 = pos - [[nodiscard]] virtual Result apply_qk_step( - ComputeEngine& /*engine*/, Tensor& /*qk*/, std::size_t /*pos*/, bool /*backward*/) + [[nodiscard]] virtual Result apply_qk_step(ComputeEngine& /*engine*/, Tensor& /*qk*/, + std::size_t /*pos*/, bool /*backward*/) { return {}; } // 绝对位置偏移（滑动窗生成；仅 RoPE 有意义） @@ -248,8 +239,8 @@ class PositionEncoder [[nodiscard]] virtual bool has_score_bias() const noexcept { return false; } // 每步准备分数偏置表（batch/seq 变化时重建） - [[nodiscard]] virtual Result prepare_score_bias( - ComputeEngine& /*engine*/, std::size_t /*batch*/, std::size_t /*seq*/) + [[nodiscard]] virtual Result + prepare_score_bias(ComputeEngine& /*engine*/, std::size_t /*batch*/, std::size_t /*seq*/) { return {}; } // fold 输入追加（斜率表）。调用方必须**先**调本函数、**再**追加掩码输入， @@ -264,15 +255,15 @@ class PositionEncoder // 掩码与偏置是两件事，融成一条 select 会让掩码策略知道位置编码的存在。 // 两种写法数值等价（被屏蔽处 -inf 加有限值仍是 -inf；未屏蔽处 s+0 精确）， // 代价是反向多一个逐元素原地 kernel。 - [[nodiscard]] virtual Result apply_score_bias( - ComputeEngine& /*engine*/, Tensor& /*score*/, - std::size_t /*BH*/, std::size_t /*seq*/) + [[nodiscard]] virtual Result apply_score_bias(ComputeEngine& /*engine*/, + Tensor& /*score*/, std::size_t /*BH*/, + std::size_t /*seq*/) { return {}; } // 增量推理的分数偏置（scores 为 (H, new_len)） - [[nodiscard]] virtual Result apply_bias_step( - ComputeEngine& /*engine*/, Tensor&& scores, std::size_t /*cur_len*/, - const PrecisionProfile& /*p*/) const + [[nodiscard]] virtual Result apply_bias_step(ComputeEngine& /*engine*/, Tensor&& scores, + std::size_t /*cur_len*/, + const PrecisionProfile& /*p*/) const { return std::move(scores); } }; @@ -300,29 +291,25 @@ class AdditivePositionEncoder : public PositionEncoder const std::size_t cols = pe.cols(); learnable_ = learnable; pos_emb_ = std::move(pe); - if (!pos_emb_.valid()) - NN_FAIL("PositionEncoder: 位置编码创建失败"); + if (!pos_emb_.valid()) NN_FAIL("PositionEncoder: 位置编码创建失败"); if (learnable_) { grad_pos_emb_ = engine.create_tensor(rows, cols, p_.param, InitSpec::zero()); - if (!grad_pos_emb_.valid()) - NN_FAIL("PositionEncoder: 梯度缓冲初始化失败"); + if (!grad_pos_emb_.valid()) NN_FAIL("PositionEncoder: 梯度缓冲初始化失败"); } return {}; } // 确保 pos_indices 缓存有效（batch-major：i = b*seq + t → position=t） - [[nodiscard]] Result ensure_pos_indices_( - ComputeEngine& engine, std::size_t batch, std::size_t seq) + [[nodiscard]] Result ensure_pos_indices_(ComputeEngine& engine, std::size_t batch, + std::size_t seq) { - if (pos_indices_batch_ == batch && pos_indices_seq_ == seq) - return {}; + if (pos_indices_batch_ == batch && pos_indices_seq_ == seq) return {}; std::vector pidx(batch * seq); // 宿主桥（17 §3 D11）：不经 Matrix for (std::size_t b = 0; b < batch; ++b) for (std::size_t t = 0; t < seq; ++t) pidx[b * seq + t] = static_cast(t); - auto pidx_t = detail::upload_span(engine, batch * seq, 1, Precision::F32, - std::span(pidx)); + auto pidx_t = detail::upload_span(engine, batch * seq, 1, Precision::F32, std::span(pidx)); NN_TRY_CHECK(pidx_t); pos_indices_cache_ = std::move(*pidx_t); pos_indices_batch_ = batch; @@ -335,38 +322,34 @@ class AdditivePositionEncoder : public PositionEncoder [[nodiscard]] bool learnable() const noexcept { return learnable_; } - [[nodiscard]] Result apply( - ComputeEngine& engine, const Tensor& token_emb_T, - std::size_t batch, std::size_t seq) override + [[nodiscard]] Result apply(ComputeEngine& engine, const Tensor& token_emb_T, + std::size_t batch, std::size_t seq) override { NN_TRY(ci, ensure_pos_indices_(engine, batch, seq)); NN_TRY(pos_gathered, engine.gather_rows(pos_emb_, pos_indices_cache_)); NN_TRY(pos_T, engine.transpose(*pos_gathered)); - auto x_with_pos = dsl::compute(engine, - dsl::leaf(token_emb_T) + dsl::leaf(*pos_T), - token_emb_T.rows(), token_emb_T.cols(), p_.compute); + auto x_with_pos = dsl::compute(engine, dsl::leaf(token_emb_T) + dsl::leaf(*pos_T), + token_emb_T.rows(), token_emb_T.cols(), p_.compute); NN_TRY_CHECK(x_with_pos); return std::move(*x_with_pos); } - [[nodiscard]] Result apply_step( - ComputeEngine& engine, const Tensor& x, std::size_t pos) override + [[nodiscard]] Result apply_step(ComputeEngine& engine, const Tensor& x, + std::size_t pos) override { std::vector pos_v(1); // 宿主桥（17 §3 D11） pos_v[0] = static_cast(pos); NN_TRY(pos_t, detail::upload_span(engine, 1, 1, Precision::F32, std::span(pos_v))); NN_TRY(pos_emb_g, engine.gather_rows(pos_emb_, *pos_t)); NN_TRY(pos_T, engine.transpose(*pos_emb_g)); - auto x_wp = dsl::compute(engine, - dsl::leaf(x) + dsl::leaf(*pos_T), - x.rows(), x.cols(), p_.compute); + auto x_wp = + dsl::compute(engine, dsl::leaf(x) + dsl::leaf(*pos_T), x.rows(), x.cols(), p_.compute); NN_TRY_CHECK(x_wp); return std::move(*x_wp); } - [[nodiscard]] Result backward( - ComputeEngine& engine, const Tensor& grad_T, - std::size_t /*batch*/, std::size_t /*seq*/) override + [[nodiscard]] Result backward(ComputeEngine& engine, const Tensor& grad_T, + std::size_t /*batch*/, std::size_t /*seq*/) override { if (!learnable_) return {}; // pos_indices 缓存由 apply() 建立，backward 直接复用（batch/seq 一致）。 @@ -377,12 +360,12 @@ class AdditivePositionEncoder : public PositionEncoder [[nodiscard]] std::vector parameters() override { if (!learnable_) return {}; - return { pos_emb_ }; + return {pos_emb_}; } [[nodiscard]] std::vector param_gradients() override { if (!learnable_) return {}; - return { grad_pos_emb_ }; + return {grad_pos_emb_}; } }; @@ -394,7 +377,9 @@ class LearnedPositionEncoder final : public AdditivePositionEncoder public: LearnedPositionEncoder(std::size_t d_model, std::size_t seq_len) - : d_model_(d_model), seq_len_(seq_len) {} + : d_model_(d_model), seq_len_(seq_len) + { + } [[nodiscard]] Result init(ComputeEngine& engine) override { @@ -402,8 +387,7 @@ class LearnedPositionEncoder final : public AdditivePositionEncoder // 分布 seed 显式传（U1）；不再自持 RNG。 auto pe = engine.create_tensor(seq_len_, d_model_, p_.param, InitSpec::normal(0, 0.02, kInitSeed)); - if (!pe.valid()) - NN_FAIL("LearnedPositionEncoder: 初始化失败"); + if (!pe.valid()) NN_FAIL("LearnedPositionEncoder: 初始化失败"); return init_(engine, std::move(pe), /*learnable=*/true); } }; @@ -417,7 +401,9 @@ class SinusoidalPositionEncoder final : public AdditivePositionEncoder public: SinusoidalPositionEncoder(std::size_t d_model, std::size_t seq_len) - : d_model_(d_model), seq_len_(seq_len) {} + : d_model_(d_model), seq_len_(seq_len) + { + } [[nodiscard]] Result init(ComputeEngine& engine) override { @@ -426,13 +412,13 @@ class SinusoidalPositionEncoder final : public AdditivePositionEncoder 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_)); + std::pow(Scalar{10000}, static_cast(2 * (i / 2)) / + static_cast(d_model_)); pe[pos * d_model_ + i] = (i % 2 == 0) ? std::sin(angle) : std::cos(angle); } // 闭式公式（非分布/常数）仍宿主计算 → span 上传进引擎（M2 裁定：InitSpec // 只收分布与常数初始化；公式数据自 M4 起经 detail::upload_span，不经 Matrix） - auto pe_t = detail::upload_span(engine, seq_len_, d_model_, p_.param, - std::span(pe)); + auto pe_t = detail::upload_span(engine, seq_len_, d_model_, p_.param, std::span(pe)); NN_TRY_CHECK(pe_t); return init_(engine, std::move(*pe_t), /*learnable=*/false); } @@ -454,10 +440,7 @@ class RopePositionEncoder final : public PositionEncoder public: explicit RopePositionEncoder(std::size_t d_k) : rope_(d_k) {} - [[nodiscard]] Result init(ComputeEngine& engine) override - { - return rope_.init(engine); - } + [[nodiscard]] Result init(ComputeEngine& engine) override { return rope_.init(engine); } void set_precision_profile(const PrecisionProfile& p) override { @@ -465,16 +448,16 @@ class RopePositionEncoder final : public PositionEncoder rope_.set_precision_profile(p); } - [[nodiscard]] Result apply_qk( - ComputeEngine& engine, Tensor& qk, std::size_t seq, bool backward) override + [[nodiscard]] Result apply_qk(ComputeEngine& engine, Tensor& qk, std::size_t seq, + bool backward) override { NN_TRY(r, rope_.apply(engine, qk, seq, backward)); qk = std::move(*r); return {}; } - [[nodiscard]] Result apply_qk_step( - ComputeEngine& engine, Tensor& qk, std::size_t pos, bool backward) override + [[nodiscard]] Result apply_qk_step(ComputeEngine& engine, Tensor& qk, std::size_t pos, + bool backward) override { NN_TRY(r, rope_.apply_step(engine, qk, pos, backward)); qk = std::move(*r); @@ -512,11 +495,10 @@ class AlibiPositionEncoder final : public PositionEncoder [[nodiscard]] std::size_t num_heads() const noexcept { return num_heads_; } [[nodiscard]] bool has_score_bias() const noexcept override { return true; } - [[nodiscard]] Result prepare_score_bias( - ComputeEngine& engine, std::size_t batch, std::size_t seq) override + [[nodiscard]] Result prepare_score_bias(ComputeEngine& engine, std::size_t batch, + std::size_t seq) override { - if (slopes_cache_.valid() && cached_batch_ == batch && cached_seq_ == seq) - return {}; + if (slopes_cache_.valid() && cached_batch_ == batch && cached_seq_ == seq) return {}; const std::size_t BH = batch * num_heads_; const std::size_t rows = BH * seq; // ① fold 输入 (1, BH)：fold 的 batch_mod(bh) 按下标 b*H+h 直读 @@ -546,8 +528,7 @@ class AlibiPositionEncoder final : public PositionEncoder void append_score_bias_inputs(std::vector& io) const override { io.push_back(slopes_cache_); } - [[nodiscard]] const Tensor* bias_slopes() const noexcept override - { return &slopes_cache_; } + [[nodiscard]] const Tensor* bias_slopes() const noexcept override { return &slopes_cache_; } // ALiBi 项：m_h · (Col − Row)（BatchMod 按下标 b*H+h 直读斜率表） // 反向：S += m_h · (Col − Row)，写回 S。 @@ -561,27 +542,26 @@ class AlibiPositionEncoder final : public PositionEncoder // （实测 |bias| 达 27，而理论界为 (seq−1)·m_0≈7，alibi gradcheck 直接失败）。 // (rows,1) 行表 + row_broadcast 按**全局行号**直读，语义与网格分解无关， // 因而独立成步也正确。表长 O(B·H·seq)，与掩码的 doc 缓存同量级。 - [[nodiscard]] Result apply_score_bias( - ComputeEngine& engine, Tensor& score, - std::size_t /*BH*/, std::size_t /*seq*/) override + [[nodiscard]] Result apply_score_bias(ComputeEngine& engine, Tensor& score, + std::size_t /*BH*/, std::size_t /*seq*/) override { // RV-01：缓存未准备时下面的 row_broadcast 会读空张量 → 无条件终止 NN_ASSERT_ALWAYS(slope_row_cache_.valid() && pos_row_cache_.valid(), "AlibiPositionEncoder: prepare_score_bias 未被调用"); NN_TRY(res, dsl::compute(engine, - dsl::leaf(score) - + dsl::row_broadcast(slope_row_cache_) - * (dsl::col() - dsl::row_broadcast(pos_row_cache_)), - score.rows(), score.cols(), score.precision())); + dsl::leaf(score) + + dsl::row_broadcast(slope_row_cache_) * + (dsl::col() - dsl::row_broadcast(pos_row_cache_)), + score.rows(), score.cols(), score.precision())); score = std::move(*res); return {}; } // 增量推理：scores 布局 (H, new_len)，new_len = cur_len + 1； // query 在位置 cur_len，key 在位置 j → 偏置 = −slope[h] · (cur_len − j) - [[nodiscard]] Result apply_bias_step( - ComputeEngine& engine, Tensor&& scores, std::size_t cur_len, - const PrecisionProfile& p) const override + [[nodiscard]] Result apply_bias_step(ComputeEngine& engine, Tensor&& scores, + std::size_t cur_len, + const PrecisionProfile& p) const override { const std::size_t new_len = cur_len + 1; std::vector bias(num_heads_ * new_len); @@ -594,12 +574,11 @@ class AlibiPositionEncoder final : public PositionEncoder bias[h * new_len + j] = -slope * static_cast(dist); } } - auto bias_t = detail::upload_span(engine, num_heads_, new_len, Precision::F32, - std::span(bias)); + auto bias_t = + detail::upload_span(engine, num_heads_, new_len, Precision::F32, std::span(bias)); NN_TRY_CHECK(bias_t); - return dsl::compute(engine, - dsl::leaf(scores) + dsl::leaf(*bias_t), - scores.rows(), scores.cols(), p.compute); + return dsl::compute(engine, dsl::leaf(scores) + dsl::leaf(*bias_t), scores.rows(), + scores.cols(), p.compute); } }; @@ -621,31 +600,31 @@ class AlibiPositionEncoder final : public PositionEncoder // 每层 2·d_k·seq·4B（本仓默认 d_model=128/H=4/layers=4/seq=256 ≈ 64KB/层， // 全模型 256KB）。层数很多时才需重新考虑共享。 // ═══════════════════════════════════════════════════════════════════════════ -[[nodiscard]] inline std::unique_ptr make_embedding_position_encoder( - PosEncodingType type, std::size_t d_model, std::size_t seq_len) +[[nodiscard]] inline std::unique_ptr +make_embedding_position_encoder(PosEncodingType type, std::size_t d_model, std::size_t seq_len) { switch (type) { - case PosEncodingType::Learned: - return std::make_unique(d_model, seq_len); - case PosEncodingType::Sinusoidal: - return std::make_unique(d_model, seq_len); - default: // ALiBi / RoPE：位置信息在注意力层注入，嵌入侧无动作 - return std::make_unique(); + case PosEncodingType::Learned: + return std::make_unique(d_model, seq_len); + case PosEncodingType::Sinusoidal: + return std::make_unique(d_model, seq_len); + default: // ALiBi / RoPE：位置信息在注意力层注入，嵌入侧无动作 + return std::make_unique(); } } -[[nodiscard]] inline std::unique_ptr make_attention_position_encoder( - PosEncodingType type, std::size_t d_k, std::size_t num_heads) +[[nodiscard]] inline std::unique_ptr +make_attention_position_encoder(PosEncodingType type, std::size_t d_k, std::size_t num_heads) { switch (type) { - case PosEncodingType::ALiBi: - return std::make_unique(num_heads); - case PosEncodingType::RoPE: - return std::make_unique(d_k); - default: // Learned / Sinusoidal：位置信息在嵌入侧注入，注意力层无动作 - return std::make_unique(); + case PosEncodingType::ALiBi: + return std::make_unique(num_heads); + case PosEncodingType::RoPE: + return std::make_unique(d_k); + default: // Learned / Sinusoidal：位置信息在嵌入侧注入，注意力层无动作 + return std::make_unique(); } } diff --git a/include/neuralnet.cpp/compute_tensor.hpp b/include/neuralnet.cpp/compute_tensor.hpp index 21d41666..6b775ab2 100644 --- a/include/neuralnet.cpp/compute_tensor.hpp +++ b/include/neuralnet.cpp/compute_tensor.hpp @@ -101,8 +101,7 @@ class Tensor // ── variant 安全访问辅助（避免 std::get UB，§6.1）──────────────────── - template - [[nodiscard]] std::shared_ptr> cpu_get() const noexcept + template [[nodiscard]] std::shared_ptr> cpu_get() const noexcept { constexpr auto slot = static_cast(P); return cpu_data_.index() == slot ? std::get(cpu_data_) : nullptr; @@ -113,22 +112,19 @@ class Tensor // cpu_matrix()；若每次拷贝 shared_ptr，则每元素两次原子引用计数操作 // （实测 ~20ns/元素，使 DSL 模板路径比等价手写循环慢 8-70 倍）。 // 调用方（cpu_matrix 的调用者）本身持有该 Tensor，故裸指针在其生存期内有效。 - template - [[nodiscard]] MatrixT

* cpu_get_ptr() noexcept + template [[nodiscard]] MatrixT

* cpu_get_ptr() noexcept { constexpr auto slot = static_cast(P); return cpu_data_.index() == slot ? std::get(cpu_data_).get() : nullptr; } - template - [[nodiscard]] const MatrixT

* cpu_get_ptr() const noexcept + template [[nodiscard]] const MatrixT

* cpu_get_ptr() const noexcept { constexpr auto slot = static_cast(P); return cpu_data_.index() == slot ? std::get(cpu_data_).get() : nullptr; } #ifdef NN_HAS_VULKAN - template - [[nodiscard]] std::shared_ptr> gpu_get() const noexcept + template [[nodiscard]] std::shared_ptr> gpu_get() const noexcept { constexpr auto slot = static_cast(P); return gpu_data_.index() == slot ? std::get(gpu_data_) : nullptr; @@ -136,14 +132,12 @@ class Tensor // 同 gpu_get，但返回**裸指针、不拷贝 shared_ptr**（设备域派生/import 热路径： // 每次拷贝 shared_ptr 都是一对原子操作）。 - template - [[nodiscard]] GpuTensorT

* gpu_get_ptr() noexcept + template [[nodiscard]] GpuTensorT

* gpu_get_ptr() noexcept { constexpr auto slot = static_cast(P); return gpu_data_.index() == slot ? std::get(gpu_data_).get() : nullptr; } - template - [[nodiscard]] const GpuTensorT

* gpu_get_ptr() const noexcept + template [[nodiscard]] const GpuTensorT

* gpu_get_ptr() const noexcept { constexpr auto slot = static_cast(P); return gpu_data_.index() == slot ? std::get(gpu_data_).get() : nullptr; @@ -156,33 +150,28 @@ class Tensor private: // ── CPU 构造（f32）──────────────────────────────────────────────────── explicit Tensor(std::shared_ptr m) - : device_(Device::CPU), precision_(Precision::F32), - rows_(m->rows()), cols_(m->cols()), - cpu_data_(std::in_place_index<1>, std::move(m)) {} + : device_(Device::CPU), precision_(Precision::F32), rows_(m->rows()), cols_(m->cols()), + cpu_data_(std::in_place_index<1>, std::move(m)) + { + } // ── CPU 构造（f16，§6.4）────────────────────────────────────────────── explicit Tensor(std::shared_ptr> m) - : device_(Device::CPU), precision_(Precision::F16), - rows_(m->rows()), cols_(m->cols()), - cpu_data_(std::in_place_index<0>, std::move(m)) {} - - // ── 从 Matrix 创建（f32 by-value）───────────────────────────────────── - static Tensor from_matrix(Matrix m) + : device_(Device::CPU), precision_(Precision::F16), rows_(m->rows()), cols_(m->cols()), + cpu_data_(std::in_place_index<0>, std::move(m)) { - return Tensor(std::make_shared(std::move(m))); } + // ── 从 Matrix 创建（f32 by-value）───────────────────────────────────── + static Tensor from_matrix(Matrix m) { return Tensor(std::make_shared(std::move(m))); } + // ── 从 MatrixT 创建（f16 by-value，§6.4）──────────────────────── static Tensor from_matrix(MatrixT m) - { - return Tensor(std::make_shared>(std::move(m))); - } + { return Tensor(std::make_shared>(std::move(m))); } // ── 创建 CPU 空张量（f32）───────────────────────────────────────────── static Tensor cpu(std::size_t rows, std::size_t cols) - { - return Tensor(std::make_shared(rows, cols)); - } + { return Tensor(std::make_shared(rows, cols)); } // ── 创建 CPU 未初始化张量（f32）─────────────────────────────────────── // 语义契约：调用方必须在任何读取之前把**全部**元素写满。 @@ -190,14 +179,11 @@ class Tensor // 输出缓冲）：省掉"分配 + 写满一遍零 + 马上被全覆盖"里的那一遍零写。 // 实测本机单线程写满 1.57MB 要 0.50ms（~3.2 GB/s），是纯浪费。 [[nodiscard]] static Tensor cpu_uninitialized(std::size_t rows, std::size_t cols) - { - return Tensor(std::make_shared(rows, cols, Matrix::uninitialized_tag{})); - } + { return Tensor(std::make_shared(rows, cols, Matrix::uninitialized_tag{})); } // ── 创建 CPU 空张量（显式 P，§6.4）─────────────────────────────────── // compile-time P；BF16/F64 → static_assert 报错 - template - [[nodiscard]] static Tensor cpu(std::size_t rows, std::size_t cols) + template [[nodiscard]] static Tensor cpu(std::size_t rows, std::size_t cols) { static_assert(P == Precision::F16 || P == Precision::F32, "Phase 1 仅支持 F16/F32（BF16/F64 为保留值）"); @@ -207,27 +193,25 @@ class Tensor #ifdef NN_HAS_VULKAN // ── GPU 构造（f32）──────────────────────────────────────────────────── explicit Tensor(std::shared_ptr t) - : device_(Device::GPU), precision_(Precision::F32), - rows_(t->rows()), cols_(t->cols()), - gpu_data_(std::in_place_index<1>, std::move(t)) {} + : device_(Device::GPU), precision_(Precision::F32), rows_(t->rows()), cols_(t->cols()), + gpu_data_(std::in_place_index<1>, std::move(t)) + { + } // ── GPU 构造（f16，§6.4）────────────────────────────────────────────── explicit Tensor(std::shared_ptr> t) - : device_(Device::GPU), precision_(Precision::F16), - rows_(t->rows()), cols_(t->cols()), - gpu_data_(std::in_place_index<0>, std::move(t)) {} + : device_(Device::GPU), precision_(Precision::F16), rows_(t->rows()), cols_(t->cols()), + gpu_data_(std::in_place_index<0>, std::move(t)) + { + } // ── 从 GpuTensor 创建（f32）─────────────────────────────────────────── static Tensor from_gpu(GpuTensor t) - { - return Tensor(std::make_shared(std::move(t))); - } + { return Tensor(std::make_shared(std::move(t))); } // ── 从 GpuTensorT 创建（f16）───────────────────────────────────── static Tensor from_gpu(GpuTensorT t) - { - return Tensor(std::make_shared>(std::move(t))); - } + { return Tensor(std::make_shared>(std::move(t))); } #endif public: @@ -258,13 +242,13 @@ class Tensor if (device_ == Device::CPU) { return precision_ == Precision::F16 - ? cpu_get() == o.cpu_get() - : cpu_get() == o.cpu_get(); + ? cpu_get() == o.cpu_get() + : cpu_get() == o.cpu_get(); } #ifdef NN_HAS_VULKAN return precision_ == Precision::F16 - ? gpu_get() == o.gpu_get() - : gpu_get() == o.gpu_get(); + ? gpu_get() == o.gpu_get() + : gpu_get() == o.gpu_get(); #else return false; #endif @@ -311,11 +295,9 @@ class Tensor private: // ── CPU 存储访问（模板化，P 默认 F32，§6.1）────────────────────────── - template - [[nodiscard]] MatrixT

& cpu_matrix() + template [[nodiscard]] MatrixT

& cpu_matrix() { - static_assert(P == Precision::F16 || P == Precision::F32, - "Phase 1 仅支持 F16/F32"); + static_assert(P == Precision::F16 || P == Precision::F32, "Phase 1 仅支持 F16/F32"); // 用 cpu_get_ptr（不拷贝 shared_ptr）：该函数在逐元素热点上会被按元素 // 调用，拷贝 shared_ptr 的原子引用计数是实测的主要开销来源。 MatrixT

* p = cpu_get_ptr

(); @@ -324,11 +306,9 @@ class Tensor return *p; } - template - [[nodiscard]] const MatrixT

& cpu_matrix() const + template [[nodiscard]] const MatrixT

& cpu_matrix() const { - static_assert(P == Precision::F16 || P == Precision::F32, - "Phase 1 仅支持 F16/F32"); + static_assert(P == Precision::F16 || P == Precision::F32, "Phase 1 仅支持 F16/F32"); const MatrixT

* p = cpu_get_ptr

(); NN_ASSERT_ALWAYS(device_ == Device::CPU && p, "cpu_matrix

() const: tensor has no P-precision CPU storage"); @@ -337,35 +317,28 @@ class Tensor template [[nodiscard]] std::shared_ptr> cpu_shared() const noexcept - { - return cpu_get

(); - } + { return cpu_get

(); } #ifdef NN_HAS_VULKAN // ── GPU 存储访问（模板化，P 默认 F32，§6.1）────────────────────────── - template - [[nodiscard]] GpuTensorT

& gpu_tensor() + template [[nodiscard]] GpuTensorT

& gpu_tensor() { - static_assert(P == Precision::F16 || P == Precision::F32, - "Phase 1 仅支持 F16/F32"); + static_assert(P == Precision::F16 || P == Precision::F32, "Phase 1 仅支持 F16/F32"); auto p = gpu_get

(); if (!p) std::fprintf(stderr, "[gpu_tensor-fail] P=%d precision_=%d gpu_slot=%zu " "dev=%d %zux%zu\n", - static_cast(P), static_cast(precision_), - gpu_data_.index(), static_cast(device_), - rows_, cols_); + static_cast(P), static_cast(precision_), gpu_data_.index(), + static_cast(device_), rows_, cols_); NN_ASSERT_ALWAYS(device_ == Device::GPU && p, "gpu_tensor

(): tensor has no P-precision GPU storage"); return *p; } - template - [[nodiscard]] const GpuTensorT

& gpu_tensor() const + template [[nodiscard]] const GpuTensorT

& gpu_tensor() const { - static_assert(P == Precision::F16 || P == Precision::F32, - "Phase 1 仅支持 F16/F32"); + static_assert(P == Precision::F16 || P == Precision::F32, "Phase 1 仅支持 F16/F32"); auto p = gpu_get

(); NN_ASSERT_ALWAYS(device_ == Device::GPU && p, "gpu_tensor

() const: tensor has no P-precision GPU storage"); @@ -374,9 +347,7 @@ class Tensor template [[nodiscard]] std::shared_ptr> gpu_shared() const noexcept - { - return gpu_get

(); - } + { return gpu_get

(); } #endif // 推进写版本：**只允许 ComputeEngine 的写入口护栏调用**（friend 限定）。 @@ -387,9 +358,8 @@ class Tensor // ── 形状描述（调试用） ──────────────────────────────────────────────── [[nodiscard]] std::string shape_str() const { - return "(" + std::to_string(rows_) + "x" + std::to_string(cols_) + ")" - + (device_ == Device::CPU ? "[CPU]" : "[GPU]") - + precision_name(precision_); + return "(" + std::to_string(rows_) + "x" + std::to_string(cols_) + ")" + + (device_ == Device::CPU ? "[CPU]" : "[GPU]") + precision_name(precision_); } private: @@ -402,11 +372,9 @@ class Tensor // 私有 + friend 引擎：不重新引入"从张量找引擎"的用法。 [[nodiscard]] DomainId device_domain_() const noexcept { - if (device_ == Device::CPU) - return valid() ? kHostDomain : 0; + if (device_ == Device::CPU) return valid() ? kHostDomain : 0; #ifdef NN_HAS_VULKAN - if (!valid()) - return 0; + if (!valid()) return 0; if (const auto* t16 = gpu_get_ptr()) return t16->buffer().backend_identity(); if (const auto* t32 = gpu_get_ptr()) @@ -423,8 +391,7 @@ class Tensor // 视图"能力，按 precision_ 复制数据到新形状的 Matrix（§6.1）。 [[nodiscard]] Tensor reshape(std::size_t new_rows, std::size_t new_cols) const { - NN_ASSERT_ALWAYS(rows_ * cols_ == new_rows * new_cols, - "reshape: element count mismatch"); + NN_ASSERT_ALWAYS(rows_ * cols_ == new_rows * new_cols, "reshape: element count mismatch"); Tensor t; t.device_ = device_; t.precision_ = precision_; @@ -537,10 +504,11 @@ inline void Tensor::touch_write_epoch_() noexcept // 绕过引擎摸张量的正道是 engine.from_matrix / to_matrix / create_tensor。 struct detail::TensorAccess { + template [[nodiscard]] static MatrixT

& cpu_matrix(Tensor& t) + { return t.cpu_matrix

(); } template - [[nodiscard]] static MatrixT

& cpu_matrix(Tensor& t) { return t.cpu_matrix

(); } - template - [[nodiscard]] static const MatrixT

& cpu_matrix(const Tensor& t) { return t.cpu_matrix

(); } + [[nodiscard]] static const MatrixT

& cpu_matrix(const Tensor& t) + { return t.cpu_matrix

(); } // 引擎域内的"写下"记账：DSL 的 CPU 编译期模板路径（eval_into_tensor_cpu） // **直接**写 dst 的 span，不经 ComputeEngine 的写入口 —— 必须自己推进写 @@ -549,22 +517,16 @@ struct detail::TensorAccess static void mark_written(Tensor& t) noexcept { t.touch_write_epoch_(); } [[nodiscard]] static Tensor from_matrix(Matrix m) { return Tensor::from_matrix(std::move(m)); } - [[nodiscard]] static Tensor from_matrix(MatrixT m) { return Tensor::from_matrix(std::move(m)); } + [[nodiscard]] static Tensor from_matrix(MatrixT m) + { return Tensor::from_matrix(std::move(m)); } [[nodiscard]] static Tensor cpu(std::size_t rows, std::size_t cols) - { - return Tensor::cpu(rows, cols); - } - template - [[nodiscard]] static Tensor cpu(std::size_t rows, std::size_t cols) - { - return Tensor::cpu

(rows, cols); - } + { return Tensor::cpu(rows, cols); } + template [[nodiscard]] static Tensor cpu(std::size_t rows, std::size_t cols) + { return Tensor::cpu

(rows, cols); } [[nodiscard]] static Tensor cpu_uninitialized(std::size_t rows, std::size_t cols) - { - return Tensor::cpu_uninitialized(rows, cols); - } + { return Tensor::cpu_uninitialized(rows, cols); } }; // ── 非拥有型引用包装器（替代裸指针） ───────────────────────────────────── diff --git a/include/neuralnet.cpp/core_env.hpp b/include/neuralnet.cpp/core_env.hpp index 23500122..52eaff63 100644 --- a/include/neuralnet.cpp/core_env.hpp +++ b/include/neuralnet.cpp/core_env.hpp @@ -52,8 +52,7 @@ namespace nn::detail #if defined(_MSC_VER) char* buf = nullptr; std::size_t len = 0; - if (_dupenv_s(&buf, &len, name) != 0 || buf == nullptr) - return {}; + if (_dupenv_s(&buf, &len, name) != 0 || buf == nullptr) return {}; std::string s(buf); std::free(buf); return s; @@ -68,15 +67,12 @@ namespace nn::detail [[nodiscard]] inline int env_int_raw(const char* name, int fallback) noexcept { const std::string s = env_string_raw(name); - if (s.empty()) - return fallback; + if (s.empty()) return fallback; const char* p = s.c_str(); char* end = nullptr; const long v = std::strtol(p, &end, 10); - if (end == p || *end != '\0') - return fallback; - if (v < -2147483647L - 1L || v > 2147483647L) - return fallback; + if (end == p || *end != '\0') return fallback; + if (v < -2147483647L - 1L || v > 2147483647L) return fallback; return static_cast(v); } diff --git a/include/neuralnet.cpp/core_observer_ptr.hpp b/include/neuralnet.cpp/core_observer_ptr.hpp index 8cde4d71..7cfc3845 100644 --- a/include/neuralnet.cpp/core_observer_ptr.hpp +++ b/include/neuralnet.cpp/core_observer_ptr.hpp @@ -27,8 +27,7 @@ namespace nn { -template -class observer_ptr +template class observer_ptr { T* ptr_ = nullptr; @@ -42,7 +41,9 @@ class observer_ptr // 从另一个 observer_ptr 转换（允许隐式向上转型） template requires std::is_convertible_v - constexpr observer_ptr(const observer_ptr& other) noexcept : ptr_(other.get()) {} + constexpr observer_ptr(const observer_ptr& other) noexcept : ptr_(other.get()) + { + } // ── 访问 ───────────────────────────────────────────────────────────── [[nodiscard]] constexpr T& operator*() const noexcept { return *ptr_; } @@ -60,22 +61,14 @@ class observer_ptr friend constexpr bool operator==(const observer_ptr&, const observer_ptr&) = default; friend constexpr bool operator==(const observer_ptr& p, std::nullptr_t) noexcept { return !p; } friend constexpr auto operator<=>(const observer_ptr& a, const observer_ptr& b) noexcept - { - return a.ptr_ <=> b.ptr_; - } + { return a.ptr_ <=> b.ptr_; } }; // ── 便捷工厂函数 ─────────────────────────────────────────────────────── -template -[[nodiscard]] constexpr observer_ptr make_observer(T& obj) noexcept -{ - return observer_ptr(obj); -} +template [[nodiscard]] constexpr observer_ptr make_observer(T& obj) noexcept +{ return observer_ptr(obj); } -template -[[nodiscard]] constexpr observer_ptr make_observer(T* ptr) noexcept -{ - return observer_ptr(ptr); -} +template [[nodiscard]] constexpr observer_ptr make_observer(T* ptr) noexcept +{ return observer_ptr(ptr); } } // namespace nn diff --git a/include/neuralnet.cpp/core_profile.hpp b/include/neuralnet.cpp/core_profile.hpp index 65a5175c..c9b8b9b4 100644 --- a/include/neuralnet.cpp/core_profile.hpp +++ b/include/neuralnet.cpp/core_profile.hpp @@ -63,7 +63,8 @@ inline const bool g_enabled = ::nn::detail::env_flag_raw("NN_PROFILE") || !::nn::detail::env_string_raw("NN_PROFILE_TRACE").empty(); inline const std::string g_trace_path = ::nn::detail::env_string_raw("NN_PROFILE_TRACE"); -[[nodiscard]] inline bool enabled() noexcept { return g_enabled; } +[[nodiscard]] inline bool enabled() noexcept +{ return g_enabled; } // ── 聚合桶 ────────────────────────────────────────────────────────────────── struct Bucket @@ -86,8 +87,7 @@ namespace detail { // key = 字符串字面量地址（见文件头"命名约束"） -using StageMap = std::unordered_map>; +using StageMap = std::unordered_map>; // Chrome trace 逐事件（仅 NN_PROFILE_TRACE 时记录） struct Event @@ -102,7 +102,7 @@ struct Event struct State { std::mutex mu; - StageMap stages; // stage → op → 桶（op=="" = stage 本身） + StageMap stages; // stage → op → 桶（op=="" = stage 本身） std::unordered_map device; // per-key 归因（"shader_key@label@stage" → 桶）：device 表按 label 聚合后 // 看不出 fused@fwd.blk 这类大桶由哪些结构贡献——此表把逐 dispatch 的 @@ -124,8 +124,8 @@ inline State& state() // 入口）天然不重复计数。 struct Frame { - const char* stage; // 静态字符串；op 帧沿用最内层 stage - const char* op; // nullptr = 纯 stage 帧 + const char* stage; // 静态字符串；op 帧沿用最内层 stage + const char* op; // nullptr = 纯 stage 帧 std::uint64_t start_ns; std::uint64_t child_ns; }; @@ -138,18 +138,17 @@ inline std::vector& stack() [[nodiscard]] inline std::uint64_t now_ns() noexcept { - return static_cast( - std::chrono::duration_cast( - std::chrono::steady_clock::now().time_since_epoch()).count()); + return static_cast(std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count()); } -inline void accumulate(const char* stage, const char* op, - std::uint64_t incl_ns, std::uint64_t excl_ns) +inline void accumulate(const char* stage, const char* op, std::uint64_t incl_ns, + std::uint64_t excl_ns) { State& s = state(); std::lock_guard lock(s.mu); - Bucket& b = s.stages[stage][op != nullptr ? std::string_view(op) - : std::string_view()]; + Bucket& b = s.stages[stage][op != nullptr ? std::string_view(op) : std::string_view()]; ++b.calls; b.incl_ns += incl_ns; b.excl_ns += excl_ns; @@ -159,13 +158,13 @@ inline void accumulate(const char* stage, const char* op, // ── Chrome trace ──────────────────────────────────────────────────────────── inline std::uint32_t current_tid() noexcept { - thread_local const std::uint32_t hashed = static_cast( - std::hash{}(std::this_thread::get_id())); + thread_local const std::uint32_t hashed = + static_cast(std::hash{}(std::this_thread::get_id())); return hashed; } -inline void trace_event(const char* name, const char* cat, - std::uint64_t ts_us, std::uint64_t dur_us) +inline void trace_event(const char* name, const char* cat, std::uint64_t ts_us, + std::uint64_t dur_us) { if (g_trace_path.empty()) return; State& s = state(); @@ -173,7 +172,7 @@ inline void trace_event(const char* name, const char* cat, s.trace.push_back(Event{name, cat, ts_us, dur_us, current_tid()}); } -} // namespace detail +} // namespace detail // ── 动态名驻留（层号等非字面量名）────────────────────────────────────────── // 返回 "prefix#idx" 的稳定 const char*（首次生成后驻留，之后零分配）； @@ -181,7 +180,7 @@ inline void trace_event(const char* name, const char* cat, [[nodiscard]] inline const char* interned(const char* prefix, std::size_t idx) { static std::mutex mu; - static std::deque pool; // 地址稳定 + static std::deque pool; // 地址稳定 static std::unordered_map seen; std::string key = std::string(prefix) + "#" + std::to_string(idx); std::lock_guard lock(mu); @@ -208,7 +207,7 @@ inline void trace_event(const char* name, const char* cat, [[nodiscard]] inline const char* interned_pair(const char* a, const char* b) { static std::mutex mu; - static std::deque pool; // 地址稳定 + static std::deque pool; // 地址稳定 static std::unordered_map seen; std::string key = std::string(a) + "@" + b; std::lock_guard lock(mu); @@ -226,13 +225,11 @@ inline void trace_event(const char* name, const char* cat, // 必须拷进稳定池才能当聚合表 key（见文件头"命名约束"）。suffix 单列是为了 // 让调用点零分配传细 tile 变体（detail=shader_name、suffix="#f"）。 // 产出 "detail[suffix]@label[@stage]"；仅剖析开启时被调用。 -[[nodiscard]] inline const char* interned_detail(std::string_view detail, - const char* suffix, - const char* label, - const char* stage) +[[nodiscard]] inline const char* interned_detail(std::string_view detail, const char* suffix, + const char* label, const char* stage) { static std::mutex mu; - static std::deque pool; // 地址稳定 + static std::deque pool; // 地址稳定 static std::unordered_map seen; std::string key(detail); if (suffix != nullptr && suffix[0] != '\0') key += suffix; @@ -254,7 +251,8 @@ inline void trace_event(const char* name, const char* cat, // ── RAII 计时器 ───────────────────────────────────────────────────────────── -namespace detail { +namespace detail +{ // 收掉栈顶帧并记账（Scope::stop 与 Phase::stop 共用）。 // 幂等契约：name != nullptr 表示"本句柄仍持有一个未结束的段"；收口后置空 @@ -369,8 +367,7 @@ class Op // ── 设备侧回填（Vulkan timestamp query 解析后调用）───────────────────────── // detail：可选的 per-key 归因串（已驻留，见 interned_detail）——非空时同时 // 计入 device_keys 表；device 表口径不变。 -inline void device_add(const char* name, std::uint64_t ns, - const char* detail = nullptr) +inline void device_add(const char* name, std::uint64_t ns, const char* detail = nullptr) { if (!g_enabled) return; detail::State& s = detail::state(); @@ -391,10 +388,8 @@ inline void device_add(const char* name, std::uint64_t ns, // ── 宏（__LINE__ 拼接出唯一变量名）────────────────────────────────────────── #define NN_PROF_CONCAT_(a, b) a##b #define NN_PROF_CONCAT(a, b) NN_PROF_CONCAT_(a, b) -#define NN_PROF_SCOPE(name) \ - ::nn::prof::Scope NN_PROF_CONCAT(nn_prof_scope_, __LINE__)(name) -#define NN_PROF_OP(name) \ - ::nn::prof::Op NN_PROF_CONCAT(nn_prof_op_, __LINE__)(name) +#define NN_PROF_SCOPE(name) ::nn::prof::Scope NN_PROF_CONCAT(nn_prof_scope_, __LINE__)(name) +#define NN_PROF_OP(name) ::nn::prof::Op NN_PROF_CONCAT(nn_prof_op_, __LINE__)(name) // ── 报告 ──────────────────────────────────────────────────────────────────── // out : 打表文件（默认 stderr） @@ -407,7 +402,7 @@ void dump(std::FILE* out, std::uint64_t wall_ns = 0, bool reset = false); // 时也会自动调用它）。 void write_trace(const std::string& path); -} // namespace nn::prof +} // namespace nn::prof // ── dump / write_trace 实现（放头尾部，header-only 需 inline）─────────────── namespace nn::prof @@ -437,17 +432,16 @@ inline void write_trace(const std::string& path) bool first = true; for (const auto& e : events) { - std::fprintf(f, "%s{\"name\":\"%s\",\"cat\":\"%s\",\"ph\":\"X\"," - "\"ts\":%llu,\"dur\":%llu,\"pid\":0,\"tid\":%u}", - first ? "" : ",", e.name, e.cat, - static_cast(e.ts_us), + std::fprintf(f, + "%s{\"name\":\"%s\",\"cat\":\"%s\",\"ph\":\"X\"," + "\"ts\":%llu,\"dur\":%llu,\"pid\":0,\"tid\":%u}", + first ? "" : ",", e.name, e.cat, static_cast(e.ts_us), static_cast(e.dur_us), e.tid); first = false; } std::fprintf(f, "]}\n"); std::fclose(f); - std::fprintf(stderr, "[profile] trace 已写入 %s（%zu 事件）\n", - path.c_str(), events.size()); + std::fprintf(stderr, "[profile] trace 已写入 %s（%zu 事件）\n", path.c_str(), events.size()); } inline void dump(std::FILE* out, std::uint64_t wall_ns, bool reset) @@ -470,8 +464,7 @@ inline void dump(std::FILE* out, std::uint64_t wall_ns, bool reset) // ── 表 1：stage（inclusive = 阶段墙钟；% 相对墙钟，嵌套阶段可读）────── std::fprintf(out, "[profile] ── 阶段构成（inclusive 墙钟）──\n"); - std::fprintf(out, "[profile] %-28s %10s %8s %12s\n", - "stage", "ms", "%wall", "calls"); + std::fprintf(out, "[profile] %-28s %10s %8s %12s\n", "stage", "ms", "%wall", "calls"); // 收集 stage 行（op 键为空串 = stage 本身），按耗时降序 std::vector> stage_rows; std::uint64_t stage_sum = 0; @@ -483,25 +476,22 @@ inline void dump(std::FILE* out, std::uint64_t wall_ns, bool reset) stage_sum += it->second.incl_ns; } std::sort(stage_rows.begin(), stage_rows.end(), - [](const auto& a, const auto& b) - { return a.second.incl_ns > b.second.incl_ns; }); + [](const auto& a, const auto& b) { return a.second.incl_ns > b.second.incl_ns; }); for (const auto& [st, b] : stage_rows) { const double ms = static_cast(b.incl_ns) / 1e6; // 分母 = 调用方给的墙钟；没给墙钟时退回 stage 合计 - const double denom = wall_ns > 0 ? static_cast(wall_ns) - : static_cast(stage_sum); - const double pct = denom > 0 - ? 100.0 * static_cast(b.incl_ns) / denom : 0.0; - std::fprintf(out, "[profile] %-28s %10.2f %7.1f%% %12llu\n", - std::string(st).c_str(), ms, pct, - static_cast(b.calls)); + const double denom = + wall_ns > 0 ? static_cast(wall_ns) : static_cast(stage_sum); + const double pct = denom > 0 ? 100.0 * static_cast(b.incl_ns) / denom : 0.0; + std::fprintf(out, "[profile] %-28s %10.2f %7.1f%% %12llu\n", std::string(st).c_str(), ms, + pct, static_cast(b.calls)); } // ── 表 2：op 跨 stage 汇总（exclusive = 净归属）──────────────────────── std::fprintf(out, "[profile] ── 原语净归属（exclusive，跨 stage 汇总）──\n"); - std::fprintf(out, "[profile] %-24s %10s %12s %10s %8s\n", - "op", "ms", "calls", "avg_us", "%wall"); + std::fprintf(out, "[profile] %-24s %10s %12s %10s %8s\n", "op", "ms", "calls", "avg_us", + "%wall"); std::unordered_map op_rollup; for (const auto& [st, ops] : stages) { @@ -516,27 +506,24 @@ inline void dump(std::FILE* out, std::uint64_t wall_ns, bool reset) if (b.max_ns > r.max_ns) r.max_ns = b.max_ns; } } - std::vector> op_rows(op_rollup.begin(), - op_rollup.end()); + std::vector> op_rows(op_rollup.begin(), op_rollup.end()); std::sort(op_rows.begin(), op_rows.end(), - [](const auto& a, const auto& b) - { return a.second.excl_ns > b.second.excl_ns; }); + [](const auto& a, const auto& b) { return a.second.excl_ns > b.second.excl_ns; }); for (const auto& [op, b] : op_rows) { const double ms = static_cast(b.excl_ns) / 1e6; - const double avg_us = b.calls > 0 - ? static_cast(b.excl_ns) / 1e3 / static_cast(b.calls) : 0.0; - const double pctw = wall_ns > 0 - ? 100.0 * static_cast(b.excl_ns) / static_cast(wall_ns) : 0.0; - std::fprintf(out, "[profile] %-24s %10.2f %12llu %10.1f %7.1f%%\n", - std::string(op).c_str(), ms, - static_cast(b.calls), avg_us, pctw); + const double avg_us = + b.calls > 0 ? static_cast(b.excl_ns) / 1e3 / static_cast(b.calls) : 0.0; + const double pctw = + wall_ns > 0 ? 100.0 * static_cast(b.excl_ns) / static_cast(wall_ns) + : 0.0; + std::fprintf(out, "[profile] %-24s %10.2f %12llu %10.1f %7.1f%%\n", std::string(op).c_str(), + ms, static_cast(b.calls), avg_us, pctw); } // ── 表 3：stage × op 明细（top 40，exclusive 降序）──────────────────── std::fprintf(out, "[profile] ── stage × op 明细（exclusive，top 40）──\n"); - std::fprintf(out, "[profile] %-28s %-20s %10s %12s\n", - "stage", "op", "ms", "calls"); + std::fprintf(out, "[profile] %-28s %-20s %10s %12s\n", "stage", "op", "ms", "calls"); std::vector> mix_rows; for (const auto& [st, ops] : stages) { @@ -547,16 +534,14 @@ inline void dump(std::FILE* out, std::uint64_t wall_ns, bool reset) } } std::sort(mix_rows.begin(), mix_rows.end(), - [](const auto& a, const auto& b) - { return a.second.excl_ns > b.second.excl_ns; }); + [](const auto& a, const auto& b) { return a.second.excl_ns > b.second.excl_ns; }); for (std::size_t i = 0; i < mix_rows.size() && i < 40; ++i) { const auto& [key, b] = mix_rows[i]; const auto pos = key.find("::"); const std::string st = key.substr(0, pos); std::fprintf(out, "[profile] %-28s %-20s %10.2f %12llu\n", - st.empty() ? "(step 外)" : st.c_str(), - key.substr(pos + 2).c_str(), + st.empty() ? "(step 外)" : st.c_str(), key.substr(pos + 2).c_str(), static_cast(b.excl_ns) / 1e6, static_cast(b.calls)); } @@ -565,8 +550,7 @@ inline void dump(std::FILE* out, std::uint64_t wall_ns, bool reset) if (!device.empty()) { std::fprintf(out, "[profile] ── GPU 设备侧 kernel（timestamp query）──\n"); - std::fprintf(out, "[profile] %-28s %10s %12s %10s\n", - "kernel", "ms", "calls", "avg_us"); + std::fprintf(out, "[profile] %-28s %10s %12s %10s\n", "kernel", "ms", "calls", "avg_us"); std::uint64_t dev_sum = 0; std::vector> dev_rows; for (const auto& [name, b] : device) @@ -574,24 +558,22 @@ inline void dump(std::FILE* out, std::uint64_t wall_ns, bool reset) dev_rows.emplace_back(name, b); dev_sum += b.total_ns; } - std::sort(dev_rows.begin(), dev_rows.end(), - [](const auto& a, const auto& b) + std::sort(dev_rows.begin(), dev_rows.end(), [](const auto& a, const auto& b) { return a.second.total_ns > b.second.total_ns; }); for (const auto& [name, b] : dev_rows) { - std::fprintf(out, "[profile] %-28s %10.2f %12llu %10.1f\n", - std::string(name).c_str(), - static_cast(b.total_ns) / 1e6, - static_cast(b.calls), - b.calls > 0 ? static_cast(b.total_ns) / 1e3 - / static_cast(b.calls) : 0.0); + std::fprintf( + out, "[profile] %-28s %10.2f %12llu %10.1f\n", std::string(name).c_str(), + static_cast(b.total_ns) / 1e6, static_cast(b.calls), + b.calls > 0 ? static_cast(b.total_ns) / 1e3 / static_cast(b.calls) + : 0.0); } if (wall_ns > 0) { - std::fprintf(out, "[profile] GPU busy%% = %.1f%%（设备 kernel %.2f ms / 墙钟 %.2f ms）\n", + std::fprintf(out, + "[profile] GPU busy%% = %.1f%%（设备 kernel %.2f ms / 墙钟 %.2f ms）\n", 100.0 * static_cast(dev_sum) / static_cast(wall_ns), - static_cast(dev_sum) / 1e6, - static_cast(wall_ns) / 1e6); + static_cast(dev_sum) / 1e6, static_cast(wall_ns) / 1e6); } // ── 表 5：per-key 归因（key@label@stage，逐结构细分 top 60）──────── @@ -601,32 +583,29 @@ inline void dump(std::FILE* out, std::uint64_t wall_ns, bool reset) if (!device_keys.empty()) { std::fprintf(out, "[profile] ── 设备侧 per-key 归因（key@label@stage，top 60）──\n"); - std::fprintf(out, "[profile] %-52s %10s %12s %10s\n", - "key@label@stage", "ms", "calls", "avg_us"); - std::vector> key_rows( - device_keys.begin(), device_keys.end()); - std::sort(key_rows.begin(), key_rows.end(), - [](const auto& a, const auto& b) + std::fprintf(out, "[profile] %-52s %10s %12s %10s\n", "key@label@stage", "ms", "calls", + "avg_us"); + std::vector> key_rows(device_keys.begin(), + device_keys.end()); + std::sort(key_rows.begin(), key_rows.end(), [](const auto& a, const auto& b) { return a.second.total_ns > b.second.total_ns; }); for (std::size_t i = 0; i < key_rows.size() && i < 60; ++i) { const auto& [name, b] = key_rows[i]; std::fprintf(out, "[profile] %-52s %10.2f %12llu %10.1f\n", - std::string(name).c_str(), - static_cast(b.total_ns) / 1e6, + std::string(name).c_str(), static_cast(b.total_ns) / 1e6, static_cast(b.calls), - b.calls > 0 ? static_cast(b.total_ns) / 1e3 - / static_cast(b.calls) : 0.0); + b.calls > 0 ? static_cast(b.total_ns) / 1e3 / + static_cast(b.calls) + : 0.0); } if (key_rows.size() > 60) - std::fprintf(out, "[profile]（共 %zu 条，仅显示 top 60）\n", - key_rows.size()); + std::fprintf(out, "[profile]（共 %zu 条，仅显示 top 60）\n", key_rows.size()); } } if (wall_ns > 0) - std::fprintf(out, "[profile] 墙钟合计 %.2f ms\n", - static_cast(wall_ns) / 1e6); + std::fprintf(out, "[profile] 墙钟合计 %.2f ms\n", static_cast(wall_ns) / 1e6); std::fprintf(out, "[profile] ════════════════════════\n"); std::fflush(out); @@ -644,4 +623,4 @@ inline void dump(std::FILE* out, std::uint64_t wall_ns, bool reset) } } -} // namespace nn::prof +} // namespace nn::prof diff --git a/include/neuralnet.cpp/dataset_samples.hpp b/include/neuralnet.cpp/dataset_samples.hpp index 17eb61e6..7757cfc6 100644 --- a/include/neuralnet.cpp/dataset_samples.hpp +++ b/include/neuralnet.cpp/dataset_samples.hpp @@ -197,16 +197,14 @@ class TextSampleSet // 只统计目标位置的 loss 有效数，不物化 batch。 // Global-batch 训练需要先知道整个 update 的归一化单位，再拆分 micro-batch。 - [[nodiscard]] Result - valid_target_count(std::span indices) const + [[nodiscard]] Result valid_target_count(std::span indices) const { std::size_t valid = 0; for (std::size_t idx : indices) { if (idx >= offsets_.size()) - NN_FAIL("TextSampleSet::valid_target_count: 样本下标越界（" + - std::to_string(idx) + " >= " + - std::to_string(offsets_.size()) + "）"); + NN_FAIL("TextSampleSet::valid_target_count: 样本下标越界（" + std::to_string(idx) + + " >= " + std::to_string(offsets_.size()) + "）"); const std::size_t offset = offsets_[idx]; const std::size_t valid_len = std::min(cfg_.seq_len, corpus_.token_flow.size() - offset); diff --git a/include/neuralnet.cpp/domain_cnn.hpp b/include/neuralnet.cpp/domain_cnn.hpp index add643cf..aeec8666 100644 --- a/include/neuralnet.cpp/domain_cnn.hpp +++ b/include/neuralnet.cpp/domain_cnn.hpp @@ -40,8 +40,8 @@ struct CnnConvSpec struct CnnConfig { std::size_t in_channels = 1; // 输入通道数（MNIST=1） - std::size_t in_size = 28; // 输入空间尺寸（方形） - std::size_t pool = 2; // 每个卷积后的 MaxPool 窗口（0=无池化） + std::size_t in_size = 28; // 输入空间尺寸（方形） + std::size_t pool = 2; // 每个卷积后的 MaxPool 窗口（0=无池化） std::vector convs; // 卷积层列表 // 展平后的全连接头（首元素 = 展平后第一隐藏层宽度，末位 = 类别数）。 // 例如 {120, 10} → Linear(flatten→120) + ReLU + Linear(120→10)。 @@ -52,7 +52,7 @@ struct CnnConfig // （BatchNorm 下即“逐特征跨 batch”的经典 Conv→BN 语义）；Head = 全连接头 // 隐藏层 Linear→Norm→ReLU（MLP 同款）。默认值由 4 epoch A/B 实测选定。 // 旧模型文件缺 norm_place 键 → spec 读回 None → 不加（旧布局逐位一致）。 - NormType norm_type = NormType::BatchNorm; + NormType norm_type = NormType::BatchNorm; NormPlace norm_place = NormPlace::Conv; // 模型级精度配置（P2-3）：默认全 F32（与历史逐位一致）。`--f16` 下由调用方 // 置 profile_f16() → Conv2D 权重/梯度按 param、融合算子按 compute 走 f16。 @@ -70,20 +70,15 @@ inline const std::vector MNIST_CNN_CONVS = { inline const std::vector MNIST_CNN_FC = {120, 10}; // ── 卷积输出尺寸（不含池化） ───────────────────────────────────────────── -[[nodiscard]] inline std::size_t conv_out_size( - std::size_t in, std::size_t k, std::size_t s, std::size_t p) -{ - return (in + 2 * p - k) / s + 1; -} +[[nodiscard]] inline std::size_t conv_out_size(std::size_t in, std::size_t k, std::size_t s, + std::size_t p) +{ return (in + 2 * p - k) / s + 1; } // ── 构建 CNN 模型（CnnConfig 版本，推荐使用） ──────────────────────────── -[[nodiscard]] inline Result build_cnn_model( - ComputeEngine& engine, const CnnConfig& cfg) +[[nodiscard]] inline Result build_cnn_model(ComputeEngine& engine, const CnnConfig& cfg) { - if (cfg.convs.empty()) - NN_FAIL("CNN: convs must not be empty"); - if (cfg.fc_dims.size() < 2) - NN_FAIL("CNN: fc_dims must have at least 2 elements"); + if (cfg.convs.empty()) NN_FAIL("CNN: convs must not be empty"); + if (cfg.fc_dims.size() < 2) NN_FAIL("CNN: fc_dims must have at least 2 elements"); if (cfg.in_channels == 0 || cfg.in_size == 0) NN_FAIL("CNN: in_channels/in_size must be positive"); if (cfg.norm_place != NormPlace::None && cfg.norm_place != NormPlace::Conv && @@ -105,8 +100,7 @@ inline const std::vector MNIST_CNN_FC = {120, 10}; if (cv.kernel > h + 2 * cv.padding || cv.kernel > w + 2 * cv.padding) NN_FAIL("CNN: kernel larger than input spatial size"); { - auto r = model.add(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); NN_TRY_CHECK(r); } h = conv_out_size(h, cv.kernel, cv.stride, cv.padding); @@ -131,12 +125,10 @@ inline const std::vector MNIST_CNN_FC = {120, 10}; } const std::size_t flattened = c * h * w; - if (flattened == 0) - NN_FAIL("CNN: flattened size is zero"); + if (flattened == 0) NN_FAIL("CNN: flattened size is zero"); // 全连接头：Linear(flatten → fc_dims[0]) → [Norm] → ReLU → ... // → Linear(→ num_classes)；Norm 只挂隐藏层，最后一层保持原始 logits。 - const bool norm_head = (cfg.norm_place == NormPlace::Head || - cfg.norm_place == NormPlace::Both); + const bool norm_head = (cfg.norm_place == NormPlace::Head || cfg.norm_place == NormPlace::Both); { auto r = model.add(flattened, cfg.fc_dims[0]); NN_TRY_CHECK(r); @@ -168,23 +160,19 @@ inline const std::vector MNIST_CNN_FC = {120, 10}; // ── 构造 CNN ModelSpec ──────────────────────────────────────────────────── // fc_dims 复用 layer_dims 字段（type 区分，与 MLP 互斥）。 -[[nodiscard]] inline ModelSpec make_cnn_spec( - std::size_t in_channels, - std::size_t in_size, - std::size_t pool, - const std::vector& convs, - const std::vector& fc_dims, - NormType norm_type = NormType::BatchNorm, - NormPlace norm_place = NormPlace::Conv) +[[nodiscard]] inline ModelSpec +make_cnn_spec(std::size_t in_channels, std::size_t in_size, std::size_t pool, + const std::vector& convs, const std::vector& fc_dims, + NormType norm_type = NormType::BatchNorm, NormPlace norm_place = NormPlace::Conv) { ModelSpec spec; - auto &c = spec.emplace(); - c.norm_type = norm_type; - c.norm_place = norm_place; + auto& c = spec.emplace(); + c.norm_type = norm_type; + c.norm_place = norm_place; c.cnn_in_channels = in_channels; - c.cnn_in_size = in_size; - c.cnn_pool = pool; - c.layer_dims = fc_dims; // 复用为 CNN 全连接头 + c.cnn_in_size = in_size; + c.cnn_pool = pool; + c.layer_dims = fc_dims; // 复用为 CNN 全连接头 for (const auto& cv : convs) { c.cnn_channels.push_back(cv.out_channels); @@ -198,22 +186,20 @@ inline const std::vector MNIST_CNN_FC = {120, 10}; // ── 从 ModelSpec 还原 CnnConfig ────────────────────────────────────────── [[nodiscard]] inline Result cnn_config_from_spec(const ModelSpec& spec) { - const auto *c = spec.as(); - if (c == nullptr) - NN_FAIL("cnn_config_from_spec: not a CNN spec"); + const auto* c = spec.as(); + if (c == nullptr) NN_FAIL("cnn_config_from_spec: not a CNN spec"); if (c->cnn_channels.size() != c->cnn_kernels.size() || c->cnn_channels.size() != c->cnn_strides.size() || c->cnn_channels.size() != c->cnn_paddings.size()) NN_FAIL("CNN spec: conv vectors length mismatch"); - if (c->cnn_channels.empty()) - NN_FAIL("CNN spec: no conv layers"); + if (c->cnn_channels.empty()) NN_FAIL("CNN spec: no conv layers"); CnnConfig cfg; cfg.in_channels = c->cnn_in_channels; - cfg.in_size = c->cnn_in_size; - cfg.pool = c->cnn_pool; - cfg.fc_dims = c->layer_dims; - cfg.norm_type = c->norm_type; - cfg.norm_place = c->norm_place; + cfg.in_size = c->cnn_in_size; + cfg.pool = c->cnn_pool; + cfg.fc_dims = c->layer_dims; + cfg.norm_type = c->norm_type; + cfg.norm_place = c->norm_place; if (cfg.norm_place != NormPlace::None && cfg.norm_place != NormPlace::Conv && cfg.norm_place != NormPlace::Head && cfg.norm_place != NormPlace::Both) NN_FAIL("CNN spec: norm_place must be none/conv/head/both (final is ViT-only)"); @@ -232,18 +218,16 @@ inline const std::vector MNIST_CNN_FC = {120, 10}; // ── 从 ModelSpec 构建 CNN 模型（用于从二进制文件加载时自动还原） ───────── // precision：模型级精度配置（P2-3）。**精度是运行期配置、不进 ModelSpec** //（与 GPT 一致：GptConfig::precision 同样只在构建期传）——由调用方显式给。 -[[nodiscard]] inline Result build_cnn_model_from_spec( - ComputeEngine& engine, const ModelSpec& spec, - const PrecisionProfile& precision = PrecisionProfile{}) +[[nodiscard]] inline Result +build_cnn_model_from_spec(ComputeEngine& engine, const ModelSpec& spec, + const PrecisionProfile& precision = PrecisionProfile{}) { auto cfg_r = cnn_config_from_spec(spec); NN_TRY_CHECK(cfg_r); cfg_r->precision = precision; auto model = build_cnn_model(engine, *cfg_r); - if (model) - model->set_spec(spec); // 记录架构规格，供 load_model 校验 + if (model) model->set_spec(spec); // 记录架构规格，供 load_model 校验 return model; } } // namespace nn - diff --git a/include/neuralnet.cpp/domain_gpt.hpp b/include/neuralnet.cpp/domain_gpt.hpp index 0837df38..d1a5d82b 100644 --- a/include/neuralnet.cpp/domain_gpt.hpp +++ b/include/neuralnet.cpp/domain_gpt.hpp @@ -17,25 +17,27 @@ #include "model_spec.hpp" #include -namespace nn { +namespace nn +{ // ── GPT 默认超参数 ──────────────────────────────────────────────────────── -inline constexpr std::size_t GPT_VOCAB_SIZE = 10000; -inline constexpr std::size_t GPT_D_MODEL = 128; -inline constexpr std::size_t GPT_NUM_HEADS = 4; -inline constexpr std::size_t GPT_D_FF = 512; -inline constexpr std::size_t GPT_NUM_LAYERS = 4; -inline constexpr std::size_t GPT_SEQ_LEN = 256; +inline constexpr std::size_t GPT_VOCAB_SIZE = 10000; +inline constexpr std::size_t GPT_D_MODEL = 128; +inline constexpr std::size_t GPT_NUM_HEADS = 4; +inline constexpr std::size_t GPT_D_FF = 512; +inline constexpr std::size_t GPT_NUM_LAYERS = 4; +inline constexpr std::size_t GPT_SEQ_LEN = 256; // ── GPT 配置结构体 ──────────────────────────────────────────────────────── // 将 build_gpt_model 的多个位置参数收拢为一个结构体，避免调用方签名过长。 -struct GptConfig { - std::size_t vocab_size = GPT_VOCAB_SIZE; - std::size_t d_model = GPT_D_MODEL; - std::size_t seq_len = GPT_SEQ_LEN; - std::size_t num_heads = GPT_NUM_HEADS; - std::size_t d_ff = GPT_D_FF; - std::size_t num_layers = GPT_NUM_LAYERS; +struct GptConfig +{ + std::size_t vocab_size = GPT_VOCAB_SIZE; + std::size_t d_model = GPT_D_MODEL; + std::size_t seq_len = GPT_SEQ_LEN; + std::size_t num_heads = GPT_NUM_HEADS; + std::size_t d_ff = GPT_D_FF; + std::size_t num_layers = GPT_NUM_LAYERS; PosEncodingType pos_enc = PosEncodingType::Learned; ActivationType activation = ActivationType::GeLU; NormType norm_type = NormType::LayerNorm; @@ -49,108 +51,87 @@ struct GptConfig { // GPTModel 是一个单一 Layer（内含 TokenEmb+PosEmb+N×GPTBlock+LN+LM Head）， // 作为 Model 的唯一层。engine 决定权重张量驻留设备。 // cfg.pos_enc 控制位置编码：Learned(默认)/Sinusoidal/ALiBi -[[nodiscard]] inline Result build_gpt_model( - ComputeEngine& engine, const GptConfig& cfg) +[[nodiscard]] inline Result build_gpt_model(ComputeEngine& engine, const GptConfig& cfg) { if (cfg.d_model == 0 || cfg.num_heads == 0 || cfg.seq_len == 0 || cfg.vocab_size == 0) NN_FAIL("GPT model parameters must be positive"); - if (cfg.d_ff == 0 || cfg.num_layers == 0) - NN_FAIL("GPT d_ff and num_layers must be positive"); - if (cfg.d_model % cfg.num_heads != 0) - NN_FAIL("GPT d_model must be divisible by num_heads"); - if (cfg.n_head_kv != 0 && (cfg.n_head_kv > cfg.num_heads || - cfg.num_heads % cfg.n_head_kv != 0)) + if (cfg.d_ff == 0 || cfg.num_layers == 0) NN_FAIL("GPT d_ff and num_layers must be positive"); + if (cfg.d_model % cfg.num_heads != 0) NN_FAIL("GPT d_model must be divisible by num_heads"); + if (cfg.n_head_kv != 0 && (cfg.n_head_kv > cfg.num_heads || cfg.num_heads % cfg.n_head_kv != 0)) NN_FAIL("GPT n_head_kv must be a divisor of num_heads (0 = MHA)"); Model model(engine); { - 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, cfg.precision, - cfg.subln, cfg.n_head_kv, cfg.tie_embeddings); + 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, + cfg.precision, cfg.subln, cfg.n_head_kv, cfg.tie_embeddings); NN_TRY_CHECK(r); } return model; } // ── 构建 GPT 模型（位置参数版本，转发至 GptConfig 版本） ─────── -[[nodiscard]] inline Result build_gpt_model( - ComputeEngine& engine, - std::size_t vocab_size = GPT_VOCAB_SIZE, - std::size_t d_model = GPT_D_MODEL, - std::size_t seq_len = GPT_SEQ_LEN, - std::size_t num_heads = GPT_NUM_HEADS, - std::size_t d_ff = GPT_D_FF, - std::size_t num_layers = GPT_NUM_LAYERS, - PosEncodingType pos_enc_type = PosEncodingType::Learned, - ActivationType activation = ActivationType::GeLU, - NormType norm_type = NormType::LayerNorm, - PrecisionProfile precision = PrecisionProfile{}, - bool subln = false, - std::size_t n_head_kv = 0, - bool tie_embeddings = false) +[[nodiscard]] inline Result +build_gpt_model(ComputeEngine& engine, std::size_t vocab_size = GPT_VOCAB_SIZE, + std::size_t d_model = GPT_D_MODEL, std::size_t seq_len = GPT_SEQ_LEN, + std::size_t num_heads = GPT_NUM_HEADS, std::size_t d_ff = GPT_D_FF, + std::size_t num_layers = GPT_NUM_LAYERS, + PosEncodingType pos_enc_type = PosEncodingType::Learned, + ActivationType activation = ActivationType::GeLU, + NormType norm_type = NormType::LayerNorm, + PrecisionProfile precision = PrecisionProfile{}, bool subln = false, + std::size_t n_head_kv = 0, bool tie_embeddings = false) { - return build_gpt_model(engine, GptConfig{ - vocab_size, d_model, seq_len, num_heads, d_ff, num_layers, - pos_enc_type, activation, norm_type, precision, subln, n_head_kv, - tie_embeddings}); + return build_gpt_model(engine, GptConfig{vocab_size, d_model, seq_len, num_heads, d_ff, + num_layers, pos_enc_type, activation, norm_type, + precision, subln, n_head_kv, tie_embeddings}); } // ── 从 ModelSpec 构建 GPT 模型 ────────────────────────────────────────── // 用于从二进制文件加载时自动还原架构。 // 统一的 GPTModel 通过 spec.pos_encoding 区分 Learned/Sinusoidal/ALiBi 模式， // 因此无需为 ALiBi 单独提供构建函数。 -[[nodiscard]] inline Result build_gpt_model_from_spec( - ComputeEngine& engine, const ModelSpec &spec, - PrecisionProfile precision = PrecisionProfile{}) +[[nodiscard]] inline Result +build_gpt_model_from_spec(ComputeEngine& engine, const ModelSpec& spec, + PrecisionProfile precision = PrecisionProfile{}) { // 接受 GPT 类型（含 ALiBi/RoPE 变体；旧 ALiBi_GPT 数值 type 同归 GptSpec） - const auto *g = spec.as(); - if (g == nullptr) - NN_FAIL("Invalid ModelSpec type for GPT: expected GPT or ALiBi_GPT"); + const auto* g = spec.as(); + if (g == nullptr) NN_FAIL("Invalid ModelSpec type for GPT: expected GPT or ALiBi_GPT"); // ── 三值（T1_58）权重由**规格**决定（权威来源，docs 21 §4.7/§4.8.3）── // 与 build_mnist_mlp_model 的加载路径同款：保存过的三值模型必须还原成 // BitLinear（参数形状相同、语义不同），命令行传进来的 precision 会被规格覆盖。 - if (g->weight_quant == WeightQuant::T1_58) - precision.param = Precision::T1_58; - auto model = build_gpt_model( - engine, - g->vocab_size, g->d_model, g->seq_len, - g->num_heads, g->d_ff, g->num_layers, - g->pos_encoding, g->activation, g->norm_type, precision, - g->subln, g->n_head_kv, g->tie_embeddings); - if (model) - model->set_spec(spec); // 记录架构规格，供 load_model 校验 + if (g->weight_quant == WeightQuant::T1_58) precision.param = Precision::T1_58; + auto model = + build_gpt_model(engine, g->vocab_size, g->d_model, g->seq_len, g->num_heads, g->d_ff, + g->num_layers, g->pos_encoding, g->activation, g->norm_type, precision, + g->subln, g->n_head_kv, g->tie_embeddings); + if (model) model->set_spec(spec); // 记录架构规格，供 load_model 校验 return model; } // ── 构造 GPT ModelSpec ────────────────────────────────────────────────── -[[nodiscard]] inline ModelSpec make_gpt_spec( - 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_encoding = PosEncodingType::Learned, - ActivationType activation = ActivationType::GeLU, - NormType norm_type = NormType::LayerNorm, - bool subln = false, - std::size_t n_head_kv = 0, - bool tie_embeddings = false) +[[nodiscard]] inline ModelSpec +make_gpt_spec(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_encoding = PosEncodingType::Learned, + ActivationType activation = ActivationType::GeLU, + NormType norm_type = NormType::LayerNorm, bool subln = false, + std::size_t n_head_kv = 0, bool tie_embeddings = false) { ModelSpec spec; - auto &g = spec.emplace(); - g.vocab_size = vocab_size; - g.d_model = d_model; - g.seq_len = seq_len; - g.num_heads = num_heads; - g.d_ff = d_ff; - g.num_layers = num_layers; + auto& g = spec.emplace(); + g.vocab_size = vocab_size; + g.d_model = d_model; + g.seq_len = seq_len; + g.num_heads = num_heads; + g.d_ff = d_ff; + g.num_layers = num_layers; g.pos_encoding = pos_encoding; - g.activation = activation; - g.norm_type = norm_type; - g.subln = subln; - g.n_head_kv = n_head_kv; + g.activation = activation; + g.norm_type = norm_type; + g.subln = subln; + g.n_head_kv = n_head_kv; g.tie_embeddings = tie_embeddings; return spec; } diff --git a/include/neuralnet.cpp/domain_mnist.hpp b/include/neuralnet.cpp/domain_mnist.hpp index 032b2538..cf77b3d0 100644 --- a/include/neuralnet.cpp/domain_mnist.hpp +++ b/include/neuralnet.cpp/domain_mnist.hpp @@ -21,7 +21,8 @@ #include "model_spec.hpp" #include "domain_cnn.hpp" -namespace nn { +namespace nn +{ // ── MNIST 常量 ────────────────────────────────────────────────────────────── inline constexpr std::size_t MNIST_INPUT_DIM = 784; @@ -29,17 +30,16 @@ inline constexpr std::size_t MNIST_NUM_CLASSES = 10; // ── MNIST Transformer (ViT) 默认超参数 ────────────────────────────────────── // 图像 28×28，patch_size=7 → 4×4=16 个 patch，patch_dim=49 -inline constexpr std::size_t MNIST_IMG_SIZE = 28; -inline constexpr std::size_t MNIST_PATCH_SIZE = 7; -inline constexpr std::size_t MNIST_TF_D_MODEL = 64; +inline constexpr std::size_t MNIST_IMG_SIZE = 28; +inline constexpr std::size_t MNIST_PATCH_SIZE = 7; +inline constexpr std::size_t MNIST_TF_D_MODEL = 64; inline constexpr std::size_t MNIST_TF_NUM_HEADS = 4; -inline constexpr std::size_t MNIST_TF_D_FF = 128; +inline constexpr std::size_t MNIST_TF_D_FF = 128; inline constexpr std::size_t MNIST_TF_NUM_LAYERS = 2; // ── 共享 CSV 行解析工具 ──────────────────────────────────────────────────── // 解析单行 CSV（逗号分隔的浮点数）为 std::vector -[[nodiscard]] inline Result> -parse_csv_line(const std::string &line) +[[nodiscard]] inline Result> parse_csv_line(const std::string& line) { std::vector values; std::stringstream ss(line); @@ -47,8 +47,7 @@ parse_csv_line(const std::string &line) while (std::getline(ss, token, ',')) { auto v = parse_number(token); - if (!v) - NN_FAIL("CSV 含无效数字 '" + token + "': " + v.error().message); + if (!v) NN_FAIL("CSV 含无效数字 '" + token + "': " + v.error().message); values.push_back(*v); } return values; @@ -59,14 +58,13 @@ parse_csv_line(const std::string &line) // 输出：(784, 1) 列向量 Matrix // 底层解析实现；上层代码应优先使用 load_image_tensor_from_csv_line // （上层统一用 Tensor）。 -[[nodiscard]] inline Result -load_image_from_csv_line(const std::string &csv_line) +[[nodiscard]] inline Result load_image_from_csv_line(const std::string& csv_line) { auto values = parse_csv_line(csv_line); NN_TRY_CHECK(values); if (values->size() != MNIST_INPUT_DIM) NN_FAIL("CSV 必须包含恰好 " + std::to_string(MNIST_INPUT_DIM) + - " 个值，实际: " + std::to_string(values->size())); + " 个值，实际: " + std::to_string(values->size())); Matrix img(MNIST_INPUT_DIM, 1); for (std::size_t i = 0; i < MNIST_INPUT_DIM; ++i) img.set_value_unchecked(i, 0, (*values)[i]); @@ -76,8 +74,8 @@ load_image_from_csv_line(const std::string &csv_line) // ── 从 CSV 行加载单张 MNIST 图片为 Tensor ──────────────────────────────── // 上层统一使用 Tensor：内部调用 load_image_from_csv_line 得到 Matrix 后， // 通过 engine.from_matrix 上传至对应设备。 -[[nodiscard]] inline Result -load_image_tensor_from_csv_line(const std::string &csv_line, ComputeEngine& engine) +[[nodiscard]] inline Result load_image_tensor_from_csv_line(const std::string& csv_line, + ComputeEngine& engine) { auto mat_r = load_image_from_csv_line(csv_line); NN_TRY_CHECK(mat_r); @@ -85,9 +83,8 @@ load_image_tensor_from_csv_line(const std::string &csv_line, ComputeEngine& engi } // 默认 MLP 网络架构：输入层 → 隐藏层1 → 隐藏层2 → 隐藏层3 → 输出层 -inline const std::vector MNIST_LAYER_DIMS = { - MNIST_INPUT_DIM, 512, 256, 128, 64, MNIST_NUM_CLASSES -}; +inline const std::vector MNIST_LAYER_DIMS = {MNIST_INPUT_DIM, 512, 256, 128, 64, + MNIST_NUM_CLASSES}; // ── 构建 MLP 模型 ────────────────────────────────────────────────────────── // 通过指定 ComputeEngine 创建同设备的权重张量；可指定自定义层维度。 @@ -102,13 +99,10 @@ inline const std::vector MNIST_LAYER_DIMS = { // 这个概念，param 槽对其无意义；不归一化会被能力校验以 // "层 ReLU 不支持 param=t1_58" 拦下，属误报）。 [[nodiscard]] inline Result build_mnist_mlp_model( - ComputeEngine& engine, - const std::vector &layer_dims = MNIST_LAYER_DIMS, - NormType norm_type = NormType::LayerNorm, - PrecisionProfile precision = PrecisionProfile{}) + ComputeEngine& engine, const std::vector& layer_dims = MNIST_LAYER_DIMS, + NormType norm_type = NormType::LayerNorm, PrecisionProfile precision = PrecisionProfile{}) { - if (layer_dims.size() < 2) - NN_FAIL("MLP layer_dims must have at least 2 elements"); + if (layer_dims.size() < 2) NN_FAIL("MLP layer_dims must have at least 2 elements"); const bool ternary = (precision.param == Precision::T1_58); PrecisionProfile plain = precision; if (ternary) plain.param = Precision::F32; @@ -117,7 +111,7 @@ inline const std::vector MNIST_LAYER_DIMS = { model.set_default_precision_profile(plain); // 须在 add 之前：权重按精度创建 for (std::size_t i = 0; i < layer_dims.size() - 1; ++i) { - std::size_t in_dim = layer_dims[i]; + std::size_t in_dim = layer_dims[i]; std::size_t out_dim = layer_dims[i + 1]; if (ternary) @@ -156,16 +150,11 @@ inline const std::vector MNIST_LAYER_DIMS = { // norm_place：Final（默认，原版 ViT 的 ln_f）/ None（显式关闭 = 旧结构）； // conv/head/both 属 CNN。旧模型文件缺 norm_place 键 → None → 不加。 [[nodiscard]] inline Result build_mnist_transformer_model( - ComputeEngine& engine, - std::size_t img_size = MNIST_IMG_SIZE, - std::size_t patch_size = MNIST_PATCH_SIZE, - std::size_t d_model = MNIST_TF_D_MODEL, - std::size_t num_heads = MNIST_TF_NUM_HEADS, - std::size_t d_ff = MNIST_TF_D_FF, - std::size_t num_layers = MNIST_TF_NUM_LAYERS, - PrecisionProfile precision = PrecisionProfile{}, - NormType norm_type = NormType::LayerNorm, - NormPlace norm_place = NormPlace::Final) + ComputeEngine& engine, std::size_t img_size = MNIST_IMG_SIZE, + std::size_t patch_size = MNIST_PATCH_SIZE, std::size_t d_model = MNIST_TF_D_MODEL, + std::size_t num_heads = MNIST_TF_NUM_HEADS, std::size_t d_ff = MNIST_TF_D_FF, + std::size_t num_layers = MNIST_TF_NUM_LAYERS, PrecisionProfile precision = PrecisionProfile{}, + NormType norm_type = NormType::LayerNorm, NormPlace norm_place = NormPlace::Final) { if (img_size % patch_size != 0) NN_FAIL("MNIST Transformer: img_size must be divisible by patch_size"); @@ -175,7 +164,7 @@ inline const std::vector MNIST_LAYER_DIMS = { NN_FAIL("MNIST Transformer: d_model must be divisible by num_heads"); if (norm_place != NormPlace::None && norm_place != NormPlace::Final) NN_FAIL("MNIST Transformer: norm_place must be none/final (conv/head/both are CNN-only)"); - const std::size_t grid_size = img_size / patch_size; + const std::size_t grid_size = img_size / patch_size; const std::size_t num_patches = grid_size * grid_size; Model model(engine); @@ -185,9 +174,8 @@ inline const std::vector MNIST_LAYER_DIMS = { NN_TRY_CHECK(r); } { - auto r = model.add(d_model, num_heads, d_ff, num_layers, - num_patches, norm_type, - norm_place == NormPlace::Final); + auto r = model.add(d_model, num_heads, d_ff, num_layers, num_patches, + norm_type, norm_place == NormPlace::Final); NN_TRY_CHECK(r); } { @@ -203,22 +191,19 @@ inline const std::vector MNIST_LAYER_DIMS = { // 默认 Final（原版 ViT = 块 pre-norm + 末端 ln_f，4 epoch A/B 实测优于无 final）。 // 旧模型文件缺 norm_place 键 → spec 读回 None → 不加（旧布局逐位一致）。 [[nodiscard]] inline ModelSpec make_mnist_transformer_spec( - std::size_t patch_size = MNIST_PATCH_SIZE, - std::size_t d_model = MNIST_TF_D_MODEL, - std::size_t num_heads = MNIST_TF_NUM_HEADS, - std::size_t d_ff = MNIST_TF_D_FF, - std::size_t num_layers = MNIST_TF_NUM_LAYERS, - NormType norm_type = NormType::LayerNorm, + std::size_t patch_size = MNIST_PATCH_SIZE, std::size_t d_model = MNIST_TF_D_MODEL, + std::size_t num_heads = MNIST_TF_NUM_HEADS, std::size_t d_ff = MNIST_TF_D_FF, + std::size_t num_layers = MNIST_TF_NUM_LAYERS, NormType norm_type = NormType::LayerNorm, NormPlace norm_place = NormPlace::Final) { ModelSpec spec; - auto &t = spec.emplace(); - t.d_model = d_model; - t.num_heads = num_heads; - t.d_ff = d_ff; + auto& t = spec.emplace(); + t.d_model = d_model; + t.num_heads = num_heads; + t.d_ff = d_ff; t.num_layers = num_layers; t.patch_size = patch_size; - t.norm_type = norm_type; + t.norm_type = norm_type; t.norm_place = norm_place; // vocab_size / seq_len 不用于 Transformer (MNIST ViT) return spec; @@ -226,18 +211,18 @@ inline const std::vector MNIST_LAYER_DIMS = { // ── 从 ModelSpec 构建模型 ───────────────────────────────────────────────── // 用于从二进制文件加载时自动还原架构（支持 MLP 和 Transformer） -[[nodiscard]] inline Result build_mnist_model_from_spec( - ComputeEngine& engine, const ModelSpec &spec, - PrecisionProfile precision = PrecisionProfile{}) +[[nodiscard]] inline Result +build_mnist_model_from_spec(ComputeEngine& engine, const ModelSpec& spec, + PrecisionProfile precision = PrecisionProfile{}) { // ── 三值（T1_58）权重由**规格**决定（不是 CLI 参数）─────────────── // 保存过的三值模型必须还原成 BitLinear（参数形状相同、语义不同）： // 规格是权威来源，CLI 的 precision.param 会被它覆盖。 // weight_quant 是各类型子类的字段（三值语义架构级，但归属各自的规格） WeightQuant wq = WeightQuant::None; - if (const auto *m = spec.as()) wq = m->weight_quant; - else if (const auto *t = spec.as()) wq = t->weight_quant; - else if (const auto *c = spec.as()) wq = c->weight_quant; + if (const auto* m = spec.as()) wq = m->weight_quant; + else if (const auto* t = spec.as()) wq = t->weight_quant; + else if (const auto* c = spec.as()) wq = c->weight_quant; if (wq == WeightQuant::T1_58) { if (!spec.is_mlp()) @@ -247,32 +232,27 @@ inline const std::vector MNIST_LAYER_DIMS = { precision.param = Precision::T1_58; } - if (const auto *m = spec.as()) + if (const auto* m = spec.as()) { - auto model = build_mnist_mlp_model(engine, m->layer_dims, m->norm_type, - precision); - if (model) - model->set_spec(spec); // 记录架构规格，供 load_model 校验 + auto model = build_mnist_mlp_model(engine, m->layer_dims, m->norm_type, precision); + if (model) model->set_spec(spec); // 记录架构规格，供 load_model 校验 return model; } - if (const auto *t = spec.as()) + if (const auto* t = spec.as()) { const std::size_t patch_size = t->patch_size != 0 ? t->patch_size : MNIST_PATCH_SIZE; - auto model = build_mnist_transformer_model( - engine, MNIST_IMG_SIZE, patch_size, - t->d_model, t->num_heads, t->d_ff, t->num_layers, precision, - t->norm_type, t->norm_place); - if (model) - model->set_spec(spec); // 记录架构规格，供 load_model 校验 + auto model = build_mnist_transformer_model(engine, MNIST_IMG_SIZE, patch_size, t->d_model, + t->num_heads, t->d_ff, t->num_layers, precision, + t->norm_type, t->norm_place); + if (model) model->set_spec(spec); // 记录架构规格，供 load_model 校验 return model; } if (spec.is_cnn()) { auto model = build_cnn_model_from_spec(engine, spec); - if (model) - model->set_spec(spec); // 记录架构规格，供 load_model 校验 + if (model) model->set_spec(spec); // 记录架构规格，供 load_model 校验 return model; } diff --git a/include/neuralnet.cpp/domain_rla.hpp b/include/neuralnet.cpp/domain_rla.hpp index a3dbdf34..f2437171 100644 --- a/include/neuralnet.cpp/domain_rla.hpp +++ b/include/neuralnet.cpp/domain_rla.hpp @@ -22,32 +22,32 @@ #include #include "domain_gpt.hpp" // GPT_VOCAB_SIZE 等共享默认超参 -namespace nn { +namespace nn +{ // ── RAPT 配置结构体 ───────────────────────────────────────────────────── -struct RAPTConfig { +struct RAPTConfig +{ std::size_t vocab_size = GPT_VOCAB_SIZE; - std::size_t d_model = GPT_D_MODEL; - std::size_t seq_len = GPT_SEQ_LEN; - std::size_t num_heads = GPT_NUM_HEADS; - std::size_t d_ff = GPT_D_FF; + std::size_t d_model = GPT_D_MODEL; + std::size_t seq_len = GPT_SEQ_LEN; + std::size_t num_heads = GPT_NUM_HEADS; + std::size_t d_ff = GPT_D_FF; std::size_t num_layers = GPT_NUM_LAYERS; - PosEncodingType pos_enc = PosEncodingType::RoPE; // RLA 约束：RoPE（或 ALiBi） + PosEncodingType pos_enc = PosEncodingType::RoPE; // RLA 约束：RoPE（或 ALiBi） ActivationType activation = ActivationType::GeLU; - NormType norm_type = NormType::LayerNorm; + NormType norm_type = NormType::LayerNorm; bool causal = true; // RAPT 是 causal LM PrecisionProfile precision; // 模型级精度配置（§9.1，默认全 F32） }; // ── 构建 RAPT 模型（推荐使用） ───────────────────────────────────────── -[[nodiscard]] inline Result build_rapt_model( - ComputeEngine& engine, const RAPTConfig& cfg) +[[nodiscard]] inline Result build_rapt_model(ComputeEngine& engine, const RAPTConfig& cfg) { - if (cfg.d_model == 0 || cfg.num_heads == 0 || cfg.seq_len == 0 || - cfg.vocab_size == 0 || cfg.d_ff == 0 || cfg.num_layers == 0) + if (cfg.d_model == 0 || cfg.num_heads == 0 || cfg.seq_len == 0 || cfg.vocab_size == 0 || + cfg.d_ff == 0 || cfg.num_layers == 0) NN_FAIL("RAPT model parameters must be positive"); - if (cfg.d_model % cfg.num_heads != 0) - NN_FAIL("RAPT d_model must be divisible by num_heads"); + if (cfg.d_model % cfg.num_heads != 0) NN_FAIL("RAPT d_model must be divisible by num_heads"); // RLA 强约束：RoPE 需偶数 d_k（旋转块）；v1 只支持 RoPE if (cfg.pos_enc != PosEncodingType::RoPE) NN_FAIL("RAPT v1 requires pos_enc == RoPE (RLA constraint)"); @@ -55,57 +55,49 @@ struct RAPTConfig { NN_FAIL("RAPT requires even d_k (RoPE) — d_model/num_heads must be even"); Model model(engine); { - 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, cfg.causal, cfg.precision); + 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, cfg.causal, cfg.precision); NN_TRY_CHECK(r); } return model; } // ── 从 ModelSpec 构建 RAPT 模型 ───────────────────────────────────────── -[[nodiscard]] inline Result build_rapt_model_from_spec( - ComputeEngine& engine, const ModelSpec& spec, - PrecisionProfile precision = PrecisionProfile{}) +[[nodiscard]] inline Result +build_rapt_model_from_spec(ComputeEngine& engine, const ModelSpec& spec, + PrecisionProfile precision = PrecisionProfile{}) { - const auto *r = spec.as(); - if (r == nullptr) - NN_FAIL("Invalid ModelSpec type for RAPT: expected RAPT"); - auto model = build_rapt_model(engine, RAPTConfig{ - r->vocab_size, r->d_model, r->seq_len, - r->num_heads, r->d_ff, r->num_layers, - r->pos_encoding, r->activation, r->norm_type, /*causal=*/true, precision}); - if (model) - model->set_spec(spec); + const auto* r = spec.as(); + if (r == nullptr) NN_FAIL("Invalid ModelSpec type for RAPT: expected RAPT"); + auto model = + build_rapt_model(engine, RAPTConfig{r->vocab_size, r->d_model, r->seq_len, r->num_heads, + r->d_ff, r->num_layers, r->pos_encoding, r->activation, + r->norm_type, /*causal=*/true, precision}); + if (model) model->set_spec(spec); return model; } // ── 构造 RAPT ModelSpec ───────────────────────────────────────────────── -[[nodiscard]] inline ModelSpec make_rapt_spec( - 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_encoding = PosEncodingType::RoPE, - ActivationType activation = ActivationType::GeLU, - NormType norm_type = NormType::LayerNorm) +[[nodiscard]] inline ModelSpec make_rapt_spec(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_encoding = PosEncodingType::RoPE, + ActivationType activation = ActivationType::GeLU, + NormType norm_type = NormType::LayerNorm) { ModelSpec spec; - auto &r = spec.emplace(); - r.vocab_size = vocab_size; - r.d_model = d_model; - r.seq_len = seq_len; - r.num_heads = num_heads; - r.d_ff = d_ff; - r.num_layers = num_layers; + auto& r = spec.emplace(); + r.vocab_size = vocab_size; + r.d_model = d_model; + r.seq_len = seq_len; + r.num_heads = num_heads; + r.d_ff = d_ff; + r.num_layers = num_layers; r.pos_encoding = pos_encoding; - r.activation = activation; - r.norm_type = norm_type; + r.activation = activation; + r.norm_type = norm_type; return spec; } } // namespace nn - diff --git a/include/neuralnet.cpp/expr_dsl.hpp b/include/neuralnet.cpp/expr_dsl.hpp index 6a301b48..df2818ee 100644 --- a/include/neuralnet.cpp/expr_dsl.hpp +++ b/include/neuralnet.cpp/expr_dsl.hpp @@ -45,9 +45,7 @@ namespace nn::dsl // ⚠ 探针门：默认构建（未 -DNN_ENABLE_PROBES=ON）里恒返回 false 且**不求值 getenv** // → 调用点被编译期折叠，"NN_F16_DEBUG" 之类的诊断开关不会在正常训练里生效。 inline bool env_flag(const char* name) -{ - return ::nn::detail::probe_flag_raw(name); -} +{ return ::nn::detail::probe_flag_raw(name); } } // namespace nn::dsl #include #include @@ -129,21 +127,13 @@ struct SpecBuilder return expr::input(static_cast(inputs.size() - 1)); } ExprOperand add_input_linear(const Tensor& t, bool dedup = true) - { - return add_input(t, expr::linear(), dedup); - } + { return add_input(t, expr::linear(), dedup); } ExprOperand add_input_rotate(const Tensor& t, std::uint32_t block) - { - return add_input(t, expr::rotate_half(block, /*negate_first_half=*/true)); - } + { return add_input(t, expr::rotate_half(block, /*negate_first_half=*/true)); } ExprOperand add_input_rowmod(const Tensor& t, std::uint32_t mod) - { - return add_input(t, expr::row_mod(mod)); - } + { return add_input(t, expr::row_mod(mod)); } ExprOperand add_input_rowaccess(const Tensor& t, std::uint32_t offset, std::uint32_t mod) - { - return add_input(t, expr::row_access(offset, mod)); - } + { return add_input(t, expr::row_access(offset, mod)); } // 分组归约视图输入：输出 (G,N)、输入 (G*R,N)，out[r][c] = Σ/max over // t(row,col)`（见 @@ -152,7 +142,7 @@ struct SpecBuilder { ExprView v; v.kind = static_cast(is_max ? ExprViewKind::GroupedReduceMax - : ExprViewKind::GroupedReduceSum); + : ExprViewKind::GroupedReduceSum); v.param = R; return add_input(t, v); } @@ -177,19 +167,17 @@ struct SpecBuilder } // 广播视图：输入本身已是 (rows,1)/(1,cols) 小向量，按行/列广播 ExprOperand add_input_rowbroadcast(const Tensor& t) - { - return add_input(t, expr::row_broadcast()); - } + { return add_input(t, expr::row_broadcast()); } ExprOperand add_input_colbroadcast(const Tensor& t) - { - return add_input(t, expr::col_broadcast()); - } + { return add_input(t, expr::col_broadcast()); } ExprOperand add_instr(ExprOp op, ExprOperand a, ExprOperand b = {}, ExprOperand c = {}) { ExprInstr in; in.op = static_cast(op); in.dst = static_cast(spec.num_regs++); - in.a = a; in.b = b; in.c = c; + in.a = a; + in.b = b; + in.c = c; spec.instrs.push_back(in); return expr::reg(in.dst); } @@ -210,18 +198,15 @@ struct SpecBuilder // 模板求值路径 → 编译器内联 + 与其它逐元素路径共用同一并行门控。 // ══════════════════════════════════════════════════════════════════════════ template -concept ColsAwareEval = requires(const T& t, std::size_t i, std::size_t cols) -{ +concept ColsAwareEval = requires (const T& t, std::size_t i, std::size_t cols) { { t.eval(i, cols) } -> std::convertible_to; }; template [[nodiscard]] Scalar eval_with_cols(const E& e, std::size_t i, std::size_t cols) { - if constexpr (ColsAwareEval) - return static_cast(e.eval(i, cols)); - else - return static_cast(e.eval(i)); + if constexpr (ColsAwareEval) return static_cast(e.eval(i, cols)); + else return static_cast(e.eval(i)); } // ══════════════════════════════════════════════════════════════════════════ @@ -247,13 +232,11 @@ template // 524288——按全局门控会**整段串行**（归一化层将只用 ~1 核），故本路径 // 取更低的专属门控。 // ══════════════════════════════════════════════════════════════════════════ -template -inline void eval_into_span(const E& e, Span d_span, std::size_t cols) noexcept +template inline void eval_into_span(const E& e, Span d_span, std::size_t cols) noexcept { Scalar* d = d_span.data(); const std::size_t n = d_span.size(); - if (n == 0) - return; + if (n == 0) return; if (n < kDslParallelThreshold) { NN_VECTORIZE_PRAGMA @@ -269,19 +252,20 @@ inline void eval_into_span(const E& e, Span d_span, std::size_t cols) noexcept // 1.7→1.9ms）——隔离内核的收益不会传递到层；再调整必须先有层级别证据。 // 并行区固定开销实测 ~100µs（32 线程）/ ~3~16µs（≤16 线程）。 constexpr std::size_t TARGET_CHUNK = std::size_t{1} << 16; - const std::size_t n_chunks = std::clamp(n / TARGET_CHUNK, std::size_t{1}, - std::max(hw, 1) * 4); + const std::size_t n_chunks = + std::clamp(n / TARGET_CHUNK, std::size_t{1}, std::max(hw, 1) * 4); const std::size_t base = n / n_chunks; - const std::size_t rem = n % n_chunks; - nn::parallel_for_samples(n_chunks, [&e, d, cols, base, rem](std::size_t c) noexcept - { - const std::size_t off = c * base + std::min(c, rem); - const std::size_t len = base + (c < rem ? 1 : 0); - Scalar* dc = d + off; - NN_VECTORIZE_PRAGMA - for (std::size_t k = 0; k < len; ++k) - dc[k] = eval_with_cols(e, off + k, cols); - }); + const std::size_t rem = n % n_chunks; + nn::parallel_for_samples(n_chunks, + [&e, d, cols, base, rem](std::size_t c) noexcept + { + const std::size_t off = c * base + std::min(c, rem); + const std::size_t len = base + (c < rem ? 1 : 0); + Scalar* dc = d + off; + NN_VECTORIZE_PRAGMA + for (std::size_t k = 0; k < len; ++k) + dc[k] = eval_with_cols(e, off + k, cols); + }); } // ══════════════════════════════════════════════════════════════════════════ // CpuViewCache — 叶子节点的 CPU 数据视图缓存 @@ -311,13 +295,11 @@ struct CpuViewCache void cache_cpu_view(const Tensor& t) { - if (!t.is_cpu()) - return; + if (!t.is_cpu()) return; // 符号张量（默认构造：未绑定、无存储）——AOT 自登记需要"从表达式类型 // 折叠出结构"，那条路径只看 to_spec 折出的视图/指令，不会 eval；此处 // 直接返回，避免在无存储张量上取 CPU 视图（原会 NN_ASSERT 硬失败）。 - if (!t.valid()) - return; + if (!t.valid()) return; if (t.precision() == Precision::F16) { auto m = std::make_shared(t.rows(), t.cols()); @@ -396,8 +378,7 @@ struct RotateHalfRef : CpuViewCache const std::size_t c = i % cols, r = i / cols; const std::size_t blk = block; const std::size_t rl = r % blk; - const std::size_t rr = (r / blk) * blk - + ((rl < blk / 2) ? (rl + blk / 2) : (rl - blk / 2)); + const std::size_t rr = (r / blk) * blk + ((rl < blk / 2) ? (rl + blk / 2) : (rl - blk / 2)); const Scalar v = at(rr * cols + c); return (rl < blk / 2) ? -v : v; } @@ -434,7 +415,8 @@ struct RowAccessRef : CpuViewCache std::uint32_t mod = 1; // 同上 RowAccessRef(Tensor tt, std::uint32_t off, std::uint32_t m) - : t(std::move(tt)), offset(off), mod(m) { cache_cpu_view(t); } + : t(std::move(tt)), offset(off), mod(m) + { cache_cpu_view(t); } // 符号实例（AOT 自登记）：结构由类型决定，视图参数取中性值（不进 key） RowAccessRef() = default; @@ -462,15 +444,12 @@ struct RowAccessRef : CpuViewCache // 而"结构 = 表达式类型"是 AOT 自登记（`FusedAnchor` 默认构造符号实例） // 的前提——运行期 bool 会让两种视图同类型、结构不可枚举。 // r_len（R）是运行期视图参数，不进 expr_spec_key（见 expr_spec.hpp）。 -template -struct GroupedReduceRef : CpuViewCache +template struct GroupedReduceRef : CpuViewCache { Tensor t; std::uint32_t r_len = 1; // 组长度 R（运行期视图参数，不进 expr_spec_key） - GroupedReduceRef(Tensor tt, std::uint32_t R) - : t(std::move(tt)), r_len(R) - { cache_cpu_view(t); } + GroupedReduceRef(Tensor tt, std::uint32_t R) : t(std::move(tt)), r_len(R) { cache_cpu_view(t); } // 符号实例（AOT 自登记）：结构由类型决定，R 取中性值（不进 key） GroupedReduceRef() = default; @@ -483,8 +462,14 @@ struct GroupedReduceRef : CpuViewCache for (std::uint32_t tt = 0; tt < r_len; ++tt) { const Scalar v = at((base + tt) * cols + c); - if (IsMax) { if (v > acc) acc = v; } - else { acc += v; } + if (IsMax) + { + if (v > acc) acc = v; + } + else + { + acc += v; + } } return acc; } @@ -510,8 +495,7 @@ using GroupedReduceMaxRef = GroupedReduceRef; // 分流到引擎 eval_expr，由引擎实现归约视图/指令语义），故 eval 仅需满足 // nn::Expression 概念约束，返回占位值。 // ══════════════════════════════════════════════════════════════════════════ -template -struct ReduceViewRef +template struct ReduceViewRef { // 归约轴与算子：行归约 → 向量长度 rows；列归约 → 长度 cols static constexpr bool reduces_rows = @@ -527,14 +511,13 @@ struct ReduceViewRef // 预计算逐字节一致（铁律 8：并行只是把独立槽分给不同线程）。 // GPU/scan 不使用本状态（to_spec 只读 t 与 Kind）。 mutable std::vector vec_{}; - mutable const Scalar *v_ = nullptr; + mutable const Scalar* v_ = nullptr; [[nodiscard]] constexpr Scalar eval(std::size_t) const noexcept { return Scalar{0}; } // CPU 模板路径（2 参）：归约视图按行/列广播 [[nodiscard]] Scalar eval(std::size_t i, std::size_t cols) const noexcept { - if (v_ == nullptr) - return Scalar{0}; + if (v_ == nullptr) return Scalar{0}; return reduces_rows ? v_[i / cols] : v_[i % cols]; } ExprOperand to_spec(SpecBuilder& b) const { return b.add_reduce_input(t, Kind); } @@ -583,8 +566,7 @@ struct MatmulRef // 头映射）。预绑定只是优化，回退不引入正确性风险。 if (a_kv_ratio > 1 || b_kv_ratio > 1) NN_FAIL("dsl matmul prepare: kv_ratio > 1 走解释器路径（CPU 预绑定不支持头映射）"); - if (c_data_ != nullptr) - return {}; // 幂等：同一棵树重复求值只物化一次 + if (c_data_ != nullptr) return {}; // 幂等：同一棵树重复求值只物化一次 // ── 精度安全（§7.2 f32 参考）──────────────────────────────────── // C 固定以 P=F32 物化，而操作数可能带 **f16 存储**（profile_f16 的 // Linear 权重/激活）：内层 GEMM 的 f32 路径按 f32 存储直接读 @@ -609,13 +591,11 @@ struct MatmulRef } constexpr Precision P = Precision::F32; Result c = (batch > 1) - ? eng.batched_matmul(*pa, *pb, batch, transA, transB, Scalar{1}, P) - : eng.matmul(*pa, *pb, transA, transB, P); + ? eng.batched_matmul(*pa, *pb, batch, transA, transB, Scalar{1}, P) + : eng.matmul(*pa, *pb, transA, transB, P); NN_TRY_CHECK(c); - if (!c->is_cpu()) - NN_FAIL("dsl matmul prepare: result not on CPU"); - if (c->precision() != Precision::F32) - NN_FAIL("dsl matmul prepare: result not f32"); + if (!c->is_cpu()) NN_FAIL("dsl matmul prepare: result not on CPU"); + if (c->precision() != Precision::F32) NN_FAIL("dsl matmul prepare: result not f32"); if (env_flag("NN_F16_DEBUG")) { const auto sp = TensorAccess::cpu_matrix(*c).span(); @@ -631,10 +611,9 @@ struct MatmulRef "[dbg][prepare] C=(%zux%zu) A=(%zux%zu)%s B=(%zux%zu)%s " "max(C)=%.6g%s <<< 异常\n", c->rows(), c->cols(), pa->rows(), pa->cols(), - pa->precision() == Precision::F16 ? "f16" : "f32", - pb->rows(), pb->cols(), - pb->precision() == Precision::F16 ? "f16" : "f32", - mx, bad ? " NONFINITE" : ""); + pa->precision() == Precision::F16 ? "f16" : "f32", pb->rows(), + pb->cols(), pb->precision() == Precision::F16 ? "f16" : "f32", mx, + bad ? " NONFINITE" : ""); } c_cache_ = std::move(*c); c_data_ = TensorAccess::cpu_matrix(c_cache_).span().data(); @@ -644,9 +623,7 @@ struct MatmulRef [[nodiscard]] constexpr Scalar eval(std::size_t) const noexcept { return Scalar{0}; } // CPU 模板路径（2 参）：读预绑定的 C（形状与输出网格一致） [[nodiscard]] Scalar eval(std::size_t i, std::size_t) const noexcept - { - return c_data_ != nullptr ? c_data_[i] : Scalar{0}; - } + { return c_data_ != nullptr ? c_data_[i] : Scalar{0}; } ExprOperand to_spec(SpecBuilder& sb) const { // 显式固定登记顺序（A 先 B 后），保证跨编译器确定（同 Binary 约定）。 @@ -661,8 +638,12 @@ struct MatmulRef // transA=1 → A.rows()/batch（A 存储 (batch*K, M)，每批 (K, M)） const std::uint32_t k = static_cast( transA ? (a.rows() / (batch / (a_kv_ratio ? a_kv_ratio : 1))) : a.cols()); - sb.spec.matmul = MatmulSpec{ai, bi, transA ? std::uint8_t{1} : std::uint8_t{0}, - transB ? std::uint8_t{1} : std::uint8_t{0}, k, batch}; + sb.spec.matmul = MatmulSpec{ai, + bi, + transA ? std::uint8_t{1} : std::uint8_t{0}, + transB ? std::uint8_t{1} : std::uint8_t{0}, + k, + batch}; // GQA（ratio>1）：把该操作数的**视图**标成 HeadGroup 并记下头分组比 // —— param 是运行期形状数据（不进 key），任意 n_rep 共用一份 shader。 if (a_kv_ratio > 1) @@ -688,9 +669,21 @@ struct MatmulRef // 用于位置相关掩码/偏置（causal select(Col > Row, -inf, 0)、ALiBi 等）。 // 走引擎 eval_expr（索引由引擎按网格推导），eval 仅满足概念约束。 // ══════════════════════════════════════════════════════════════════════════ -struct RowIdxLeaf { [[nodiscard]] constexpr Scalar eval(std::size_t) const noexcept { return Scalar{0}; } ExprOperand to_spec(SpecBuilder&) const { return expr::row(); } }; -struct ColIdxLeaf { [[nodiscard]] constexpr Scalar eval(std::size_t) const noexcept { return Scalar{0}; } ExprOperand to_spec(SpecBuilder&) const { return expr::col(); } }; -struct BatchIdxLeaf{ [[nodiscard]] constexpr Scalar eval(std::size_t) const noexcept { return Scalar{0}; } ExprOperand to_spec(SpecBuilder&) const { return expr::batch(); } }; +struct RowIdxLeaf +{ + [[nodiscard]] constexpr Scalar eval(std::size_t) const noexcept { return Scalar{0}; } + ExprOperand to_spec(SpecBuilder&) const { return expr::row(); } +}; +struct ColIdxLeaf +{ + [[nodiscard]] constexpr Scalar eval(std::size_t) const noexcept { return Scalar{0}; } + ExprOperand to_spec(SpecBuilder&) const { return expr::col(); } +}; +struct BatchIdxLeaf +{ + [[nodiscard]] constexpr Scalar eval(std::size_t) const noexcept { return Scalar{0}; } + ExprOperand to_spec(SpecBuilder&) const { return expr::batch(); } +}; // ══════════════════════════════════════════════════════════════════════════ // 视图叶子：标签行收集 / 按批次索引 @@ -737,9 +730,7 @@ struct BatchColRef [[nodiscard]] constexpr Scalar eval(std::size_t) const noexcept { return Scalar{0}; } ExprOperand to_spec(SpecBuilder& sb) const - { - return sb.add_input(t, expr::batch_col(per_batch_cols)); - } + { return sb.add_input(t, expr::batch_col(per_batch_cols)); } }; // 广播视图叶子：输入 (rows,1)/(1,cols) 小向量，按行/列广播参与算术。 @@ -750,8 +741,7 @@ struct BatchColRef // - 行广播：输入 (rows,1)，读第 i/cols 个标量 —— 需要"输出列数"，输入自身 // 推不出，故 1 参 eval(i)（无 cols）无法给出正确值，返回占位值；行广播 // 一律经 2 参 eval(i, cols)（eval_with_cols 自动选择）求值。 -template -struct BroadcastRef : CpuViewCache +template struct BroadcastRef : CpuViewCache { Tensor t; @@ -763,22 +753,17 @@ struct BroadcastRef : CpuViewCache { if constexpr (Kind == ExprViewKind::ColBroadcast) return at(i % cols); // 输入 (1,cols) → cols 即输出列数 - else - return Scalar{0}; // 行广播需输出列数，见 eval(i, out_cols) + else return Scalar{0}; // 行广播需输出列数，见 eval(i, out_cols) } [[nodiscard]] Scalar eval(std::size_t i, std::size_t out_cols) const { - if constexpr (Kind == ExprViewKind::ColBroadcast) - return at(i % cols); - else - return at(i / out_cols); // 输入 (rows,1)：读第 r 个标量 + if constexpr (Kind == ExprViewKind::ColBroadcast) return at(i % cols); + else return at(i / out_cols); // 输入 (rows,1)：读第 r 个标量 } ExprOperand to_spec(SpecBuilder& b) const { - if constexpr (Kind == ExprViewKind::RowBroadcast) - return b.add_input_rowbroadcast(t); - else - return b.add_input_colbroadcast(t); + if constexpr (Kind == ExprViewKind::RowBroadcast) return b.add_input_rowbroadcast(t); + else return b.add_input_colbroadcast(t); } }; @@ -794,8 +779,7 @@ using ColBroadcastRef = BroadcastRef; // 候选，误用在编译期即暴露。 // ══════════════════════════════════════════════════════════════════════════ template -concept DslExpr = nn::Expression && requires(SpecBuilder& b, const T& t) -{ +concept DslExpr = nn::Expression && requires (SpecBuilder& b, const T& t) { { t.to_spec(b) } -> std::convertible_to; }; @@ -803,23 +787,21 @@ concept DslExpr = nn::Expression && requires(SpecBuilder& b, const T& t) // 节点模板（每个都同时支持 CPU 求值 eval(i) 与 GPU 折叠 to_spec） // ══════════════════════════════════════════════════════════════════════════ -template -struct Unary +template struct Unary { C child; [[nodiscard]] auto eval(std::size_t i) const { return Op::apply(child.eval(i)); } [[nodiscard]] auto eval(std::size_t i, std::size_t cols) const { return Op::apply(eval_with_cols(child, i, cols)); } - ExprOperand to_spec(SpecBuilder& b) const - { return b.add_instr(Op::op_id(), child.to_spec(b)); } + ExprOperand to_spec(SpecBuilder& b) const { return b.add_instr(Op::op_id(), child.to_spec(b)); } }; // 比较算子（ops::Gt 等）使 eval 返回 bool → 该 Binary 是 BoolExpression， // 供 select() 使用；to_spec 仍折叠成单条比较指令（输出 0.0/1.0 到寄存器）。 -template -struct Binary +template struct Binary { - L l; R r; + L l; + R r; [[nodiscard]] auto eval(std::size_t i) const { return Op::apply(l.eval(i), r.eval(i)); } [[nodiscard]] auto eval(std::size_t i, std::size_t cols) const { return Op::apply(eval_with_cols(l, i, cols), eval_with_cols(r, i, cols)); } @@ -836,16 +818,17 @@ struct Binary }; // select(cond, then, else) —— cond 为 BoolExpression，then/else 为值表达式 -template -struct Select +template struct Select { - C cond; T then_e; E else_e; + C cond; + T then_e; + E else_e; [[nodiscard]] Scalar eval(std::size_t i) const { return cond.eval(i) ? then_e.eval(i) : else_e.eval(i); } [[nodiscard]] Scalar eval(std::size_t i, std::size_t cols) const { - return (eval_with_cols(cond, i, cols) != Scalar{0}) - ? eval_with_cols(then_e, i, cols) : eval_with_cols(else_e, i, cols); + return (eval_with_cols(cond, i, cols) != Scalar{0}) ? eval_with_cols(then_e, i, cols) + : eval_with_cols(else_e, i, cols); } ExprOperand to_spec(SpecBuilder& b) const { @@ -860,8 +843,7 @@ struct Select // 归约节点：对子表达式结果做归约（归约**指令**）→ 隐式标量向量（Reduce 操作数）。 // row_reduce_sum(exp(x)) 这类"表达式内部归约"；对输入 Tensor 直接归约请用 // ReduceViewRef（归约视图，GPU 融合更友好）。 -template -struct ReduceRef +template struct ReduceRef { // 归约轴与算子：行归约 → 向量长度 rows；列归约 → 长度 cols static constexpr bool reduces_rows = (Rop == ExprOp::RowSum || Rop == ExprOp::RowMax); @@ -874,14 +856,13 @@ struct ReduceRef // 对子表达式求值并归约到 vec_。槽内沿归约轴升序累加 → 与解释器 // replay_prefix 的累加顺序逐字节一致（铁律 8）。 mutable std::vector vec_{}; - mutable const Scalar *v_ = nullptr; + mutable const Scalar* v_ = nullptr; [[nodiscard]] constexpr Scalar eval(std::size_t) const noexcept { return Scalar{0}; } // CPU 模板路径（2 参）：归约指令按行/列广播 [[nodiscard]] Scalar eval(std::size_t i, std::size_t cols) const noexcept { - if (v_ == nullptr) - return Scalar{0}; + if (v_ == nullptr) return Scalar{0}; return reduces_rows ? v_[i / cols] : v_[i % cols]; } ExprOperand to_spec(SpecBuilder& b) const { return b.add_reduce_instr(Rop, child.to_spec(b)); } @@ -899,35 +880,24 @@ struct ReduceRef template inline constexpr bool has_reduction_v = false; // 归约叶子需走 eval_expr（模板求值无法表达全行/全列归约） -template -inline constexpr bool has_reduction_v> = true; +template inline constexpr bool has_reduction_v> = true; // 广播视图**不需要**走 eval_expr：它是纯索引映射（无跨元素依赖），2 参 // eval(i, cols) 即可在模板求值路径上正确求值（见 BroadcastRef 与 // eval_with_cols）→ 留在编译期模板路径（内联 + 并行）。 -template -inline constexpr bool has_reduction_v> = false; +template inline constexpr bool has_reduction_v> = false; // 分组归约视图：每输出元素独立归约 R 行、元素间零依赖 → 非归约（留 // 编译期模板路径；fail-safe 显式声明，防默认值漂移把它推进解释器） -template <> -inline constexpr bool has_reduction_v = false; -template <> -inline constexpr bool has_reduction_v = false; +template <> inline constexpr bool has_reduction_v = false; +template <> inline constexpr bool has_reduction_v = false; // matmul 叶子同样需走 eval_expr（matmul 预计算 + 逐元素链，引擎实现） -template <> -inline constexpr bool has_reduction_v = true; +template <> inline constexpr bool has_reduction_v = true; // 索引/收集叶子：行/列/批次下标与标签收集由引擎按网格推导 -template <> -inline constexpr bool has_reduction_v = true; -template <> -inline constexpr bool has_reduction_v = true; -template <> -inline constexpr bool has_reduction_v = true; -template <> -inline constexpr bool has_reduction_v = true; -template <> -inline constexpr bool has_reduction_v = true; -template <> -inline constexpr bool has_reduction_v = true; +template <> inline constexpr bool has_reduction_v = true; +template <> inline constexpr bool has_reduction_v = true; +template <> inline constexpr bool has_reduction_v = true; +template <> inline constexpr bool has_reduction_v = true; +template <> inline constexpr bool has_reduction_v = true; +template <> inline constexpr bool has_reduction_v = true; template inline constexpr bool has_reduction_v> = true; template @@ -935,32 +905,31 @@ inline constexpr bool has_reduction_v> = has_reduction_v; template inline constexpr bool has_reduction_v> = has_reduction_v || has_reduction_v; template -inline constexpr bool has_reduction_v> - = has_reduction_v || has_reduction_v || has_reduction_v; +inline constexpr bool has_reduction_v> = + has_reduction_v || has_reduction_v || has_reduction_v; // ══════════════════════════════════════════════════════════════════════════ // 叶子构造（普通写法入口） // ══════════════════════════════════════════════════════════════════════════ -[[nodiscard]] inline TensorRef leaf(Tensor t) { return TensorRef{std::move(t)}; } +[[nodiscard]] inline TensorRef leaf(Tensor t) +{ return TensorRef{std::move(t)}; } // 运行时标量参数（RParam）：值不限、结构固定 → 同结构共享融合 shader。 // 用于优化器超参（lr/eps/β）、偏差修正系数等每步会变但结构不变的标量。 -[[nodiscard]] inline constexpr RParamLeaf rparam(Scalar v) { return RParamLeaf{v}; } +[[nodiscard]] inline constexpr RParamLeaf rparam(Scalar v) +{ return RParamLeaf{v}; } // matmul 叶子：C = op(A,B)，折叠为 ExprSpec 的前置 matmul 段。 // 逐元素链（如 +bias、激活）自动与 matmul 融合成一个 kernel（GPU AOT）。 // batch：A/B 按 batch 垂直切分（batched_matmul 同布局），形状参数。 -[[nodiscard]] inline MatmulRef matmul(Tensor a, Tensor b, - bool transA = false, bool transB = false, +[[nodiscard]] inline MatmulRef matmul(Tensor a, Tensor b, bool transA = false, bool transB = false, std::uint32_t batch = 1) { return MatmulRef{std::move(a), std::move(b), transA, transB, batch}; } // GQA 变体：a_kv_ratio / b_kv_ratio > 1 时该操作数只存 batch/ratio 份行块 // （多个连续 query 头共享一个 KV 头；块映射见 docs/internal/22 §3.2）。 -[[nodiscard]] inline MatmulRef matmul_kv(Tensor a, Tensor b, - bool transA, bool transB, - std::uint32_t batch, - std::uint32_t a_kv_ratio, +[[nodiscard]] inline MatmulRef matmul_kv(Tensor a, Tensor b, bool transA, bool transB, + std::uint32_t batch, std::uint32_t a_kv_ratio, std::uint32_t b_kv_ratio) { MatmulRef m{std::move(a), std::move(b), transA, transB, batch}; @@ -970,9 +939,12 @@ inline constexpr bool has_reduction_v> } // 网格索引叶子：行号/列号/批次下标 -[[nodiscard]] inline RowIdxLeaf row() { return {}; } -[[nodiscard]] inline ColIdxLeaf col() { return {}; } -[[nodiscard]] inline BatchIdxLeaf batch() { return {}; } +[[nodiscard]] inline RowIdxLeaf row() +{ return {}; } +[[nodiscard]] inline ColIdxLeaf col() +{ return {}; } +[[nodiscard]] inline BatchIdxLeaf batch() +{ return {}; } // 视图叶子：标签行收集（稀疏 CE loss）/ 按批次取模索引（ALiBi 斜率）/ // 按 (batch,col) 切片（doc_ids） @@ -1004,70 +976,104 @@ inline constexpr bool has_reduction_v> // 对表达式结果归约 → 归约**指令**（重载按参数类型自动选择，见 ReduceRef）。 // 返回的归约叶子参与算术时自动按行/按列广播。 // ══════════════════════════════════════════════════════════════════════════ -[[nodiscard]] inline RowReduceSumRef row_reduce_sum(Tensor t) { return {std::move(t)}; } -[[nodiscard]] inline RowReduceMaxRef row_reduce_max(Tensor t) { return {std::move(t)}; } -[[nodiscard]] inline ColReduceSumRef col_reduce_sum(Tensor t) { return {std::move(t)}; } -[[nodiscard]] inline ColReduceMaxRef col_reduce_max(Tensor t) { return {std::move(t)}; } - -template [[nodiscard]] auto row_reduce_sum(const E& e) { return ReduceRef{e}; } -template [[nodiscard]] auto row_reduce_max(const E& e) { return ReduceRef{e}; } -template [[nodiscard]] auto col_reduce_sum(const E& e) { return ReduceRef{e}; } -template [[nodiscard]] auto col_reduce_max(const E& e) { return ReduceRef{e}; } +[[nodiscard]] inline RowReduceSumRef row_reduce_sum(Tensor t) +{ return {std::move(t)}; } +[[nodiscard]] inline RowReduceMaxRef row_reduce_max(Tensor t) +{ return {std::move(t)}; } +[[nodiscard]] inline ColReduceSumRef col_reduce_sum(Tensor t) +{ return {std::move(t)}; } +[[nodiscard]] inline ColReduceMaxRef col_reduce_max(Tensor t) +{ return {std::move(t)}; } + +template [[nodiscard]] auto row_reduce_sum(const E& e) +{ return ReduceRef{e}; } +template [[nodiscard]] auto row_reduce_max(const E& e) +{ return ReduceRef{e}; } +template [[nodiscard]] auto col_reduce_sum(const E& e) +{ return ReduceRef{e}; } +template [[nodiscard]] auto col_reduce_max(const E& e) +{ return ReduceRef{e}; } // 广播视图自由函数：gamma/beta (F,1) 按行广播、std_inv (1,B) 按列广播 -[[nodiscard]] inline RowBroadcastRef row_broadcast(Tensor t) { return {std::move(t)}; } -[[nodiscard]] inline ColBroadcastRef col_broadcast(Tensor t) { return {std::move(t)}; } +[[nodiscard]] inline RowBroadcastRef row_broadcast(Tensor t) +{ return {std::move(t)}; } +[[nodiscard]] inline ColBroadcastRef col_broadcast(Tensor t) +{ return {std::move(t)}; } // ══════════════════════════════════════════════════════════════════════════ // 一元函数（普通数学写法） // ══════════════════════════════════════════════════════════════════════════ -template [[nodiscard]] constexpr auto neg(const E& e) { return Unary{e}; } -template [[nodiscard]] auto abs(const E& e) { return Unary{e}; } -template [[nodiscard]] auto exp(const E& e) { return Unary{e}; } -template [[nodiscard]] auto log(const E& e) { return Unary{e}; } -template [[nodiscard]] auto sqrt(const E& e) { return Unary{e}; } -template [[nodiscard]] auto rsqrt(const E& e) { return Unary{e}; } -template [[nodiscard]] auto tanh(const E& e) { return Unary{e}; } +template [[nodiscard]] constexpr auto neg(const E& e) +{ return Unary{e}; } +template [[nodiscard]] auto abs(const E& e) +{ return Unary{e}; } +template [[nodiscard]] auto exp(const E& e) +{ return Unary{e}; } +template [[nodiscard]] auto log(const E& e) +{ return Unary{e}; } +template [[nodiscard]] auto sqrt(const E& e) +{ return Unary{e}; } +template [[nodiscard]] auto rsqrt(const E& e) +{ return Unary{e}; } +template [[nodiscard]] auto tanh(const E& e) +{ return Unary{e}; } // ══════════════════════════════════════════════════════════════════════════ // 二元函数：max / min / select // ══════════════════════════════════════════════════════════════════════════ -template [[nodiscard]] constexpr auto max(const L& l, const R& r) { return Binary{l, r}; } -template [[nodiscard]] constexpr auto max(const E& e, Scalar s) { return Binary{e, ConstLeaf{s}}; } -template [[nodiscard]] constexpr auto max(Scalar s, const E& e) { return Binary{ConstLeaf{s}, e}; } -template [[nodiscard]] constexpr auto min(const L& l, const R& r) { return Binary{l, r}; } -template [[nodiscard]] constexpr auto min(const E& e, Scalar s) { return Binary{e, ConstLeaf{s}}; } -template [[nodiscard]] constexpr auto min(Scalar s, const E& e) { return Binary{ConstLeaf{s}, e}; } +template +[[nodiscard]] constexpr auto max(const L& l, const R& r) +{ return Binary{l, r}; } +template [[nodiscard]] constexpr auto max(const E& e, Scalar s) +{ + return Binary{e, ConstLeaf{s}}; +} +template [[nodiscard]] constexpr auto max(Scalar s, const E& e) +{ return Binary{ConstLeaf{s}, e}; } +template +[[nodiscard]] constexpr auto min(const L& l, const R& r) +{ return Binary{l, r}; } +template [[nodiscard]] constexpr auto min(const E& e, Scalar s) +{ + return Binary{e, ConstLeaf{s}}; +} +template [[nodiscard]] constexpr auto min(Scalar s, const E& e) +{ return Binary{ConstLeaf{s}, e}; } /// relu(x) = max(x, 0) -template [[nodiscard]] auto relu(const E& e) { return nn::dsl::max(e, ConstLeaf{Scalar{0}}); } +template [[nodiscard]] auto relu(const E& e) +{ return nn::dsl::max(e, ConstLeaf{Scalar{0}}); } template [[nodiscard]] constexpr auto select(const C& c, const T& t, const E& e) { return Select{c, t, e}; } template [[nodiscard]] constexpr auto select(const C& c, const T& t, Scalar ev) -{ return Select{c, t, ConstLeaf{ev}}; } +{ + return Select{c, t, ConstLeaf{ev}}; +} template [[nodiscard]] constexpr auto select(const C& c, Scalar tv, const E& e) { return Select{c, ConstLeaf{tv}, e}; } template [[nodiscard]] constexpr auto select(const C& c, Scalar tv, Scalar ev) -{ return Select{c, ConstLeaf{tv}, ConstLeaf{ev}}; } +{ + return Select{c, ConstLeaf{tv}, ConstLeaf{ev}}; +} // ══════════════════════════════════════════════════════════════════════════ // 运算符（Sym 生成，覆盖 Expr×Expr / Expr×Scalar / Scalar×Expr） // 值运算（+ - * /）返回值表达式；比较（> < >= <= == !=）返回布尔表达式。 // ══════════════════════════════════════════════════════════════════════════ -#define NN_DSL_BINARY(OP, OpT) \ - template \ - [[nodiscard]] constexpr auto operator OP(const L& l, const R& r) \ - { return Binary{l, r}; } \ - template \ - [[nodiscard]] constexpr auto operator OP(const E& e, Scalar s) \ - { return Binary{e, ConstLeaf{s}}; } \ - template \ - [[nodiscard]] constexpr auto operator OP(Scalar s, const E& e) \ +#define NN_DSL_BINARY(OP, OpT) \ + template \ + [[nodiscard]] constexpr auto operator OP(const L& l, const R& r) \ + { return Binary{l, r}; } \ + template [[nodiscard]] constexpr auto operator OP(const E& e, Scalar s) \ + { \ + return Binary{e, ConstLeaf{s}}; \ + } \ + template [[nodiscard]] constexpr auto operator OP(Scalar s, const E& e) \ { return Binary{ConstLeaf{s}, e}; } NN_DSL_BINARY(+, ops::Add) @@ -1128,8 +1134,7 @@ struct GridInfer bool conflict = false; // 广播规则不成立（某维既不等、又都不为 1） bool unknown = false; // 有贡献者存在但形状推不出来（如行块不整除） - [[nodiscard]] static std::size_t merge_dim(std::size_t a, std::size_t b, - bool& conflict) + [[nodiscard]] static std::size_t merge_dim(std::size_t a, std::size_t b, bool& conflict) { if (a == b) return a; if (a == 1) return b; @@ -1140,40 +1145,63 @@ struct GridInfer void merge(std::size_t r, std::size_t c) { if (conflict || unknown) return; - if (!any) { rows = r; cols = c; any = true; return; } + if (!any) + { + rows = r; + cols = c; + any = true; + return; + } rows = merge_dim(rows, r, conflict); cols = merge_dim(cols, c, conflict); } }; // 缺省：不贡献（标量叶 / 网格索引叶 / 稀疏与批视图） -template -inline void grid_contrib(const E&, GridInfer&) {} +template inline void grid_contrib(const E&, GridInfer&) {} // 形状保持的叶子/视图 inline void grid_contrib(const TensorRef& e, GridInfer& g) -{ if (e.t.valid()) g.merge(e.t.rows(), e.t.cols()); else g.unknown = true; } +{ + if (e.t.valid()) g.merge(e.t.rows(), e.t.cols()); + else g.unknown = true; +} inline void grid_contrib(const RotateHalfRef& e, GridInfer& g) -{ if (e.t.valid()) g.merge(e.t.rows(), e.t.cols()); else g.unknown = true; } +{ + if (e.t.valid()) g.merge(e.t.rows(), e.t.cols()); + else g.unknown = true; +} inline void grid_contrib(const RowModRef& e, GridInfer& g) -{ if (e.t.valid()) g.merge(e.t.rows(), e.t.cols()); else g.unknown = true; } +{ + if (e.t.valid()) g.merge(e.t.rows(), e.t.cols()); + else g.unknown = true; +} inline void grid_contrib(const RowAccessRef& e, GridInfer& g) -{ if (e.t.valid()) g.merge(e.t.rows(), e.t.cols()); else g.unknown = true; } -template -inline void grid_contrib(const BroadcastRef& e, GridInfer& g) -{ if (e.t.valid()) g.merge(e.t.rows(), e.t.cols()); else g.unknown = true; } +{ + if (e.t.valid()) g.merge(e.t.rows(), e.t.cols()); + else g.unknown = true; +} +template inline void grid_contrib(const BroadcastRef& e, GridInfer& g) +{ + if (e.t.valid()) g.merge(e.t.rows(), e.t.cols()); + else g.unknown = true; +} // 归约视图：compute 下是广播读 → 贡献输入整网格（compute_reduce 要的也是 // 输入网格，两者同口径） -template -inline void grid_contrib(const ReduceViewRef& e, GridInfer& g) -{ if (e.t.valid()) g.merge(e.t.rows(), e.t.cols()); else g.unknown = true; } +template inline void grid_contrib(const ReduceViewRef& e, GridInfer& g) +{ + if (e.t.valid()) g.merge(e.t.rows(), e.t.cols()); + else g.unknown = true; +} // 分组归约：输出 (rows/R, cols) -template -inline void grid_contrib(const GroupedReduceRef& e, GridInfer& g) +template inline void grid_contrib(const GroupedReduceRef& e, GridInfer& g) { if (!e.t.valid() || e.r_len == 0 || e.t.rows() % e.r_len != 0) - { g.unknown = true; return; } + { + g.unknown = true; + return; + } g.merge(e.t.rows() / e.r_len, e.t.cols()); } @@ -1183,39 +1211,54 @@ inline void grid_contrib(const MatmulRef& e, GridInfer& g) const std::size_t batch = e.batch; const std::size_t ra = e.a_kv_ratio ? e.a_kv_ratio : 1; const std::size_t rb = e.b_kv_ratio ? e.b_kv_ratio : 1; - if (batch == 0 || !e.a.valid() || !e.b.valid() - || batch % ra != 0 || batch % rb != 0) - { g.unknown = true; return; } + if (batch == 0 || !e.a.valid() || !e.b.valid() || batch % ra != 0 || batch % rb != 0) + { + g.unknown = true; + return; + } const std::size_t a_blocks = batch / ra; // GQA：该操作数只存 batch/ratio 份行块 const std::size_t b_blocks = batch / rb; if (e.a.rows() % a_blocks != 0 || e.b.rows() % b_blocks != 0) - { g.unknown = true; return; } + { + g.unknown = true; + return; + } const std::size_t a_rows_per = e.a.rows() / a_blocks; const std::size_t b_rows_per = e.b.rows() / b_blocks; const std::size_t M = e.transA ? e.a.cols() : a_rows_per; const std::size_t N = e.transB ? b_rows_per : e.b.cols(); - if (M == 0 || N == 0) { g.unknown = true; return; } + if (M == 0 || N == 0) + { + g.unknown = true; + return; + } g.merge(batch * M, N); } // 复合节点：递进子节点合并 -template -inline void grid_contrib(const ReduceRef& e, GridInfer& g) +template inline void grid_contrib(const ReduceRef& e, GridInfer& g) { grid_contrib(e.child, g); } template inline void grid_contrib(const Unary& e, GridInfer& g) { grid_contrib(e.child, g); } template inline void grid_contrib(const Binary& e, GridInfer& g) -{ grid_contrib(e.l, g); grid_contrib(e.r, g); } +{ + grid_contrib(e.l, g); + grid_contrib(e.r, g); +} template inline void grid_contrib(const Select& e, GridInfer& g) -{ grid_contrib(e.cond, g); grid_contrib(e.then_e, g); grid_contrib(e.else_e, g); } +{ + grid_contrib(e.cond, g); + grid_contrib(e.then_e, g); + grid_contrib(e.else_e, g); +} // 推导出口：失败一律给出"可操作"的错误（点名入口 + 提示显式 rows/cols） template -[[nodiscard]] inline Result> -infer_grid_shape(const E& e, const char* entry) +[[nodiscard]] inline Result> infer_grid_shape(const E& e, + const char* entry) { GridInfer g; grid_contrib(e, g); @@ -1266,22 +1309,22 @@ template inline constexpr bool cpu_preparable_leaf_v = false; // （与 CpuEngine::eval_expr_into 的 expr_spec_reduce_axis != -1 守卫一致）。 template inline constexpr bool has_real_reduction_v = false; template inline constexpr bool has_real_reduction_v> = true; -template inline constexpr bool has_real_reduction_v> = true; +template +inline constexpr bool has_real_reduction_v> = true; template inline constexpr bool has_real_reduction_v> = has_real_reduction_v; template -inline constexpr bool has_real_reduction_v> - = has_real_reduction_v || has_real_reduction_v; +inline constexpr bool has_real_reduction_v> = + has_real_reduction_v || has_real_reduction_v; template -inline constexpr bool has_real_reduction_v> - = has_real_reduction_v || has_real_reduction_v || has_real_reduction_v; +inline constexpr bool has_real_reduction_v> = + has_real_reduction_v || has_real_reduction_v || has_real_reduction_v; // matmul：用引擎通用 GEMM 原语物化 C 后绑定指针，尾链交模板路径内联 template <> inline constexpr bool cpu_preparable_leaf_v = true; // 归约视图：按输出槽预计算行/列归约向量后交模板路径（语义与解释器一致） -template -inline constexpr bool cpu_preparable_leaf_v> = true; +template inline constexpr bool cpu_preparable_leaf_v> = true; // 归约指令：内部会先递归备好子表达式 template @@ -1295,19 +1338,18 @@ inline constexpr bool cpu_preparable_v = !has_reduction_v || cpu_preparable_l template inline constexpr bool cpu_preparable_v> = cpu_preparable_v; template -inline constexpr bool cpu_preparable_v> - = cpu_preparable_v && cpu_preparable_v; +inline constexpr bool cpu_preparable_v> = + cpu_preparable_v && cpu_preparable_v; template -inline constexpr bool cpu_preparable_v> - = cpu_preparable_v && cpu_preparable_v && cpu_preparable_v; +inline constexpr bool cpu_preparable_v> = + cpu_preparable_v && cpu_preparable_v && cpu_preparable_v; // 归约指令：可预绑定 ⇔ 子表达式可预绑定（其自身归约由 cpu_prepare 负责） template inline constexpr bool cpu_preparable_v> = cpu_preparable_v; // 默认：普通叶子 / 视图 / 索引叶子无需预绑定 template -[[nodiscard]] inline Result cpu_prepare(const T&, ComputeEngine&, - std::size_t, std::size_t) +[[nodiscard]] inline Result cpu_prepare(const T&, ComputeEngine&, std::size_t, std::size_t) { return {}; } template @@ -1319,8 +1361,7 @@ template [[nodiscard]] inline Result cpu_prepare(const Binary& b, ComputeEngine& eng, std::size_t rows, std::size_t cols) { - if (auto r = cpu_prepare(b.l, eng, rows, cols); !r) - return r; + if (auto r = cpu_prepare(b.l, eng, rows, cols); !r) return r; return cpu_prepare(b.r, eng, rows, cols); } @@ -1328,16 +1369,14 @@ template [[nodiscard]] inline Result cpu_prepare(const Select& s, ComputeEngine& eng, std::size_t rows, std::size_t cols) { - if (auto r = cpu_prepare(s.cond, eng, rows, cols); !r) - return r; - if (auto r = cpu_prepare(s.then_e, eng, rows, cols); !r) - return r; + if (auto r = cpu_prepare(s.cond, eng, rows, cols); !r) return r; + if (auto r = cpu_prepare(s.then_e, eng, rows, cols); !r) return r; return cpu_prepare(s.else_e, eng, rows, cols); } // matmul：把 C 物化一次（引擎通用 GEMM 原语）并绑定指针 -[[nodiscard]] inline Result cpu_prepare(const MatmulRef& m, ComputeEngine& eng, - std::size_t, std::size_t) +[[nodiscard]] inline Result cpu_prepare(const MatmulRef& m, ComputeEngine& eng, std::size_t, + std::size_t) { return m.prepare_cpu(eng); } // 归约视图：按输出槽算出行/列归约向量（槽内沿归约轴升序累加） @@ -1345,10 +1384,8 @@ template [[nodiscard]] inline Result cpu_prepare(const ReduceViewRef& r, ComputeEngine&, std::size_t rows, std::size_t cols) { - if (r.v_ != nullptr) - return {}; - if (!r.t.is_cpu()) - NN_FAIL("dsl reduce-view prepare: input not on CPU"); + if (r.v_ != nullptr) return {}; + if (!r.t.is_cpu()) NN_FAIL("dsl reduce-view prepare: input not on CPU"); if (r.t.rows() != rows || r.t.cols() != cols) NN_FAIL("dsl reduce-view prepare: shape mismatch"); using R = ReduceViewRef; @@ -1357,10 +1394,7 @@ template // 精度-存储不匹配约束）。 Matrix mirror; std::span s; - if (r.t.precision() == Precision::F32) - { - s = TensorAccess::cpu_matrix(r.t).span(); - } + if (r.t.precision() == Precision::F32) { s = TensorAccess::cpu_matrix(r.t).span(); } else if (r.t.precision() == Precision::F16) { const auto s16 = TensorAccess::cpu_matrix(r.t).span(); @@ -1383,17 +1417,14 @@ template Scalar a = init; if constexpr (R::reduces_rows) for (std::size_t c = 0; c < cols; ++c) - a = R::reduces_max ? std::max(a, s[slot * cols + c]) - : a + s[slot * cols + c]; + a = R::reduces_max ? std::max(a, s[slot * cols + c]) : a + s[slot * cols + c]; else for (std::size_t rr = 0; rr < rows; ++rr) - a = R::reduces_max ? std::max(a, s[rr * cols + slot]) - : a + s[rr * cols + slot]; + a = R::reduces_max ? std::max(a, s[rr * cols + slot]) : a + s[rr * cols + slot]; r.vec_[slot] = a; }; // 槽间独立（累加方向固定）→ 并行与串行逐字节一致，无需归并 - if (rows * cols >= kDslParallelThreshold && len > 1) - nn::parallel_for_samples(len, slot_kernel); + if (rows * cols >= kDslParallelThreshold && len > 1) nn::parallel_for_samples(len, slot_kernel); else for (std::size_t k = 0; k < len; ++k) slot_kernel(k); @@ -1408,10 +1439,8 @@ template [[nodiscard]] inline Result cpu_prepare(const ReduceRef& r, ComputeEngine& eng, std::size_t rows, std::size_t cols) { - if (r.v_ != nullptr) - return {}; - if (auto pr = cpu_prepare(r.child, eng, rows, cols); !pr) - return pr; + if (r.v_ != nullptr) return {}; + if (auto pr = cpu_prepare(r.child, eng, rows, cols); !pr) return pr; using R = ReduceRef; const Scalar init = R::reduces_max ? std::numeric_limits::lowest() : Scalar{0}; @@ -1435,8 +1464,7 @@ template } r.vec_[slot] = a; }; - if (rows * cols >= kDslParallelThreshold && len > 1) - nn::parallel_for_samples(len, slot_kernel); + if (rows * cols >= kDslParallelThreshold && len > 1) nn::parallel_for_samples(len, slot_kernel); else for (std::size_t k = 0; k < len; ++k) slot_kernel(k); @@ -1452,11 +1480,11 @@ template inline constexpr bool cpu_reduce_root_v = false; template inline constexpr bool cpu_reduce_root_v> = true; template inline constexpr bool cpu_reduce_root_v> = true; -[[nodiscard]] inline Tensor copy_reduce_vector(const Scalar* v, bool reduces_rows, - std::size_t rows, std::size_t cols) +[[nodiscard]] inline Tensor copy_reduce_vector(const Scalar* v, bool reduces_rows, std::size_t rows, + std::size_t cols) { - Matrix out = reduces_rows ? Matrix::make_uninitialized(rows, 1) - : Matrix::make_uninitialized(1, cols); + Matrix out = + reduces_rows ? Matrix::make_uninitialized(rows, 1) : Matrix::make_uninitialized(1, cols); const std::size_t len = reduces_rows ? rows : cols; for (std::size_t k = 0; k < len; ++k) out.span()[k] = v[k]; @@ -1464,13 +1492,13 @@ template inline constexpr bool cpu_reduce_root_v -[[nodiscard]] inline Tensor reduce_vector_tensor(const ReduceViewRef& r, - std::size_t rows, std::size_t cols) +[[nodiscard]] inline Tensor reduce_vector_tensor(const ReduceViewRef& r, std::size_t rows, + std::size_t cols) { return copy_reduce_vector(r.v_, ReduceViewRef::reduces_rows, rows, cols); } template -[[nodiscard]] inline Tensor reduce_vector_tensor(const ReduceRef& r, - std::size_t rows, std::size_t cols) +[[nodiscard]] inline Tensor reduce_vector_tensor(const ReduceRef& r, std::size_t rows, + std::size_t cols) { return copy_reduce_vector(r.v_, ReduceRef::reduces_rows, rows, cols); } // ══════════════════════════════════════════════════════════════════════════ @@ -1551,15 +1579,14 @@ template // 锚点：每个被扫描的表达式**类型**实例化一个 `FusedAnchor`，其静态 // 初始化期把"从类型默认构造的符号实例"折叠出的结构登记进**锚点注册表** // （scan_exprs 结尾与 dry-run 注册表合并，见下）。 -template -struct FusedAnchor +template struct FusedAnchor { - static inline const int reg = [] { + static inline const int reg = [] + { const Expr sym{}; auto [spec, inputs] = to_expr_spec(sym); (void)inputs; // 符号实例：inputs 全是占位张量，注册只用结构 - if (auto v = validate_expr_spec(spec, spec.views.size()); !v) - scan_reject(v.error()); + if (auto v = validate_expr_spec(spec, spec.views.size()); !v) scan_reject(v.error()); anchor_registry().add(spec); return 0; }(); @@ -1574,8 +1601,7 @@ struct FoldAnchor { explicit FoldAnchor(ExprSpec s) { - if (auto v = validate_expr_spec(s, s.views.size()); !v) - scan_reject(v.error()); + if (auto v = validate_expr_spec(s, s.views.size()); !v) scan_reject(v.error()); anchor_registry().add(std::move(s)); } }; @@ -1594,15 +1620,13 @@ struct FoldAnchor // 覆盖，且只在 GPU 闭合世界报错时才暴露）。锚点路径的输入是**符号占位张量** // （默认 f32）→ 其签名恒 0，不会污染特化集合。 template -[[nodiscard]] inline ExprSpec scan_register(const Expr& e, - Precision P = Precision::F32) +[[nodiscard]] inline ExprSpec scan_register(const Expr& e, Precision P = Precision::F32) { #if !defined(NN_SCAN_NO_ANCHOR) (void)FusedAnchor::reg; #endif auto [spec, inputs] = to_expr_spec(e); - if (auto v = validate_expr_spec(spec, inputs.size()); !v) - scan_reject(v.error()); + if (auto v = validate_expr_spec(spec, inputs.size()); !v) scan_reject(v.error()); fused::global_registry().add(spec); // 签名位：bit i = 第 i 个输入是 f16，bit16 = 输出是 f16（与运行时 // expr_prec_sig_of / 生成阶段同源）。这里就地算，避免依赖 compute_engine @@ -1611,10 +1635,8 @@ template std::uint32_t bits = 0; const std::size_t n = inputs.size() < 16u ? inputs.size() : 16u; for (std::size_t i = 0; i < n; ++i) - if (inputs[i].precision() == Precision::F16) - bits |= (1u << i); - fused::global_registry().add_variant( - spec, expr_prec_sig_make(bits, P == Precision::F16)); + if (inputs[i].precision() == Precision::F16) bits |= (1u << i); + fused::global_registry().add_variant(spec, expr_prec_sig_make(bits, P == Precision::F16)); } return spec; } @@ -1634,8 +1656,8 @@ template // · **不做 Auto 推导**：P 是唯一可见实参，来源可追溯（§8.5 G4）。 // ══════════════════════════════════════════════════════════════════════════ template -[[nodiscard]] Result compute(ComputeEngine& eng, const E& e, - std::size_t rows, std::size_t cols, +[[nodiscard]] Result compute(ComputeEngine& eng, const E& e, std::size_t rows, + std::size_t cols, Precision P) // 默认实参在 compute_engine.hpp 前置声明处 { NN_PROF_OP("dsl::compute"); // CPU 模板路径不经引擎入口，须在此计时（剖析器） @@ -1645,8 +1667,7 @@ template // 不拦会静默按 f32 建张量（两端语义不一致）。 if (auto ok = check_precision_supported(P); !ok) { - std::fprintf(stderr, "[dsl] compute 拒绝非存储精度: %s\n", - ok.error().message.c_str()); + std::fprintf(stderr, "[dsl] compute 拒绝非存储精度: %s\n", ok.error().message.c_str()); return Tensor{}; } #ifdef NN_EXPR_SCAN @@ -1663,20 +1684,18 @@ template if (auto* g = fused::recording_graph()) { auto [spec, inputs] = to_expr_spec(e); - if (auto v = validate_expr_spec(spec, inputs.size()); !v) - scan_reject(v.error()); + if (auto v = validate_expr_spec(spec, inputs.size()); !v) scan_reject(v.error()); const int node = g->add_node(spec, inputs, rows, cols, /*vector_out=*/false, P); - Tensor ph = (P == Precision::F16 - ? TensorAccess::cpu(rows, cols) - : TensorAccess::cpu(rows, cols)); + Tensor ph = (P == Precision::F16 ? TensorAccess::cpu(rows, cols) + : TensorAccess::cpu(rows, cols)); ph.set_virtual_tag(g->tag_of_node[static_cast(node)]); g->node_outputs[node] = ph; return ph; } (void)reg_spec; return (P == Precision::F16 ? TensorAccess::cpu(rows, cols) - : TensorAccess::cpu(rows, cols)); + : TensorAccess::cpu(rows, cols)); #else if (eng.device() == Device::CPU) { @@ -1717,16 +1736,14 @@ template return compute(eng, e, rc->first, rc->second, P); } -template -[[nodiscard]] Result compute(ComputeEngine& eng, const E& e) +template [[nodiscard]] Result compute(ComputeEngine& eng, const E& e) { return compute(eng, e, Precision::F32); } // ── CPU 侧"写进目标张量"（f16 目标 → f32 参考求值 + 舍入写回，§7.2）──────── // f32 目标 = 直通 span 写入（零额外拷贝）。 // 写版本：本函数是 DSL 在 CPU 上**绕过 ComputeEngine 写入口**的唯一落点， // 故出口处自己记账（TensorAccess::mark_written）—— 参数派生物缓存依赖它。 -template -inline void eval_into_tensor_cpu(const E& e, Tensor& dst) +template inline void eval_into_tensor_cpu(const E& e, Tensor& dst) { TensorAccess::mark_written(dst); static const auto dbg_max = [](const Matrix& m) @@ -1795,8 +1812,7 @@ template if (auto* g = fused::recording_graph()) { auto [tspec, tinputs] = to_expr_spec(e); - if (auto v = validate_expr_spec(tspec, tinputs.size()); !v) - scan_reject(v.error()); + if (auto v = validate_expr_spec(tspec, tinputs.size()); !v) scan_reject(v.error()); const int node = g->add_node(tspec, tinputs, dst.rows(), dst.cols(), /*vector_out=*/false, dst.precision(), // RV-03：写目标存储标识（同目标写序的显式建边依据） @@ -1809,8 +1825,7 @@ template #else if (eng.device() == Device::CPU) { - if (!dst.is_cpu()) - NN_FAIL("dsl::compute_into: dst not on CPU"); + if (!dst.is_cpu()) NN_FAIL("dsl::compute_into: dst not on CPU"); if constexpr (nn::dsl::has_reduction_v) { // 注意：只在"不含真归约"时才走预绑定路径。含真归约的表达式必须落到 @@ -1861,12 +1876,12 @@ template // P：输出向量精度（§8.1；归约恒 f32 累加 + 输出舍入到 P，§7.3）。 // ══════════════════════════════════════════════════════════════════════════ template -[[nodiscard]] Result compute_reduce(ComputeEngine& eng, const E& e, - std::size_t rows, std::size_t cols, - Precision P) // 默认实参在 compute_engine.hpp 前置声明处 +[[nodiscard]] Result compute_reduce(ComputeEngine& eng, const E& e, std::size_t rows, + std::size_t cols, + Precision P) // 默认实参在 compute_engine.hpp 前置声明处 { - NN_PROF_OP("dsl::compute_reduce"); // CPU 模板路径不经引擎入口，须在此计时（剖析器） - NN_TRY(pc, check_precision_supported(P)); // 同 compute()：拒绝非存储精度 + NN_PROF_OP("dsl::compute_reduce"); // CPU 模板路径不经引擎入口，须在此计时（剖析器） + NN_TRY(pc, check_precision_supported(P)); // 同 compute()：拒绝非存储精度 #ifdef NN_EXPR_SCAN // 构建期扫描：同 compute()，登记结构（归约轴供下面取占位张量形状）。 // 占位张量按归约轴取向量形状 (rows,1)/(1,cols)，使 Layer 后续 @@ -1879,18 +1894,17 @@ template if (auto* g = fused::recording_graph()) { auto [rspec, rinputs] = to_expr_spec(e); - if (auto v = validate_expr_spec(rspec, rinputs.size()); !v) - scan_reject(v.error()); + if (auto v = validate_expr_spec(rspec, rinputs.size()); !v) scan_reject(v.error()); const int raxis0 = expr_spec_reduce_axis(rspec); const int node = g->add_node(rspec, rinputs, rows, cols, /*vector_out=*/true, P); - Tensor ph = (P == Precision::F16 - ? ((raxis0 == 0) ? TensorAccess::cpu(rows, 1) - : (raxis0 == 1) ? TensorAccess::cpu(1, cols) - : TensorAccess::cpu(rows, cols)) - : ((raxis0 == 0) ? TensorAccess::cpu(rows, 1) - : (raxis0 == 1) ? TensorAccess::cpu(1, cols) - : TensorAccess::cpu(rows, cols))); + Tensor ph = + (P == Precision::F16 ? ((raxis0 == 0) ? TensorAccess::cpu(rows, 1) + : (raxis0 == 1) ? TensorAccess::cpu(1, cols) + : TensorAccess::cpu(rows, cols)) + : ((raxis0 == 0) ? TensorAccess::cpu(rows, 1) + : (raxis0 == 1) ? TensorAccess::cpu(1, cols) + : TensorAccess::cpu(rows, cols))); ph.set_virtual_tag(g->tag_of_node[static_cast(node)]); g->node_outputs[node] = ph; return ph; @@ -1898,12 +1912,12 @@ template const ExprSpec spec = reg_spec; const int raxis = expr_spec_reduce_axis(spec); if (P == Precision::F16) - return ((raxis == 0) ? TensorAccess::cpu(rows, 1) - : (raxis == 1) ? TensorAccess::cpu(1, cols) - : TensorAccess::cpu(rows, cols)); - return ((raxis == 0) ? TensorAccess::cpu(rows, 1) - : (raxis == 1) ? TensorAccess::cpu(1, cols) - : TensorAccess::cpu(rows, cols)); + return ((raxis == 0) ? TensorAccess::cpu(rows, 1) + : (raxis == 1) ? TensorAccess::cpu(1, cols) + : TensorAccess::cpu(rows, cols)); + return ((raxis == 0) ? TensorAccess::cpu(rows, 1) + : (raxis == 1) ? TensorAccess::cpu(1, cols) + : TensorAccess::cpu(rows, cols)); #else if (eng.device() == Device::CPU) { @@ -1915,8 +1929,7 @@ template if (auto r = cpu_prepare(e, eng, rows, cols); r) { auto v = reduce_vector_tensor(e, rows, cols); - if (P == Precision::F32) - return (std::move(v)); + if (P == Precision::F32) return (std::move(v)); // f16 目标：归约向量求值在 f32 参考空间，结果舍入到 f16（§7.3） return eng.cast(std::move(v), P); } @@ -1938,8 +1951,7 @@ template return compute_reduce(eng, e, rc->first, rc->second, P); } -template -[[nodiscard]] Result compute_reduce(ComputeEngine& eng, const E& e) +template [[nodiscard]] Result compute_reduce(ComputeEngine& eng, const E& e) { return compute_reduce(eng, e, Precision::F32); } // ══════════════════════════════════════════════════════════════════════════ @@ -1953,28 +1965,25 @@ template // - leaf(d))); // ⚠ 与 ComputeEngine::begin_expr/end_expr（**跨表达式**图录制）不是一回事。 // ══════════════════════════════════════════════════════════════════════════ -template -struct ExprBlock +template struct ExprBlock { // 强引用：ExprBlock 存活期间引擎不会析构。 std::shared_ptr eng; - std::size_t rows, cols; - E expr; + std::size_t rows, cols; + E expr; }; template -[[nodiscard]] ExprBlock start_expr(ComputeEngine& eng, std::size_t rows, - std::size_t cols, const E& e) +[[nodiscard]] ExprBlock start_expr(ComputeEngine& eng, std::size_t rows, std::size_t cols, + const E& e) { auto keepalive = eng.weak_from_this().lock(); return {std::move(keepalive), rows, cols, e}; } -template -[[nodiscard]] Result end_expr(const ExprBlock& b) +template [[nodiscard]] Result end_expr(const ExprBlock& b) { - if (!b.eng) - NN_FAIL("dsl::end_expr: ExprBlock 未绑定引擎"); + if (!b.eng) NN_FAIL("dsl::end_expr: ExprBlock 未绑定引擎"); return compute(*b.eng, b.expr, b.rows, b.cols); } diff --git a/include/neuralnet.cpp/expr_emitter.hpp b/include/neuralnet.cpp/expr_emitter.hpp index e2f1c2db..e8444ace 100644 --- a/include/neuralnet.cpp/expr_emitter.hpp +++ b/include/neuralnet.cpp/expr_emitter.hpp @@ -51,14 +51,12 @@ class ExprEmitter // float16，全链 float16_t 计算）。仅当 expr_prec_sig_native16(spec, sig) // 为真时请求；否则返回空串（上层跳过）。f32 算术变体（native16=false） // 与 native16 变体**并存**——设备按 shaderFloat16 能力选 pipeline。 - [[nodiscard]] virtual std::string generate( - const std::string& name_, const ExprSpec& spec, - ExprPrecSig sig = 0, bool native16 = false) = 0; + [[nodiscard]] virtual std::string generate(const std::string& name_, const ExprSpec& spec, + ExprPrecSig sig = 0, bool native16 = false) = 0; // 归约 kernel 源码；不支持（混合轴/超槽/带类型变体）返回空串 - [[nodiscard]] virtual std::string generate_reduce( - const std::string& name_, const ExprSpec& spec, - ExprPrecSig sig = 0) = 0; + [[nodiscard]] virtual std::string + generate_reduce(const std::string& name_, const ExprSpec& spec, ExprPrecSig sig = 0) = 0; // 批量变体 kernel 源码（键 = key#b，P3 跨链批量派发）：同结构多实例 // 经实例表（uvec2 数组，stride = 输入+输出+尾槽）+ buffer device address @@ -68,8 +66,9 @@ class ExprEmitter // 编码里 → 带 extras 的 f16 结构仍逐 kernel 降级）。运行期分派哨兵不支持。 // 默认返回空串 = 该后端不支持 → 运行期回退逐 kernel 派发（闭合世界降级， // 不是错误）。 - [[nodiscard]] virtual std::string generate_batched( - const std::string&, const ExprSpec&, ExprPrecSig = 0) const { return {}; } + [[nodiscard]] virtual std::string generate_batched(const std::string&, const ExprSpec&, + ExprPrecSig = 0) const + { return {}; } // matmul 细 tile 变体 kernel 源码（键 = key#f，P0 wgrad 网格饥饿）： // 输出块 64×64 → 16×16、16×16 线程 → 8×8、每线程 4×4 → 2×2 输出 @@ -81,8 +80,9 @@ class ExprEmitter // EXPR_PREC_STORAGE_SUFFIX / expr_prec_sig_key）。运行期分派哨兵不支持。 // 默认返回空串 = 该后端不支持 → 构建期不产出条目、运行期留在 64-tile // 路径（降级，不是错误）。 - [[nodiscard]] virtual std::string generate_matmul_fine( - const std::string&, const ExprSpec&, ExprPrecSig = 0) const { return {}; } + [[nodiscard]] virtual std::string generate_matmul_fine(const std::string&, const ExprSpec&, + ExprPrecSig = 0) const + { return {}; } // matmul K 切分变体 kernel 源码（键 = key#ks，wgrad 二次归约）： // 细 tile 网格仍低于 EXPR_MATMUL_KSPLIT_MIN_WGS 时，把求和维 K 切成 @@ -94,16 +94,18 @@ class ExprEmitter // 由伴生 #kr（generate_matmul_ksplit_reduce）完成二次归约 + 尾链。 // 默认返回空串 = 该后端不支持 → 构建期不产出条目、运行期留细 tile // 路径（降级，不是错误）。 - [[nodiscard]] virtual std::string generate_matmul_ksplit( - const std::string&, const ExprSpec&, ExprPrecSig = 0) const { return {}; } + [[nodiscard]] virtual std::string generate_matmul_ksplit(const std::string&, const ExprSpec&, + ExprPrecSig = 0) const + { return {}; } // K 切分伴生归约 kernel 源码（键 = key#kr）：绑定 = 原输入 0..N-1 + // 部分和（binding N）+ 输出（binding N+1）；读段数经 PC（mm_k 槽位）。 // 求和按段号升序左折（确定性），随后跑与 V0 同源的尾链。 // 与 generate_matmul_ksplit 配对产出（有 #ks 才有 #kr）。 // 默认返回空串 = 该后端不支持 → 构建期不产出条目（#ks 同步不产出）。 - [[nodiscard]] virtual std::string generate_matmul_ksplit_reduce( - const std::string&, const ExprSpec&, ExprPrecSig = 0) const { return {}; } + [[nodiscard]] virtual std::string + generate_matmul_ksplit_reduce(const std::string&, const ExprSpec&, ExprPrecSig = 0) const + { return {}; } }; // ── 简单注册表：按后端名选择 emitter 工厂 ──────────────────────────────── @@ -114,7 +116,7 @@ namespace emitter_registry { // 工厂返回 unique_ptr（铁律 2：无裸指针所有权）；make() 直接转发。 -using Factory = std::unique_ptr(*)(); +using Factory = std::unique_ptr (*)(); // 后端表（静态初始化，见各 emitter 头文件的登记语句） inline std::unordered_map& backends() @@ -135,8 +137,7 @@ inline bool register_backend(std::string_view name, Factory factory) { const auto& m = backends(); const auto it = m.find(std::string(name)); - if (it == m.end()) - return nullptr; + if (it == m.end()) return nullptr; return it->second(); } @@ -153,4 +154,3 @@ inline bool register_backend(std::string_view name, Factory factory) } // namespace emitter_registry } // namespace nn - diff --git a/include/neuralnet.cpp/expr_fold.hpp b/include/neuralnet.cpp/expr_fold.hpp index 1a494cf8..5d2dbe0f 100644 --- a/include/neuralnet.cpp/expr_fold.hpp +++ b/include/neuralnet.cpp/expr_fold.hpp @@ -31,14 +31,13 @@ namespace nn::expr { // 末指令 = 状态自复制（dst 自身 Add 常量 0），把"输出 = 末指令 dst"落到状态上 -[[nodiscard]] inline ExprInstr fold_emit_state(std::uint8_t state_reg, - std::uint8_t const_zero_slot) +[[nodiscard]] inline ExprInstr fold_emit_state(std::uint8_t state_reg, std::uint8_t const_zero_slot) { ExprInstr ins; - ins.op = static_cast(ExprOp::Add); + ins.op = static_cast(ExprOp::Add); ins.dst = state_reg; - ins.a = reg(state_reg); - ins.b = cst(const_zero_slot); + ins.a = reg(state_reg); + ins.b = cst(const_zero_slot); return ins; } @@ -50,25 +49,25 @@ namespace nn::expr { ExprSpec s; FoldSpec f; - f.k = k; - f.num_state = 1; - f.inits = { -std::numeric_limits::infinity() }; + f.k = k; + f.num_state = 1; + f.inits = {-std::numeric_limits::infinity()}; ExprInstr b0; // blk = RowMax(input 0) - b0.op = static_cast(ExprOp::RowMax); + b0.op = static_cast(ExprOp::RowMax); b0.dst = 1; - b0.a = input(0); - b0.b = {}; + b0.a = input(0); + b0.b = {}; ExprInstr b1; // m = max(m, Reduce(blk)) - b1.op = static_cast(ExprOp::Max); + b1.op = static_cast(ExprOp::Max); b1.dst = 0; - b1.a = reg(0); - b1.b = reduce(1); - f.body = { b0, b1 }; - f.finalize = { fold_emit_state(0, 0) }; - s.fold = std::move(f); - s.consts = { Scalar{0} }; + b1.a = reg(0); + b1.b = reduce(1); + f.body = {b0, b1}; + f.finalize = {fold_emit_state(0, 0)}; + s.fold = std::move(f); + s.consts = {Scalar{0}}; s.num_regs = 2; // 状态 0 + 临时 1 - s.views = { linear() }; + s.views = {linear()}; return s; } @@ -77,24 +76,24 @@ namespace nn::expr { ExprSpec s; FoldSpec f; - f.k = k; - f.num_state = 1; - f.inits = { Scalar{0} }; + f.k = k; + f.num_state = 1; + f.inits = {Scalar{0}}; ExprInstr b0; - b0.op = static_cast(ExprOp::RowSum); + b0.op = static_cast(ExprOp::RowSum); b0.dst = 1; - b0.a = input(0); + b0.a = input(0); ExprInstr b1; - b1.op = static_cast(ExprOp::Add); + b1.op = static_cast(ExprOp::Add); b1.dst = 0; - b1.a = reg(0); - b1.b = reduce(1); - f.body = { b0, b1 }; - f.finalize = { fold_emit_state(0, 0) }; - s.fold = std::move(f); - s.consts = { Scalar{0} }; + b1.a = reg(0); + b1.b = reduce(1); + f.body = {b0, b1}; + f.finalize = {fold_emit_state(0, 0)}; + s.fold = std::move(f); + s.consts = {Scalar{0}}; s.num_regs = 2; - s.views = { linear() }; + s.views = {linear()}; return s; } @@ -115,37 +114,37 @@ namespace nn::expr // 参考语义：l_final = Σ_k exp(x(row,k) − max_k x)（浮点结合序容差内等价）。 [[nodiscard]] inline ExprSpec make_fold_softmax_denom(std::uint32_t k) { - const auto ins = [](ExprOp op, std::uint8_t dst, ExprOperand a, - ExprOperand b = {}) { + const auto ins = [](ExprOp op, std::uint8_t dst, ExprOperand a, ExprOperand b = {}) + { ExprInstr i; - i.op = static_cast(op); + i.op = static_cast(op); i.dst = dst; - i.a = a; - i.b = b; + i.a = a; + i.b = b; return i; }; ExprSpec s; FoldSpec f; - f.k = k; + f.k = k; f.num_state = 2; - f.inits = { -std::numeric_limits::infinity(), Scalar{0} }; + f.inits = {-std::numeric_limits::infinity(), Scalar{0}}; f.body = { - ins(ExprOp::Add, 2, reg(0), cst(0)), // m_old 拷贝 + ins(ExprOp::Add, 2, reg(0), cst(0)), // m_old 拷贝 ins(ExprOp::RowMax, 3, input(0)), // 块内 blk_m - ins(ExprOp::Max, 0, reg(2), reduce(3)),// m' ← 状态更新 - ins(ExprOp::Sub, 4, reg(2), reg(0)), // m_old − m' - ins(ExprOp::Exp, 5, reg(4)), // α - ins(ExprOp::Sub, 7, input(0), reg(0)), // x − m' - ins(ExprOp::Exp, 8, reg(7)), // exp(x − m') + ins(ExprOp::Max, 0, reg(2), reduce(3)),// m' ← 状态更新 + ins(ExprOp::Sub, 4, reg(2), reg(0)), // m_old − m' + ins(ExprOp::Exp, 5, reg(4)), // α + ins(ExprOp::Sub, 7, input(0), reg(0)), // x − m' + ins(ExprOp::Exp, 8, reg(7)), // exp(x − m') ins(ExprOp::RowSum, 6, reg(8)), // 块内 blk_l - ins(ExprOp::Mul, 9, reg(1), reg(5)), // l · α - ins(ExprOp::Add, 1, reg(9), reduce(6)),// l' ← 状态更新 + ins(ExprOp::Mul, 9, reg(1), reg(5)), // l · α + ins(ExprOp::Add, 1, reg(9), reduce(6)),// l' ← 状态更新 }; - f.finalize = { fold_emit_state(1, 0) }; - s.fold = std::move(f); - s.consts = { Scalar{0} }; + f.finalize = {fold_emit_state(1, 0)}; + s.fold = std::move(f); + s.consts = {Scalar{0}}; s.num_regs = 10; // 状态 0..1 + 临时 2..9 - s.views = { linear() }; + s.views = {linear()}; return s; } @@ -159,17 +158,16 @@ namespace nn::expr // l ≥ 1 ⇒ log l 有限；dst 用临时寄存器，不覆写被读状态） // 参考语义：lse_final = log Σ_k exp(x(row,k) − max_k x) + max_k x // （浮点结合序容差内等价于直接 logsumexp）。 -[[nodiscard]] inline ExprSpec make_fold_logsumexp(std::uint32_t k, - ExprView input_view = linear()) +[[nodiscard]] inline ExprSpec make_fold_logsumexp(std::uint32_t k, ExprView input_view = linear()) { ExprSpec s = make_fold_softmax_denom(k); // body/states/inits/consts 同源 - const auto ins = [](ExprOp op, std::uint8_t dst, ExprOperand a, - ExprOperand b = {}) { + const auto ins = [](ExprOp op, std::uint8_t dst, ExprOperand a, ExprOperand b = {}) + { ExprInstr i; - i.op = static_cast(op); + i.op = static_cast(op); i.dst = dst; - i.a = a; - i.b = b; + i.a = a; + i.b = b; return i; }; // finalize：lse = m + log(l)——寄存器 10/11 为 finalize 专用临时 @@ -178,7 +176,7 @@ namespace nn::expr ins(ExprOp::Add, 11, reg(0), reg(10)), }; s.num_regs = 12; // 状态 0..1 + body 临时 2..9 + finalize 10..11 - s.views = { input_view }; + s.views = {input_view}; return s; } @@ -200,9 +198,9 @@ namespace nn::expr // k（= 本头的唯一形状参数）不进 key → 一个代表值覆盖所有收缩长度。 inline const bool g_fold_sample_anchors = [] { - (void)dsl::FoldAnchor{ make_fold_rowmax(64) }; - (void)dsl::FoldAnchor{ make_fold_rowsum(64) }; - (void)dsl::FoldAnchor{ make_fold_softmax_denom(64) }; + (void)dsl::FoldAnchor{make_fold_rowmax(64)}; + (void)dsl::FoldAnchor{make_fold_rowsum(64)}; + (void)dsl::FoldAnchor{make_fold_softmax_denom(64)}; return true; }(); } // namespace nn::expr diff --git a/include/neuralnet.cpp/expr_glsl_gen.hpp b/include/neuralnet.cpp/expr_glsl_gen.hpp index e095f985..896b668a 100644 --- a/include/neuralnet.cpp/expr_glsl_gen.hpp +++ b/include/neuralnet.cpp/expr_glsl_gen.hpp @@ -38,28 +38,25 @@ namespace nn // 因此改用整数位运算在 f32 域做精确 RNE：返回值本身落在 f16 网格上（f16 ⊂ f32 // 可精确表示），随后的 float16_t() 收窄无论取哪种舍入模式都不再产生偏差。 [[nodiscard]] inline std::string glsl_f16_store(const std::string& v) -{ - return "float16_t(nn_f16_rne(" + v + "))"; -} +{ return "float16_t(nn_f16_rne(" + v + "))"; } // 上面 glsl_f16_store 依赖的 GLSL 辅助函数源码；凡会写 f16 的生成 shader // 前导必须发射一次（声明必须先于使用）。用整数位运算实现精确 RNE，不依赖 // 驱动 OpFConvert / packHalf2x16 的舍入模式（两者实测都向零截断）。 [[nodiscard]] inline const char* glsl_f16_rne_source() { - return - "// f32 → f16 网格的 RNE 预舍入：返回值本身落在 f16 网格上（f16 ⊂ f32\n" - "// 可精确表示），随后的 float16_t() 收窄不再产生任何舍入偏差。\n" - "float nn_f16_rne(float x)\n" - "{\n" - " uint u = floatBitsToUint(x);\n" - " const uint a = u & 0x7fffffffu;\n" - " if (a >= 0x7f800000u || a < 0x38800000u) // Inf/NaN，或 |x| < 2^-14\n" - " return roundEven(x * 16777216.0) / 16777216.0;\n" - " u += 0x0fffu + ((u >> 13u) & 1u); // 半个 ULP + tie-to-even\n" - " u &= 0xffffe000u; // 截到 f16 的 10 位尾数\n" - " return uintBitsToFloat(u);\n" - "}\n\n"; + return "// f32 → f16 网格的 RNE 预舍入：返回值本身落在 f16 网格上（f16 ⊂ f32\n" + "// 可精确表示），随后的 float16_t() 收窄不再产生任何舍入偏差。\n" + "float nn_f16_rne(float x)\n" + "{\n" + " uint u = floatBitsToUint(x);\n" + " const uint a = u & 0x7fffffffu;\n" + " if (a >= 0x7f800000u || a < 0x38800000u) // Inf/NaN，或 |x| < 2^-14\n" + " return roundEven(x * 16777216.0) / 16777216.0;\n" + " u += 0x0fffu + ((u >> 13u) & 1u); // 半个 ULP + tie-to-even\n" + " u &= 0xffffe000u; // 截到 f16 的 10 位尾数\n" + " return uintBitsToFloat(u);\n" + "}\n\n"; } // ── 算子 → GLSL 右值片段（形态/符号来自 expr_op_table.hpp，RV-07）────────── @@ -71,10 +68,9 @@ namespace nn // wrap ：内联上下文需额外一层括号（与语句形态逐字一致） // b/c 以惰性回调传入：一元 / 函数算子不触碰不使用的操作数（保持原求值副作用序）。 template -[[nodiscard]] inline std::string glsl_alu_expr(const detail::ExprOpInfo& info, - const std::string& a, - GetB&& b, GetC&& c, - bool vec4, bool native16, bool wrap) +[[nodiscard]] inline std::string glsl_alu_expr(const detail::ExprOpInfo& info, const std::string& a, + GetB&& b, GetC&& c, bool vec4, bool native16, + bool wrap) { switch (info.emit) { @@ -86,12 +82,11 @@ template case detail::ExprEmitKind::Func2: return std::string(info.sym) + "(" + a + ", " + b() + ")"; case detail::ExprEmitKind::Compare: - if (vec4) - return std::string("vec4(") + info.sym_vec + "(" + a + ", " + b() + "))"; + if (vec4) return std::string("vec4(") + info.sym_vec + "(" + a + ", " + b() + "))"; { - const std::string e = "(" + a + " " + info.sym + " " + b() + ") ? " - + (native16 ? "float16_t(1.0)" : "1.0") + " : " - + (native16 ? "float16_t(0.0)" : "0.0"); + const std::string e = "(" + a + " " + info.sym + " " + b() + ") ? " + + (native16 ? "float16_t(1.0)" : "1.0") + " : " + + (native16 ? "float16_t(0.0)" : "0.0"); return wrap ? ("(" + e + ")") : e; } case detail::ExprEmitKind::Neg: @@ -106,8 +101,8 @@ template // mix(x,y,t) = x(1-t) + y·t：t=(a≠0) 时须取 b → b 在 y 位 return "mix(" + c() + ", " + b() + ", notEqual(" + a + ", vec4(0.0)))"; { - const std::string e = "(" + a + " != " - + (native16 ? "float16_t(0.0)" : "0.0") + ") ? " + b() + " : " + c(); + const std::string e = "(" + a + " != " + (native16 ? "float16_t(0.0)" : "0.0") + + ") ? " + b() + " : " + c(); return wrap ? ("(" + e + ")") : e; } case detail::ExprEmitKind::None: @@ -129,17 +124,13 @@ template // 单次构建：给定缓冲区名与该视图是否 f16 视图，产出读取表达式字符串。 // **非分派路径的输出与历史逐字一致**（RotateHalf 内部转 f32、其余分支裸读， // 外层转换仍由调用方 float(...) 负责）——这是 V0/V2 文本 byte-identical 的前提。 -inline std::string glsl_view_read_(const ExprView& v, std::uint32_t buf_id, - const std::string& buf, bool f16_buf, - const std::string& idx_var, +inline std::string glsl_view_read_(const ExprView& v, std::uint32_t buf_id, const std::string& buf, + bool f16_buf, const std::string& idx_var, const std::string& row_var, const std::string& col_var, - std::uint32_t vp_slot, - ExprPrecSig sig) + std::uint32_t vp_slot, ExprPrecSig sig) { - const auto rw = [&](const std::string& idx) -> std::string { - return f16_buf ? ("float(" + buf + "[" + idx + "])") - : (buf + "[" + idx + "]"); - }; + const auto rw = [&](const std::string& idx) -> std::string + { return f16_buf ? ("float(" + buf + "[" + idx + "])") : (buf + "[" + idx + "]"); }; const bool disp = expr_prec_sig_is_dispatch(sig); switch (v.kind) { @@ -150,31 +141,29 @@ inline std::string glsl_view_read_(const ExprView& v, std::uint32_t buf_id, { // block 为运行时视图参数（push constant 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 rl = "(" + row_var + " % vp" + std::to_string(vp_slot) + ")"; const std::string half = "(vp" + std::to_string(vp_slot) + " / 2u)"; // 列主序索引 = 源行 * cols + 列号（cols 为 push constant 列数） - 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 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 ? "-" : ""; return "((" + rl + " < " + half + ") ? " + neg + rw(hi) + " : " + rw(lo) + ")"; } case static_cast(ExprViewKind::RowMod): // 周期为运行时视图参数（push constant vpN），形状无关融合 - return buf + "[(" + row_var + " % vp" + std::to_string(vp_slot) + ") * cols + " - + col_var + "]"; + return buf + "[(" + row_var + " % vp" + std::to_string(vp_slot) + ") * cols + " + col_var + + "]"; case static_cast(ExprViewKind::RowAccess): // 行偏移+取模：data[(offset + r % mod)*cols + c]。 // vp_slot=mod（op），vp_slot+1=offset（param2），双 vp 槽（形状无关）。 - return buf + "[(" + "vp" + std::to_string(vp_slot + 1) + " + (" - + row_var + " % vp" + std::to_string(vp_slot) + ")) * cols + " - + col_var + "]"; + return buf + "[(" + "vp" + std::to_string(vp_slot + 1) + " + (" + row_var + " % vp" + + std::to_string(vp_slot) + ")) * cols + " + col_var + "]"; case static_cast(ExprViewKind::Transpose): // 转置读：data[col*stride + row]（stride = 存储行距，运行期 vp 槽）。 // 与 Linear 同为裸读（f16 外层转换由调用方负责）。 - return buf + "[" + col_var + " * vp" + std::to_string(vp_slot) + " + " - + row_var + "]"; + return buf + "[" + col_var + " * vp" + std::to_string(vp_slot) + " + " + row_var + "]"; case static_cast(ExprViewKind::GroupedReduceSum): case static_cast(ExprViewKind::GroupedReduceMax): // 分组归约：输出行 r → 输入组行 [r*R, (r+1)*R)，列不变。 @@ -197,12 +186,11 @@ inline std::string glsl_view_read_(const ExprView& v, std::uint32_t buf_id, const std::string lb = "b" + std::to_string(v.param); std::string lab; if (disp) - lab = "((prec >> " + std::to_string(v.param) + "u & 1u) != 0u ? uint(float(" - + lb + "h[" + col_var + "])) : uint(" + lb + "[" + col_var + "]))"; + lab = "((prec >> " + std::to_string(v.param) + "u & 1u) != 0u ? uint(float(" + lb + + "h[" + col_var + "])) : uint(" + lb + "[" + col_var + "]))"; else if (expr_prec_sig_in_f16(sig, v.param)) lab = "uint(float(" + lb + "[" + col_var + "]))"; - else - lab = "uint(" + lb + "[" + col_var + "])"; + else lab = "uint(" + lb + "[" + col_var + "])"; return buf + "[" + lab + " * cols + " + col_var + "]"; } case static_cast(ExprViewKind::BatchMod): @@ -216,13 +204,10 @@ inline std::string glsl_view_read_(const ExprView& v, std::uint32_t buf_id, } } -inline void glsl_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, - ExprPrecSig sig = 0, - bool f16_native = false) +inline void glsl_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, + ExprPrecSig sig = 0, bool f16_native = false) { const std::string buf = "b" + std::to_string(buf_id); // 该输入的存储类型：f16 时**视图内部**的算术/索引读取必须显式转 f32 @@ -234,8 +219,7 @@ inline void glsl_view_read(std::ostringstream& os, const ExprViewKind vkind = static_cast(v.kind); if (!expr_prec_sig_is_dispatch(sig)) { - os << glsl_view_read_(v, buf_id, buf, f16_buf, idx_var, row_var, col_var, - vp_slot, sig); + os << glsl_view_read_(v, buf_id, buf, f16_buf, idx_var, row_var, col_var, vp_slot, sig); return; } if (vkind == ExprViewKind::GroupedReduceSum || vkind == ExprViewKind::GroupedReduceMax) @@ -245,34 +229,30 @@ inline void glsl_view_read(std::ostringstream& os, // 产生 `prec ? … : …` 里嵌同样 `prec ? … : …` 的重复表达式（语义相同、 // 文本翻倍）。RowGather **不能**走这条快路：它自带分派的只有**标签槽**， // 取值仍需外层三元在 `b` / `bh` 之间选择。 - os << glsl_view_read_(v, buf_id, buf, false, idx_var, row_var, col_var, - vp_slot, sig); + os << glsl_view_read_(v, buf_id, buf, false, idx_var, row_var, col_var, vp_slot, sig); return; } // ── 运行期精度分派（V1）：同一索引同时作两种视图解引用，uniform 选择 ── // f16 分支取 float(...)（行内 rw 已转的表达式再套一层 float 是无害恒等）。 - const std::string f32e = glsl_view_read_(v, buf_id, buf, false, idx_var, - row_var, col_var, vp_slot, sig); + const std::string f32e = + glsl_view_read_(v, buf_id, buf, false, idx_var, row_var, col_var, vp_slot, sig); const std::string hbuf = buf + "h"; - const std::string f16e = glsl_view_read_(v, buf_id, hbuf, true, idx_var, - row_var, col_var, vp_slot, sig); - os << "((prec >> " << buf_id << "u & 1u) != 0u ? float(" << f16e - << ") : " << f32e << ")"; + const std::string f16e = + glsl_view_read_(v, buf_id, hbuf, true, idx_var, row_var, col_var, vp_slot, sig); + os << "((prec >> " << buf_id << "u & 1u) != 0u ? float(" << f16e << ") : " << f32e << ")"; } // ── 视图的运行时参数（vp）槽位 ────────────────────────────────────────── // 槽位 = 排在该视图之前的 runtime-param 视图槽数之和（RowAccess 占 2 槽）。 // **唯一实现**：六个生成路径此前各写一遍同一段循环，任一处漂移就会读写错位的 // push constant 字（后端 `pc_base` 的注释记录过这类"静默错值"事故）。 -[[nodiscard]] inline std::uint32_t glsl_vp_slot(const ExprSpec& spec, - std::size_t view_idx) +[[nodiscard]] inline std::uint32_t glsl_vp_slot(const ExprSpec& spec, std::size_t view_idx) { std::uint32_t vp = 0; for (std::size_t j = 0; j < view_idx; ++j) { const auto k = static_cast(spec.views[j].kind); - if (expr_view_has_runtime_param(k)) - vp += expr_view_runtime_param_slots(k); + if (expr_view_has_runtime_param(k)) vp += expr_view_runtime_param_slots(k); } return vp; } @@ -284,8 +264,7 @@ inline void glsl_view_read(std::ostringstream& os, [[nodiscard]] inline std::string gqa_blk_expr(const ExprSpec& spec, std::size_t view_idx, const std::string& base = "batch") { - if (view_idx >= spec.views.size()) - return base; + if (view_idx >= spec.views.size()) return base; if (static_cast(spec.views[view_idx].kind) != ExprViewKind::HeadGroup) return base; return "(" + base + " / vp" + std::to_string(glsl_vp_slot(spec, view_idx)) + ")"; @@ -300,40 +279,34 @@ inline void glsl_view_read(std::ostringstream& os, // **每个可能发射 `gr_r` 的生成器都必须调用本函数**：glsl_view_read 对 // GroupedReduce 视图无条件发 `gr_r(...)`，目前 elementwise / matmul / // 归约三处调用它。无该视图时本函数不产出任何文本（调用零成本）。 -inline void glsl_emit_grouped_reduce_helpers(std::ostringstream& os, - const ExprSpec& spec, - ExprPrecSig sig = 0, - bool f16_native = false) +inline void glsl_emit_grouped_reduce_helpers(std::ostringstream& os, const ExprSpec& spec, + ExprPrecSig sig = 0, bool f16_native = false) { const char* ty = f16_native ? "float16_t" : "float"; for (std::size_t i = 0; i < spec.views.size(); ++i) { const auto k = static_cast(spec.views[i].kind); - if (k != ExprViewKind::GroupedReduceSum && - k != ExprViewKind::GroupedReduceMax) - continue; + if (k != ExprViewKind::GroupedReduceSum && k != ExprViewKind::GroupedReduceMax) continue; // vp 槽 = 按视图序累计（与 glsl_view_read 调用点同源，见 glsl_vp_slot） const std::uint32_t vp = glsl_vp_slot(spec, i); const bool is_max = (k == ExprViewKind::GroupedReduceMax); const bool disp = expr_prec_sig_is_dispatch(sig); const bool f16_buf = expr_prec_sig_in_f16(sig, i) && !f16_native; const std::string buf = "b" + std::to_string(i); - const auto elem = [&](const std::string& tt) -> std::string { + const auto elem = [&](const std::string& tt) -> std::string + { const std::string idx = "((row * R + " + tt + ") * cols + col)"; if (disp) - return "((prec >> " + std::to_string(i) + "u & 1u) != 0u ? float(" - + buf + "h[" + idx + "]) : " + buf + "[" + idx + "])"; - return f16_buf ? ("float(" + buf + "[" + idx + "])") - : (buf + "[" + idx + "]"); + return "((prec >> " + std::to_string(i) + "u & 1u) != 0u ? float(" + buf + "h[" + + idx + "]) : " + buf + "[" + idx + "])"; + return f16_buf ? ("float(" + buf + "[" + idx + "])") : (buf + "[" + idx + "]"); }; os << "\n" << ty << " gr_r" << i << "(uint row, uint col)\n{\n"; os << " const uint R = vp" << vp << ";\n"; os << " " << ty << " acc = " << elem("0u") << ";\n"; os << " for (uint tt = 1u; tt < R; ++tt)\n"; - if (is_max) - os << " acc = max(acc, " << elem("tt") << ");\n"; - else - os << " acc += " << elem("tt") << ";\n"; + if (is_max) os << " acc = max(acc, " << elem("tt") << ");\n"; + else os << " acc += " << elem("tt") << ";\n"; os << " return acc;\n}\n"; } } @@ -365,8 +338,7 @@ inline bool glsl_spec_uses_batch(const ExprSpec& spec) const ExprOperand* ops[3] = {&ins.a, &ins.b, &ins.c}; const std::size_t nops = expr_instr_num_operands(static_cast(ins.op)); for (std::size_t oi = 0; oi < nops; ++oi) - if (ops[oi]->kind == static_cast(ExprOperandKind::Batch)) - return true; + if (ops[oi]->kind == static_cast(ExprOperandKind::Batch)) return true; } return false; } @@ -379,8 +351,7 @@ inline bool glsl_spec_uses_batch(const ExprSpec& spec) // 任意形状仍走同 kernel 内的标量回退——正确性不依赖编译期形状。 inline bool glsl_vec4_eligible(const ExprSpec& spec) { - if (spec.matmul) - return false; + if (spec.matmul) return false; for (const auto& v : spec.views) { const auto k = static_cast(v.kind); @@ -409,11 +380,10 @@ inline bool glsl_vec4_eligible(const ExprSpec& spec) // 输出调用）共用同一条尾链——两者的输入绑定/PC 布局同源（rows/mm_batch/ // cols/vp/consts/rparams），故 #ks 只写部分和、#kr 才跑尾链（尾链非线性， // 不能逐段先跑再加）。n_inputs = 原结构输入数（跳过 A/B 槽）。 -inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec, - ExprPrecSig sig, std::size_t n_inputs) +inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec, ExprPrecSig sig, + std::size_t n_inputs) { - const std::size_t last_dst = spec.instrs.empty() - ? 0u : spec.instrs.back().dst; + const std::size_t last_dst = spec.instrs.empty() ? 0u : spec.instrs.back().dst; L << "float eval_tail(float mm, uint row, uint col, uint batch)\n{\n"; if (!spec.instrs.empty()) { @@ -451,21 +421,21 @@ inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec L << " const uint grow = batch * m_per + row;\n"; for (std::size_t i = 0; i < n_inputs; ++i) { - if (i == spec.matmul->a_input || i == spec.matmul->b_input) - continue; + if (i == spec.matmul->a_input || i == spec.matmul->b_input) continue; const std::uint32_t vp = glsl_vp_slot(spec, i); L << " const float v" << i << " = "; const bool tail_f16 = sig != 0 && expr_prec_sig_in_f16(sig, i); if (tail_f16) L << "float("; // f16 读 → 显式转 f32（算术仍在 f32） - glsl_view_read(L, spec.views[i], static_cast(i), - "grow*cols + col", "grow", "col", vp, sig); + glsl_view_read(L, spec.views[i], static_cast(i), "grow*cols + col", + "grow", "col", vp, sig); if (tail_f16) L << ")"; L << ";\n"; } if (spec.num_regs > 0) { L << " float r0"; - for (std::uint32_t r = 1; r < spec.num_regs; ++r) L << ", r" << r; + for (std::uint32_t r = 1; r < spec.num_regs; ++r) + L << ", r" << r; L << ";\n"; } for (const auto& ins : spec.instrs) @@ -478,9 +448,9 @@ inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec return; // 归约指令：由 generate_glsl_reduce 处理 → 上层报错 // 发射形态/符号查表（expr_op_table）→ 单处生成，与标量链逐字一致 L << " " << dst << " = " - << glsl_alu_expr(info, a, [&] { return operand(ins.b); }, - [&] { return operand(ins.c); }, - /*vec4=*/false, /*native16=*/false, /*wrap=*/false) + << glsl_alu_expr( + info, a, [&] { return operand(ins.b); }, [&] { return operand(ins.c); }, + /*vec4=*/false, /*native16=*/false, /*wrap=*/false) << ";\n"; } L << " return r" << last_dst << ";\n"; @@ -537,9 +507,10 @@ inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec // 尾链。段界替换了全部 k 上界守卫（kend 取代 mm_k）：非末段的上界是 // 段界而非 mm_k，否则会把下一段的 k 拉进本段的部分和（静默错值）。 // ═══════════════════════════════════════════════════════════════════════════ -[[nodiscard]] inline std::string generate_glsl_matmul( - const std::string& name, const ExprSpec& spec, ExprPrecSig sig = 0, - std::uint32_t block = EXPR_MATMUL_BLOCK, bool k_split = false) +[[nodiscard]] inline std::string generate_glsl_matmul(const std::string& name, const ExprSpec& spec, + ExprPrecSig sig = 0, + std::uint32_t block = EXPR_MATMUL_BLOCK, + bool k_split = false) { // block = 输出块每维元素数（64 = 默认 / 16 = 细 tile 变体 "#f"）。 // T 按 block 选：64 → 16 线程/维（R=4 = 4×4 vec4 寄存器分块）； @@ -547,8 +518,7 @@ inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec // R = B/T = 每线程输出边长：R=2 在"网格摊满 SM（48 WG）"与"线程内 // ILP（4 条独立 K 链）"之间取甜点（实测见 EXPR_MATMUL_BLOCK_FINE 注释）。 // 其余 block 值无对应计算路径 → 拒绝。 - if (block != EXPR_MATMUL_BLOCK && block != EXPR_MATMUL_BLOCK_FINE) - return {}; + if (block != EXPR_MATMUL_BLOCK && block != EXPR_MATMUL_BLOCK_FINE) return {}; // K 切分变体（键 #ks）：仅细 tile 网格 + 仅**静态签名**（sig==0 或具体 // 签名；运行期分派哨兵 0x20000 需要 `prec` PC 槽 + 输入/输出双视图，与 // "免尾链 + 写部分和"的绑定/PC 形态不相容 → 防御性拒绝）。 @@ -556,9 +526,7 @@ inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec // 后端按 f32 分配该缓冲（run_fused_gpu 的 use_ks 分支），故此处必须把 // 输出声明钉成 float、写原始 acc（不做 f16 舍入）——否则 f16 签名会按 // 2B/元素写进 4B/元素的缓冲（静默错值）。最终精度由伴生 #kr 落盘。 - if (k_split && (block != EXPR_MATMUL_BLOCK_FINE || - expr_prec_sig_is_dispatch(sig))) - return {}; + if (k_split && (block != EXPR_MATMUL_BLOCK_FINE || expr_prec_sig_is_dispatch(sig))) return {}; const MatmulSpec& mm = *spec.matmul; const std::size_t n_inputs = spec.views.size(); const std::uint32_t a_slot = mm.a_input; @@ -585,8 +553,8 @@ inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec // 全部下游公式（local_size、VPT、共享维度 B/4、块锚 wg×B、写回）均以 // T/B 为参 → 两个变体共用同一份生成代码，只有网格粒度与寄存器分块不同。 const std::uint32_t B = block; // 输出块每维（64 或 16） - const std::uint32_t T = (B == EXPR_MATMUL_BLOCK) - ? EXPR_MATMUL_TILE : EXPR_MATMUL_TILE_FINE; // 16 或 8 + const std::uint32_t T = + (B == EXPR_MATMUL_BLOCK) ? EXPR_MATMUL_TILE : EXPR_MATMUL_TILE_FINE; // 16 或 8 const std::uint32_t R = B / T; // 每线程输出边长（4 或 2） // BK 取值：16 + 双缓冲（16KB 共享）——同参数四点交错 A/B 全点最优 // （结论见头注释 trade-off 段）。 @@ -601,26 +569,22 @@ inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec std::ostringstream L; L << "// ── 自动生成（AOT 算子融合 · matmul 分块 4×4 双缓冲），请勿手动编辑 ──\n"; L << "// 表达式: " << name; - if (sign_f16 && !disp) - L << " [精度 " << expr_prec_sig_str(sig, n_inputs) << "]"; - else if (disp) - L << " [运行期精度分派]"; + if (sign_f16 && !disp) L << " [精度 " << expr_prec_sig_str(sig, n_inputs) << "]"; + else if (disp) L << " [运行期精度分派]"; L << "\n"; L << "#version 450\n\n"; - if (sign_f16) - L << "#extension GL_EXT_shader_16bit_storage : require\n\n"; - if (sign_f16) - L << glsl_f16_rne_source(); + if (sign_f16) L << "#extension GL_EXT_shader_16bit_storage : require\n\n"; + if (sign_f16) L << glsl_f16_rne_source(); L << "layout(local_size_x = " << T << ", local_size_y = " << T << ") in;\n\n"; for (std::size_t i = 0; i < n_inputs; ++i) { - L << "layout(std430, binding = " << i << ") readonly buffer Buf" << i - << " { " << ((disp || (sign_f16 && expr_prec_sig_in_f16(sig, i))) - ? (disp ? "float" : "float16_t") : "float") + L << "layout(std430, binding = " << i << ") readonly buffer Buf" << i << " { " + << ((disp || (sign_f16 && expr_prec_sig_in_f16(sig, i))) ? (disp ? "float" : "float16_t") + : "float") << " b" << i << "[]; };\n"; if (disp) - L << "layout(std430, binding = " << (n_inputs + i) - << ") readonly buffer Buf" << i << "h { float16_t b" << i << "h[]; };\n"; + L << "layout(std430, binding = " << (n_inputs + i) << ") readonly buffer Buf" << i + << "h { float16_t b" << i << "h[]; };\n"; } if (disp) { @@ -632,10 +596,8 @@ inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec else // k_split 时输出 = 部分和缓冲（后端按 f32 分配）→ 声明恒 float //（见本函数头 k_split 段：写的是原始部分和，不做精度落盘）。 - L << "layout(std430, binding = " << n_inputs - << ") writeonly buffer BufOut { " - << ((out_f16 && !k_split) ? "float16_t" : "float") - << " bout[]; };\n"; + L << "layout(std430, binding = " << n_inputs << ") writeonly buffer BufOut { " + << ((out_f16 && !k_split) ? "float16_t" : "float") << " bout[]; };\n"; // vec4 别名视图（A/B 全局快路径专用；同 binding 双声明 std430，op 级 // matmul_tiled 同款）。a_slot==b_slot 时不发别名、快路径整体关闭。 // f16 槽的别名用 uvec2（4×half = 8B，std430 步长恰好 8）→ 由 @@ -656,9 +618,8 @@ inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec } else { - const auto alias_type = [&](std::uint32_t slot) -> const char* { - return (sign_f16 && expr_prec_sig_in_f16(sig, slot)) ? "uvec2" : "vec4"; - }; + const auto alias_type = [&](std::uint32_t slot) -> const char* + { return (sign_f16 && expr_prec_sig_in_f16(sig, slot)) ? "uvec2" : "vec4"; }; L << "layout(std430, binding = " << a_slot << ") readonly buffer BufAv4 { " << alias_type(a_slot) << " b" << a_slot << "v4[]; };\n"; L << "layout(std430, binding = " << b_slot << ") readonly buffer BufBv4 { " @@ -667,13 +628,15 @@ inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec } L << "\n"; // 别名槽的 vec4 读取表达式（f16 → unpackHalf2x16；f32 → 直接别名） - const auto vec4_load = [&](std::uint32_t slot, const std::string& idx) -> std::string { + const auto vec4_load = [&](std::uint32_t slot, const std::string& idx) -> std::string + { const std::string arr = "b" + std::to_string(slot) + "v4[" + idx + "]"; if (disp) - return "((prec >> " + std::to_string(slot) + "u & 1u) != 0u ? vec4(" - "unpackHalf2x16(b" + std::to_string(slot) + "v4h[" + idx - + "].x), unpackHalf2x16(b" + std::to_string(slot) + "v4h[" + idx - + "].y)) : " + arr + ")"; + return "((prec >> " + std::to_string(slot) + + "u & 1u) != 0u ? vec4(" + "unpackHalf2x16(b" + + std::to_string(slot) + "v4h[" + idx + "].x), unpackHalf2x16(b" + + std::to_string(slot) + "v4h[" + idx + "].y)) : " + arr + ")"; if (sign_f16 && expr_prec_sig_in_f16(sig, slot)) return "vec4(unpackHalf2x16(" + arr + ".x), unpackHalf2x16(" + arr + ".y))"; return arr; @@ -688,8 +651,7 @@ inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec if (k_split) L << " uint ks_span; // K 切分段宽（段 k 上界 = 段起点 + 本槽；后端保证 " "16 对齐 → vec4 快路径 k 连续判据成立）\n"; - if (disp) - L << " uint prec; // 运行期精度分派位（bit i=输入 i 为 f16，bit16=输出）\n"; + if (disp) L << " uint prec; // 运行期精度分派位（bit i=输入 i 为 f16，bit16=输出）\n"; const std::uint32_t n_vp = expr_spec_runtime_view_param_count(spec); for (std::uint32_t i = 0; i < n_vp; ++i) L << " uint vp" << i << ";\n"; @@ -721,29 +683,27 @@ inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec // 行偏移 / batch*mm_k 为 B k 偏移；row/col 为 **batch 内**坐标，block // 偏移由调用方并入）。trans 由调用方按运行期 mm_trans 的两条分支分别传 // 常量（true/false）——两条路径都发射，运行时经 uniform 分支选择 ── - const auto a_load = [&](const std::string& row_e, const std::string& k_e, - bool trans) + const auto a_load = [&](const std::string& row_e, const std::string& k_e, bool trans) { - const std::string idx = trans - ? "((" + a_blk + "*mm_k + (" + k_e + "))*m_per + (" + row_e + "))" - : "((" + a_blk + "*m_per + (" + row_e + "))*mm_k + (" + k_e + "))"; + const std::string idx = + trans ? "((" + a_blk + "*mm_k + (" + k_e + "))*m_per + (" + row_e + "))" + : "((" + a_blk + "*m_per + (" + row_e + "))*mm_k + (" + k_e + "))"; const std::string an = "b" + std::to_string(a_slot); if (disp) - return "((prec >> " + std::to_string(a_slot) + "u & 1u) != 0u ? float(" - + an + "h[" + idx + "]) : " + an + "[" + idx + "])"; + return "((prec >> " + std::to_string(a_slot) + "u & 1u) != 0u ? float(" + an + "h[" + + idx + "]) : " + an + "[" + idx + "])"; const std::string raw = an + "[" + idx + "]"; return a_f16 ? ("float(" + raw + ")") : raw; }; - const auto b_load = [&](const std::string& col_e, const std::string& k_e, - bool trans) + const auto b_load = [&](const std::string& col_e, const std::string& k_e, bool trans) { - const std::string idx = trans - ? "((" + b_blk + "*cols + (" + col_e + "))*mm_k + (" + k_e + "))" - : "((" + b_blk + "*mm_k + (" + k_e + "))*cols + (" + col_e + "))"; + const std::string idx = + trans ? "((" + b_blk + "*cols + (" + col_e + "))*mm_k + (" + k_e + "))" + : "((" + b_blk + "*mm_k + (" + k_e + "))*cols + (" + col_e + "))"; const std::string bn = "b" + std::to_string(b_slot); if (disp) - return "((prec >> " + std::to_string(b_slot) + "u & 1u) != 0u ? float(" - + bn + "h[" + idx + "]) : " + bn + "[" + idx + "])"; + return "((prec >> " + std::to_string(b_slot) + "u & 1u) != 0u ? float(" + bn + "h[" + + idx + "]) : " + bn + "[" + idx + "])"; const std::string raw = bn + "[" + idx + "]"; return b_f16 ? ("float(" + raw + ")") : raw; }; @@ -761,8 +721,7 @@ inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec // 标量回退 → 每个 (k,m) 恰被写一次（快路径与回退的写集合相同）。共享槽位恒用 // tile 内坐标，越界判/全局地址才用 block_row/col + tile 偏移。 L << "void load_tiles(uint t, uint stage)\n{\n"; - L << " const uint tid = gl_LocalInvocationID.y * " << T - << "u + gl_LocalInvocationID.x;\n"; + L << " const uint tid = gl_LocalInvocationID.y * " << T << "u + gl_LocalInvocationID.x;\n"; L << " const uint batch = gl_WorkGroupID.z;\n"; L << " const uint m_per = rows / mm_batch;\n"; L << " const uint block_row = gl_WorkGroupID.y * " << B << "u;\n"; @@ -925,10 +884,7 @@ inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec // 累加器：R=4 → vec4 acc[4]（16 输出/线程，每 k 4 条 VFMA.128）； // R=2 → 4 个标量（2×2 输出/线程，4 条独立 K 链 = 线程内 ILP）。 // 两者每输出都是逐 k 标量升序乘加 → 累加序一致（逐位一致）。 - if (R == 2) - { - L << " float a00 = 0.0, a01 = 0.0, a10 = 0.0, a11 = 0.0;\n"; - } + if (R == 2) { L << " float a00 = 0.0, a01 = 0.0, a10 = 0.0, a11 = 0.0;\n"; } else { L << " vec4 acc[4];\n"; @@ -970,15 +926,15 @@ inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec } else { - L << " for (uint k = 0u; k < " << BK << "u; ++k)\n"; - L << " {\n"; - L << " const vec4 av = Ash[stage][k][ty];\n"; - L << " const vec4 bv = Bsh[stage][k][tx];\n"; - L << " acc[0] += av.x * bv;\n"; - L << " acc[1] += av.y * bv;\n"; - L << " acc[2] += av.z * bv;\n"; - L << " acc[3] += av.w * bv;\n"; - L << " }\n"; + L << " for (uint k = 0u; k < " << BK << "u; ++k)\n"; + L << " {\n"; + L << " const vec4 av = Ash[stage][k][ty];\n"; + L << " const vec4 bv = Bsh[stage][k][tx];\n"; + L << " acc[0] += av.x * bv;\n"; + L << " acc[1] += av.y * bv;\n"; + L << " acc[2] += av.z * bv;\n"; + L << " acc[3] += av.w * bv;\n"; + L << " }\n"; } L << " barrier();\n"; L << " stage = nstage;\n"; @@ -1035,33 +991,33 @@ inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec } else { - L << " #pragma unroll\n"; - L << " for (uint i = 0u; i < 4u; ++i)\n"; - L << " for (uint j = 0u; j < 4u; ++j)\n"; - L << " {\n"; - L << " const uint rr = block_row + ty * 4u + i;\n"; - L << " const uint cc = block_col + tx * 4u + j;\n"; - L << " if (rr < m_per && cc < cols)\n"; - if (disp) - { - L << " {\n"; - L << " const uint oi = (batch * m_per + rr) * cols + cc;\n"; - L << " if ((prec >> 16u & 1u) != 0u)\n"; - L << " bout16[oi] = " - << glsl_f16_store("eval_tail(acc[i][j], rr, cc, batch)") << ";\n"; - L << " else\n"; - L << " bout[oi] = eval_tail(acc[i][j], rr, cc, batch);\n"; - L << " }\n"; - } - else - { - L << " bout[(batch * m_per + rr) * cols + cc]\n"; - L << " = " - << (out_f16 ? glsl_f16_store("eval_tail(acc[i][j], rr, cc, batch)") - : std::string("eval_tail(acc[i][j], rr, cc, batch)")) - << ";\n"; - } - L << " }\n"; + L << " #pragma unroll\n"; + L << " for (uint i = 0u; i < 4u; ++i)\n"; + L << " for (uint j = 0u; j < 4u; ++j)\n"; + L << " {\n"; + L << " const uint rr = block_row + ty * 4u + i;\n"; + L << " const uint cc = block_col + tx * 4u + j;\n"; + L << " if (rr < m_per && cc < cols)\n"; + if (disp) + { + L << " {\n"; + L << " const uint oi = (batch * m_per + rr) * cols + cc;\n"; + L << " if ((prec >> 16u & 1u) != 0u)\n"; + L << " bout16[oi] = " + << glsl_f16_store("eval_tail(acc[i][j], rr, cc, batch)") << ";\n"; + L << " else\n"; + L << " bout[oi] = eval_tail(acc[i][j], rr, cc, batch);\n"; + L << " }\n"; + } + else + { + L << " bout[(batch * m_per + rr) * cols + cc]\n"; + L << " = " + << (out_f16 ? glsl_f16_store("eval_tail(acc[i][j], rr, cc, batch)") + : std::string("eval_tail(acc[i][j], rr, cc, batch)")) + << ";\n"; + } + L << " }\n"; } L << "}\n"; return L.str(); @@ -1088,8 +1044,9 @@ inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec // （浮点加法不结合）→ 移动 GPU 字节锚；裁定放行（torch 文档同样不保证 // 数学等价计算逐位一致，见 history「K 切分」条）。 // ═══════════════════════════════════════════════════════════════════════════ -[[nodiscard]] inline std::string generate_glsl_matmul_ksplit_reduce( - const std::string& name, const ExprSpec& spec, ExprPrecSig sig = 0) +[[nodiscard]] inline std::string generate_glsl_matmul_ksplit_reduce(const std::string& name, + const ExprSpec& spec, + ExprPrecSig sig = 0) { // 资格 = 与 #ks/细 tile 同源：纯 matmul、无 fold、raxis<0、**静态签名** // （sig==0 或具体签名；运行期分派哨兵与 #ks 同款被拒——它的 `prec` PC @@ -1100,8 +1057,7 @@ inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec expr_spec_reduce_axis(spec) != -1 || !spec.extras.empty()) return {}; const std::size_t n_inputs = spec.views.size(); - if (n_inputs + 2 > EXPR_MAX_INPUTS + 2) - return {}; + if (n_inputs + 2 > EXPR_MAX_INPUTS + 2) return {}; // 输入/输出按注入签名选元素类型；**部分和缓冲恒 f32**（后端 use_ks 分支 // 按 f32 分配，与 sig 的输出精度无关）。尾链的读取表达式由 @@ -1113,21 +1069,17 @@ inline void glsl_emit_matmul_tail_fn(std::ostringstream& L, const ExprSpec& spec L << "// ── 自动生成（AOT 算子融合 · matmul K 切分二次归约），请勿手动编辑 ──\n"; L << "// 表达式: " << name << " [K 切分归约 #kr：段升序求和 + 尾链]\n"; L << "#version 450\n\n"; - if (sign_f16) - L << "#extension GL_EXT_shader_16bit_storage : require\n\n"; - if (out_f16) - L << glsl_f16_rne_source(); // glsl_f16_store 依赖（RNE 预舍入） + if (sign_f16) L << "#extension GL_EXT_shader_16bit_storage : require\n\n"; + if (out_f16) L << glsl_f16_rne_source(); // glsl_f16_store 依赖（RNE 预舍入） 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 - << " { " - << ((sign_f16 && expr_prec_sig_in_f16(sig, i)) ? "float16_t" : "float") - << " b" << i << "[]; };\n"; + L << "layout(std430, binding = " << i << ") readonly buffer Buf" << i << " { " + << ((sign_f16 && expr_prec_sig_in_f16(sig, i)) ? "float16_t" : "float") << " b" << i + << "[]; };\n"; L << "layout(std430, binding = " << n_inputs << ") readonly buffer BufPart { float bpart[]; };\n"; - L << "layout(std430, binding = " << (n_inputs + 1) - << ") writeonly buffer BufOut { " << (out_f16 ? "float16_t" : "float") - << " bout[]; };\n"; + L << "layout(std430, binding = " << (n_inputs + 1) << ") writeonly buffer BufOut { " + << (out_f16 ? "float16_t" : "float") << " bout[]; };\n"; L << "\n"; // 分组归约视图辅助（尾链含 GroupedReduce 视图时需要；纯 Linear 时为空） glsl_emit_grouped_reduce_helpers(L, spec, sig, /*f16_native=*/false); @@ -1196,12 +1148,10 @@ inline std::string generate_glsl_fold(const std::string& name, const ExprSpec& s ExprPrecSig sig = 0) { // 形态由 validate_expr_spec 保证；防御性拒绝一切越界组合 - if (!spec.fold || spec.matmul || !spec.instrs.empty()) - return {}; + if (!spec.fold || spec.matmul || !spec.instrs.empty()) return {}; const FoldSpec& f = *spec.fold; auto cls = expr_fold_classify(f, spec.num_regs); - if (!cls) - return {}; + if (!cls) return {}; const std::vector& is_elem = *cls; const std::size_t n_inputs = spec.views.size(); // 带类型变体：f16 输入直读（float() 转 f32 归约——状态/累加全 @@ -1210,14 +1160,11 @@ inline std::string generate_glsl_fold(const std::string& name, const ExprSpec& s const bool disp = expr_prec_sig_is_dispatch(sig); const bool out_f16 = sig != 0 && expr_prec_sig_out_f16(sig); const auto in_type = [&](std::size_t i) -> const char* - { - return (sig != 0 && expr_prec_sig_in_f16(sig, i)) ? "float16_t" : "float"; - }; + { return (sig != 0 && expr_prec_sig_in_f16(sig, i)) ? "float16_t" : "float"; }; std::vector is_block_reduce(spec.num_regs, 0); for (const auto& ins : f.body) - if (expr_op_is_reduce(static_cast(ins.op))) - is_block_reduce[ins.dst] = 1; + if (expr_op_is_reduce(static_cast(ins.op))) is_block_reduce[ins.dst] = 1; // ── 归约输入内联（通用生成优化）───────────────────────── // 块归约的元素源若**只被归约消费**（可沿 def 链传递），把其定义表达式 @@ -1234,8 +1181,7 @@ inline std::string generate_glsl_fold(const std::string& name, const ExprSpec& s { const auto& ins = f.body[ii]; const ExprOp op = static_cast(ins.op); - const std::size_t nops = expr_op_is_reduce(op) ? 1 - : expr_instr_num_operands(op); + const std::size_t nops = expr_op_is_reduce(op) ? 1 : expr_instr_num_operands(op); const ExprOperand* ops[3] = {&ins.a, &ins.b, &ins.c}; for (std::size_t oi = 0; oi < nops; ++oi) { @@ -1255,11 +1201,9 @@ inline std::string generate_glsl_fold(const std::string& name, const ExprSpec& s changed = false; for (std::uint32_t r = 0; r < spec.num_regs; ++r) { - if (!is_elem[r] || inlineable[r] || use_count[r] != 1) - continue; + if (!is_elem[r] || inlineable[r] || use_count[r] != 1) continue; const ExprInstr& u = f.body[static_cast(single_use[r])]; - const bool via_reduce = - expr_op_is_reduce(static_cast(u.op)); + const bool via_reduce = expr_op_is_reduce(static_cast(u.op)); const bool via_inlined = is_elem[u.dst] && inlineable[u.dst]; if (via_reduce || via_inlined) { @@ -1274,18 +1218,16 @@ inline std::string generate_glsl_fold(const std::string& name, const ExprSpec& s L << "// ── 自动生成（AOT 算子融合 · fold 分块状态归约），请勿手动编辑 ──\n"; L << "// 表达式: " << name << "\n"; L << "#version 450\n\n"; - if (sig != 0) - L << "#extension GL_EXT_shader_16bit_storage : require\n\n"; - if (sig != 0) - L << glsl_f16_rne_source(); + if (sig != 0) L << "#extension GL_EXT_shader_16bit_storage : require\n\n"; + if (sig != 0) L << glsl_f16_rne_source(); 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 - << " { " << (disp ? "float" : in_type(i)) << " b" << i << "[]; };\n"; + L << "layout(std430, binding = " << i << ") readonly buffer Buf" << i << " { " + << (disp ? "float" : in_type(i)) << " b" << i << "[]; };\n"; if (disp) - L << "layout(std430, binding = " << (n_inputs + i) - << ") readonly buffer Buf" << i << "h { float16_t b" << i << "h[]; };\n"; + L << "layout(std430, binding = " << (n_inputs + i) << ") readonly buffer Buf" << i + << "h { float16_t b" << i << "h[]; };\n"; } if (disp) { @@ -1295,17 +1237,15 @@ inline std::string generate_glsl_fold(const std::string& name, const ExprSpec& s << ") writeonly buffer BufOuth { float16_t bout16[]; };\n\n"; } else - L << "layout(std430, binding = " << n_inputs - << ") writeonly buffer BufOut { " << (out_f16 ? "float16_t" : "float") - << " bout[]; };\n\n"; + L << "layout(std430, binding = " << n_inputs << ") writeonly buffer BufOut { " + << (out_f16 ? "float16_t" : "float") << " bout[]; };\n\n"; L << "layout(push_constant) uniform PC {\n"; L << " uint count;\n"; L << " uint cols;\n"; L << " uint rows;\n"; L << " uint vector_out;\n"; L << " uint fold_k;\n"; - if (disp) - L << " uint prec; // 运行期精度分派位（bit i=输入 i 为 f16，bit16=输出）\n"; + if (disp) L << " uint prec; // 运行期精度分派位（bit i=输入 i 为 f16，bit16=输出）\n"; const std::uint32_t n_vp = expr_spec_runtime_view_param_count(spec); for (std::uint32_t i = 0; i < n_vp; ++i) L << " uint vp" << i << ";\n"; @@ -1323,22 +1263,18 @@ inline std::string generate_glsl_fold(const std::string& name, const ExprSpec& s // （可内联元素寄存器不物化：消费点归约循环内逐 kb 重算） for (std::uint32_t r = 0; r < spec.num_regs; ++r) { - if (is_elem[r] && inlineable[r]) - continue; - if (is_elem[r]) - L << " float e" << r << "[" << EXPR_FOLD_BLOCK << "];\n"; - else - L << " float r" << r << ";\n"; + if (is_elem[r] && inlineable[r]) continue; + if (is_elem[r]) L << " float e" << r << "[" << EXPR_FOLD_BLOCK << "];\n"; + else L << " float r" << r << ";\n"; } for (std::uint32_t r = 0; r < spec.num_regs; ++r) - if (is_block_reduce[r]) - L << " float red" << r << ";\n"; + if (is_block_reduce[r]) L << " float red" << r << ";\n"; // 状态初值（bit_cast → uintBitsToFloat：±inf 与任意 float 均零精度损失） for (std::uint8_t s = 0; s < f.num_state; ++s) { const std::uint32_t bits = std::bit_cast(f.inits[s]); - L << " r" << static_cast(s) << " = uintBitsToFloat(0x" - << std::hex << bits << std::dec << "u);\n"; + L << " r" << static_cast(s) << " = uintBitsToFloat(0x" << std::hex << bits + << std::dec << "u);\n"; } // 操作数 → GLSL 片段；kbv = 块内下标变量（标量上下文传 "0u"——类别分析 @@ -1350,9 +1286,8 @@ inline std::string generate_glsl_fold(const std::string& name, const ExprSpec& s default: case static_cast(ExprOperandKind::Reg): case static_cast(ExprOperandKind::Fanout): - return is_elem[op.idx] - ? ("e" + std::to_string(op.idx) + "[" + kbv + "]") - : ("r" + std::to_string(op.idx)); + return is_elem[op.idx] ? ("e" + std::to_string(op.idx) + "[" + kbv + "]") + : ("r" + std::to_string(op.idx)); case static_cast(ExprOperandKind::Const): return "c" + std::to_string(op.idx); case static_cast(ExprOperandKind::RParam): @@ -1365,8 +1300,7 @@ inline std::string generate_glsl_fold(const std::string& name, const ExprSpec& s std::ostringstream os; const bool f16_in = sig != 0 && expr_prec_sig_in_f16(sig, op.idx); if (f16_in) os << "float("; - glsl_view_read(os, spec.views[op.idx], op.idx, - "(row * cols + " + gk + ")", "row", gk, + glsl_view_read(os, spec.views[op.idx], op.idx, "(row * cols + " + gk + ")", "row", gk, glsl_vp_slot(spec, op.idx), sig); if (f16_in) os << ")"; return os.str(); @@ -1374,8 +1308,8 @@ inline std::string generate_glsl_fold(const std::string& name, const ExprSpec& s } }; // 指令 → 赋值语句（dst_str 已含下标/变量名），直接写入 L - const auto emit_assign = [&](const ExprInstr& ins, const std::string& dst_str, - const std::string& kbv) + const auto emit_assign = + [&](const ExprInstr& ins, const std::string& dst_str, const std::string& kbv) { const ExprOp op = static_cast(ins.op); const std::string a = operand(ins.a, kbv); @@ -1387,9 +1321,9 @@ inline std::string generate_glsl_fold(const std::string& name, const ExprSpec& s return; } L << dst_str << " = " - << glsl_alu_expr(info, a, [&] { return operand(ins.b, kbv); }, - [&] { return operand(ins.c, kbv); }, - /*vec4=*/false, /*native16=*/false, /*wrap=*/false) + << glsl_alu_expr( + info, a, [&] { return operand(ins.b, kbv); }, [&] { return operand(ins.c, kbv); }, + /*vec4=*/false, /*native16=*/false, /*wrap=*/false) << ";\n"; }; @@ -1397,8 +1331,7 @@ inline std::string generate_glsl_fold(const std::string& name, const ExprSpec& s const auto def_of = [&](std::uint32_t r) -> const ExprInstr* { for (const auto& d : f.body) - if (d.dst == r) - return &d; + if (d.dst == r) return &d; return nullptr; }; std::function opnd_inl; @@ -1409,8 +1342,7 @@ inline std::string generate_glsl_fold(const std::string& name, const ExprSpec& s o.kind == static_cast(ExprOperandKind::Fanout)) && o.idx < spec.num_regs && is_elem[o.idx] && inlineable[o.idx]) { - if (const ExprInstr* d = def_of(o.idx)) - return expr_inl(*d, kbv); + if (const ExprInstr* d = def_of(o.idx)) return expr_inl(*d, kbv); } return operand(o, kbv); }; @@ -1422,9 +1354,9 @@ inline std::string generate_glsl_fold(const std::string& name, const ExprSpec& s const detail::ExprOpInfo& info = detail::expr_op_info(op); if (info.emit == detail::ExprEmitKind::None) return "(" + a + ")"; // 归约指令：不经内联（调用方分流） - return glsl_alu_expr(info, a, [&] { return opnd_inl(ins.b, kbv); }, - [&] { return opnd_inl(ins.c, kbv); }, - /*vec4=*/false, /*native16=*/false, /*wrap=*/true); + return glsl_alu_expr( + info, a, [&] { return opnd_inl(ins.b, kbv); }, [&] { return opnd_inl(ins.c, kbv); }, + /*vec4=*/false, /*native16=*/false, /*wrap=*/true); }; // ── 块循环（k0 步进 BLOCK；valid gate 尾块——与 CPU 同序）────────── @@ -1449,12 +1381,9 @@ inline std::string generate_glsl_fold(const std::string& name, const ExprSpec& s const bool src_inlined = (ins.a.kind == static_cast(ExprOperandKind::Reg) || ins.a.kind == static_cast(ExprOperandKind::Fanout)) && - ins.a.idx < spec.num_regs && is_elem[ins.a.idx] && - inlineable[ins.a.idx]; - const auto src_at = [&](const std::string& kbv) { - return src_inlined ? expr_inl(*def_of(ins.a.idx), kbv) - : operand(ins.a, kbv); - }; + ins.a.idx < spec.num_regs && is_elem[ins.a.idx] && inlineable[ins.a.idx]; + const auto src_at = [&](const std::string& kbv) + { return src_inlined ? expr_inl(*def_of(ins.a.idx), kbv) : operand(ins.a, kbv); }; // ── 循环剥离 + 预取（软件流水，通用）───────────────────────── // 直写 `acc = acc + f(load(kb))` 是逐迭代 load→f→add 的**串行依赖 // 链**（动态 valid 抑制编译器自动流水）：实测 fold v1 因此比同工作 @@ -1462,18 +1391,15 @@ inline std::string generate_glsl_fold(const std::string& name, const ExprSpec& s // 不了）。剥离首迭代、把 kb+1 的源表达式先算进 vn 再累加 v（上一 // 迭代值）——访存/表达式求值与累加重叠。**累加序仍为 kb 升序逐个 // 累加 → 与 CPU eval_fold 逐位一致**（铁律 8）。 - L << " { float acc = " - << (is_max ? "uintBitsToFloat(0xFF7FFFFFu)" : "0.0") + L << " { float acc = " << (is_max ? "uintBitsToFloat(0xFF7FFFFFu)" : "0.0") << ";\n"; L << " float v = " << src_at("0u") << ";\n"; L << " for (uint kb = 1u; kb < valid; ++kb) {\n"; L << " const float vn = " << src_at("kb") << ";\n"; - L << " acc = " - << (is_max ? "max(acc, v)" : "(acc + v)") << ";\n"; + L << " acc = " << (is_max ? "max(acc, v)" : "(acc + v)") << ";\n"; L << " v = vn;\n"; L << " }\n"; - L << " acc = " - << (is_max ? "max(acc, v)" : "(acc + v)") << ";\n"; + L << " acc = " << (is_max ? "max(acc, v)" : "(acc + v)") << ";\n"; L << " red" << dst_r << " = acc; }\n"; } else if (!is_elem[ins.dst]) @@ -1506,16 +1432,15 @@ inline std::string generate_glsl_fold(const std::string& name, const ExprSpec& s { L << " if ((prec >> 16u & 1u) != 0u)\n"; L << " bout16[row] = " - << glsl_f16_store("r" + std::to_string(static_cast(f.finalize.back().dst))) - << ";\n"; + << glsl_f16_store("r" + std::to_string(static_cast(f.finalize.back().dst))) << ";\n"; L << " else\n"; L << " bout[row] = r" << static_cast(f.finalize.back().dst) << ";\n"; } else L << " bout[row] = " - << (out_f16 ? glsl_f16_store( - "r" + std::to_string(static_cast(f.finalize.back().dst))) - : ("r" + std::to_string(static_cast(f.finalize.back().dst)))) + << (out_f16 + ? glsl_f16_store("r" + std::to_string(static_cast(f.finalize.back().dst))) + : ("r" + std::to_string(static_cast(f.finalize.back().dst)))) << ";\n"; L << "}\n"; return L.str(); @@ -1548,12 +1473,10 @@ inline std::string generate_glsl_fold(const std::string& name, const ExprSpec& s inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec& spec, ExprPrecSig sig = 0) { - if (!spec.fold || spec.matmul || !spec.instrs.empty()) - return {}; + if (!spec.fold || spec.matmul || !spec.instrs.empty()) return {}; const FoldSpec& f = *spec.fold; auto cls = expr_fold_classify(f, spec.num_regs); - if (!cls) - return {}; + if (!cls) return {}; const std::vector& is_elem = *cls; std::vector eidx(spec.num_regs, -1); int ne = 0; @@ -1574,9 +1497,7 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec const bool disp = expr_prec_sig_is_dispatch(sig); const bool out_f16 = sig != 0 && expr_prec_sig_out_f16(sig); const auto in_type = [&](std::size_t i) -> const char* - { - return (sig != 0 && expr_prec_sig_in_f16(sig, i)) ? "float16_t" : "float"; - }; + { return (sig != 0 && expr_prec_sig_in_f16(sig, i)) ? "float16_t" : "float"; }; // 直接下标读（mm 段 A/B、vecacc B——非视图，平铺索引）：f16 输入包 float()。 // 入参是**纯下标表达式**（如 "(batch*mm_k + d)*cols + col"），缓冲名由本 // 函数按分派位拼出 f32 / f16 两个视图——此前实现是"在下标串里做子串 @@ -1587,10 +1508,9 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec const std::string an = "b" + std::to_string(i); const std::string f32e = an + "[" + idx + "]"; if (disp) - return "((prec >> " + std::to_string(i) + "u & 1u) != 0u ? float(" - + an + "h[" + idx + "]) : " + f32e + ")"; - return (sig != 0 && expr_prec_sig_in_f16(sig, i)) ? ("float(" + f32e + ")") - : f32e; + return "((prec >> " + std::to_string(i) + "u & 1u) != 0u ? float(" + an + "h[" + idx + + "]) : " + f32e + ")"; + return (sig != 0 && expr_prec_sig_in_f16(sig, i)) ? ("float(" + f32e + ")") : f32e; }; // 输出列数 **运行时读 PC 的 vector_out 槽**（形状无关：同 key 服务任意 // d_k——veclen 若进 key 则每个 dk 一个 shader、闭合世界必然缺登记） @@ -1604,18 +1524,16 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec // 书面前提）；`: require` 与 matmul_gemv 的限定符口径统一 L << "#extension GL_KHR_shader_subgroup_basic : require\n"; // gl_SubgroupSize L << "#extension GL_KHR_shader_subgroup_shuffle_relative : require\n\n"; // subgroupShuffleDown - if (sig != 0) - L << "#extension GL_EXT_shader_16bit_storage : require\n\n"; - if (sig != 0) - L << glsl_f16_rne_source(); + if (sig != 0) L << "#extension GL_EXT_shader_16bit_storage : require\n\n"; + if (sig != 0) L << glsl_f16_rne_source(); 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 - << " { " << (disp ? "float" : in_type(i)) << " b" << i << "[]; };\n"; + L << "layout(std430, binding = " << i << ") readonly buffer Buf" << i << " { " + << (disp ? "float" : in_type(i)) << " b" << i << "[]; };\n"; if (disp) - L << "layout(std430, binding = " << (n_inputs + i) - << ") readonly buffer Buf" << i << "h { float16_t b" << i << "h[]; };\n"; + L << "layout(std430, binding = " << (n_inputs + i) << ") readonly buffer Buf" << i + << "h { float16_t b" << i << "h[]; };\n"; } if (disp) { @@ -1625,16 +1543,13 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec << ") writeonly buffer BufOuth { float16_t bout16[]; };\n\n"; } else - L << "layout(std430, binding = " << n_inputs - << ") writeonly buffer BufOut { " << (out_f16 ? "float16_t" : "float") - << " bout[]; };\n\n"; + L << "layout(std430, binding = " << n_inputs << ") writeonly buffer BufOut { " + << (out_f16 ? "float16_t" : "float") << " bout[]; };\n\n"; L << "layout(push_constant) uniform PC {\n"; L << " uint count;\n uint cols;\n uint rows;\n" " uint vector_out;\n uint fold_k;\n"; - if (has_mm) - L << " uint mm_k;\n uint mm_batch;\n"; - if (disp) - L << " uint prec; // 运行期精度分派位（bit i=输入 i 为 f16，bit16=输出）\n"; + if (has_mm) L << " uint mm_k;\n uint mm_batch;\n"; + if (disp) L << " uint prec; // 运行期精度分派位（bit i=输入 i 为 f16，bit16=输出）\n"; const std::uint32_t n_vp = expr_spec_runtime_view_param_count(spec); for (std::uint32_t i = 0; i < n_vp; ++i) L << " uint vp" << i << ";\n"; @@ -1646,7 +1561,7 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec if (ne > 0) { L << "shared float sreg[" << ne << "][" << EXPR_FOLD_BLOCK << "];\n"; - L << "shared float sred[1];\n"; // subgroup 归约 lane0 落点 → barrier 广播 + L << "shared float sred[1];\n"; // subgroup 归约 lane0 落点 → barrier 广播 } if (has_mm) { @@ -1655,8 +1570,7 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec // jm = tid & (BLOCK-1)、grp = tid >> log2(BLOCK)（组数 = 256/BLOCK）、 // d 步进 = 256/BLOCK。 L << "shared float Ash[" << FOLD_MAX_MMK << "];\n"; - L << "shared float spart[" << EXPR_FOLD_BLOCK << "][" - << (256u / EXPR_FOLD_BLOCK) << "];\n"; + L << "shared float spart[" << EXPR_FOLD_BLOCK << "][" << (256u / EXPR_FOLD_BLOCK) << "];\n"; L << "shared float smm[" << EXPR_FOLD_BLOCK << "];\n"; } if (f.vecacc && !has_mm) @@ -1675,8 +1589,7 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec // 跨行复用 → shared 零增长。ri 循环边界对全 WG 均匀（体内含屏障， // 禁分支发散）；越界行 clamp 到末行重复算、写回由 row_ok 挡掉。── L << " for (uint ri = 0u; ri < " << EXPR_FOLD_ROWS_PER_WG << "u; ++ri) {\n"; - L << " const uint row_r = gl_WorkGroupID.x * " << EXPR_FOLD_ROWS_PER_WG - << "u + ri;\n"; + L << " const uint row_r = gl_WorkGroupID.x * " << EXPR_FOLD_ROWS_PER_WG << "u + ri;\n"; L << " const uint row = min(row_r, rows - 1u);\n"; L << " const bool row_ok = row_r < rows;\n"; if (has_mm) @@ -1689,11 +1602,11 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec // GQA：A 侧行块号——HeadGroup 视图时按 vp 槽里的头分组比取（1 = 恒等） const std::string a_blk = gqa_blk_expr(spec, f.matmul->a_input); if (f.matmul->transA) - L << " Ash[d] = " << in_rd(f.matmul->a_input, - "(" + a_blk + " * mm_k + d) * m_per + row_in") << ";\n"; + L << " Ash[d] = " + << in_rd(f.matmul->a_input, "(" + a_blk + " * mm_k + d) * m_per + row_in") << ";\n"; else - L << " Ash[d] = " << in_rd(f.matmul->a_input, - "(" + a_blk + " * m_per + row_in) * mm_k + d") << ";\n"; + L << " Ash[d] = " + << in_rd(f.matmul->a_input, "(" + a_blk + " * m_per + row_in) * mm_k + d") << ";\n"; L << " }\n"; L << " barrier(); // Ash 就绪（每行预载，块循环前一次）\n"; } @@ -1701,8 +1614,8 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec for (std::uint8_t s = 0; s < f.num_state; ++s) { const std::uint32_t bits = std::bit_cast(f.inits[s]); - L << " r" << static_cast(s) << " = uintBitsToFloat(0x" - << std::hex << bits << std::dec << "u);\n"; + L << " r" << static_cast(s) << " = uintBitsToFloat(0x" << std::hex << bits + << std::dec << "u);\n"; } // 行向量态（P2-2 双形态）： // 路径 B（0 std::string + const auto operand = [&](const ExprOperand& op, char ctx, const std::string& jv) -> std::string { switch (op.kind) { @@ -1743,7 +1655,7 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec case static_cast(ExprOperandKind::RParam): return "rp" + std::to_string(op.idx); case static_cast(ExprOperandKind::Reduce): - return "r" + std::to_string(op.idx); // 归约 dst 每线程副本 + return "r" + std::to_string(op.idx); // 归约 dst 每线程副本 case static_cast(ExprOperandKind::Matmul): // mm 段结果已在块首协作算入 smm[]（见块循环头）——链内直接读， // 不做现场串行 Σ_d（现场串行结构是吞吐回退主因） @@ -1768,8 +1680,7 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec std::ostringstream os; const bool f16_in = sig != 0 && expr_prec_sig_in_f16(sig, op.idx); if (f16_in) os << "float("; - glsl_view_read(os, spec.views[op.idx], op.idx, - "(row * cols + " + gk + ")", "row", gk, + glsl_view_read(os, spec.views[op.idx], op.idx, "(row * cols + " + gk + ")", "row", gk, glsl_vp_slot(spec, op.idx), sig); if (f16_in) os << ")"; return os.str(); @@ -1777,25 +1688,25 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec } }; // 指令 → 赋值语句（dst_str = 目标表达式） - const auto emit_assign = [&](const ExprInstr& ins, const std::string& dst_str, - char ctx, const std::string& jv) -> std::string + const auto emit_assign = [&](const ExprInstr& ins, const std::string& dst_str, char ctx, + const std::string& jv) -> std::string { const ExprOp op = static_cast(ins.op); const std::string a = operand(ins.a, ctx, jv); // 发射形态/符号查表（expr_op_table） const detail::ExprOpInfo& info = detail::expr_op_info(op); if (info.emit == detail::ExprEmitKind::None) - return " " + dst_str + " = (" + a + ");\n"; // 归约：调用方分流 - return " " + dst_str + " = " - + glsl_alu_expr(info, a, [&] { return operand(ins.b, ctx, jv); }, - [&] { return operand(ins.c, ctx, jv); }, - /*vec4=*/false, /*native16=*/false, /*wrap=*/false) - + ";\n"; + return " " + dst_str + " = (" + a + ");\n"; // 归约：调用方分流 + return " " + dst_str + " = " + + glsl_alu_expr( + info, a, [&] { return operand(ins.b, ctx, jv); }, + [&] { return operand(ins.c, ctx, jv); }, + /*vec4=*/false, /*native16=*/false, /*wrap=*/false) + + ";\n"; }; // ── 块循环 + body 段扫描发射 ───────────────────────────────────────── - L << "\n for (uint k0 = 0u; k0 < fold_k; k0 += " << EXPR_FOLD_BLOCK - << "u) {\n"; + L << "\n for (uint k0 = 0u; k0 < fold_k; k0 += " << EXPR_FOLD_BLOCK << "u) {\n"; if (f.tri_skip) { // 行界整块/边界跳过（qt = row%m_per，与链内 select 谓词同源）： @@ -1818,19 +1729,23 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec } else { - L << " const uint valid = min(" << EXPR_FOLD_BLOCK - << "u, fold_k - k0);\n"; + L << " const uint valid = min(" << EXPR_FOLD_BLOCK << "u, fold_k - k0);\n"; } if (has_mm) { // mm 段收缩协作化（映射随 BLOCK 参数化，见 shared 声明处注释） const MatmulSpec& m = *f.matmul; - constexpr std::uint32_t LG = []() { + constexpr std::uint32_t LG = []() + { std::uint32_t v = EXPR_FOLD_BLOCK, g = 0; - while (v > 1) { v >>= 1; ++g; } + while (v > 1) + { + v >>= 1; + ++g; + } return g; }(); - constexpr std::uint32_t NG = 256u / EXPR_FOLD_BLOCK; // d 组数 + constexpr std::uint32_t NG = 256u / EXPR_FOLD_BLOCK; // d 组数 L << " const uint jm = tid & " << (EXPR_FOLD_BLOCK - 1) << "u;\n"; L << " const uint grp = tid >> " << LG << "u;\n"; L << " float part = 0.0;\n"; @@ -1839,11 +1754,11 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec // GQA：B 侧行块号——HeadGroup 视图时按 vp 槽里的头分组比取 const std::string b_blk = gqa_blk_expr(spec, m.b_input); if (m.transB) - L << " part += Ash[d] * " << in_rd(m.b_input, - "(" + b_blk + " * cols + (k0 + jm)) * mm_k + d") << ";\n"; + L << " part += Ash[d] * " + << in_rd(m.b_input, "(" + b_blk + " * cols + (k0 + jm)) * mm_k + d") << ";\n"; else - L << " part += Ash[d] * " << in_rd(m.b_input, - "(" + b_blk + " * mm_k + d) * cols + (k0 + jm)") << ";\n"; + L << " part += Ash[d] * " + << in_rd(m.b_input, "(" + b_blk + " * mm_k + d) * cols + (k0 + jm)") << ";\n"; L << " }\n"; L << " }\n"; L << " spart[jm][grp] = part;\n"; @@ -1858,8 +1773,7 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec std::vector chain; const auto flush_chain = [&]() { - if (chain.empty()) - return; + if (chain.empty()) return; // 元素 j 循环：**按读集载入、按写集存回**（不再是全元素寄存器对称 // round-trip）。动机：注意力 chain1 实测 9 次 sreg 载入全为死加载 // （全部覆写前载入 / 本段未用）、chain2 为 9 中 8 死——纯浪费 @@ -1875,12 +1789,10 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec std::vector wr(spec.num_regs, false); const auto note_read = [&](const ExprOperand& op) { - const bool as_local = - op.kind == static_cast(ExprOperandKind::Reg) || - op.kind == static_cast(ExprOperandKind::Fanout) || - op.kind == static_cast(ExprOperandKind::Reduce); - if (as_local && is_elem[op.idx] && !wr[op.idx]) - need_load[op.idx] = true; + const bool as_local = op.kind == static_cast(ExprOperandKind::Reg) || + op.kind == static_cast(ExprOperandKind::Fanout) || + op.kind == static_cast(ExprOperandKind::Reduce); + if (as_local && is_elem[op.idx] && !wr[op.idx]) need_load[op.idx] = true; }; for (const auto& ins : chain) { @@ -1892,11 +1804,9 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec L << " for (uint j = tid; j < valid; j += 256u) {\n"; for (std::uint32_t r = 0; r < spec.num_regs; ++r) { - if (!is_elem[r] || (!need_load[r] && !wr[r])) - continue; + if (!is_elem[r] || (!need_load[r] && !wr[r])) continue; L << " float r" << r; - if (need_load[r]) - L << " = sreg[" << eidx[r] << "][j]"; + if (need_load[r]) L << " = sreg[" << eidx[r] << "][j]"; L << ";\n"; } for (const auto& ins : chain) @@ -1905,8 +1815,7 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec L << emit_assign(ins, dst, 'E', "j"); } for (std::uint32_t r = 0; r < spec.num_regs; ++r) - if (wr[r]) - L << " sreg[" << eidx[r] << "][j] = r" << r << ";\n"; + if (wr[r]) L << " sreg[" << eidx[r] << "][j] = r" << r << ";\n"; L << " }\n"; L << " barrier();\n"; chain.clear(); @@ -1919,9 +1828,9 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec flush_chain(); const bool is_max = (op == ExprOp::RowMax || op == ExprOp::ColMax); const auto skind = static_cast(ins.a.kind); - const bool elem_src = (skind == ExprOperandKind::Reg || - skind == ExprOperandKind::Fanout) - && is_elem[ins.a.idx]; + const bool elem_src = + (skind == ExprOperandKind::Reg || skind == ExprOperandKind::Fanout) && + is_elem[ins.a.idx]; if (elem_src) { // 归约本体（subgroup shuffle 蝶式，fold v2 吞吐路径）； @@ -1947,15 +1856,13 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec // eval_fold_impl(numeric_limits::lowest()) 对齐——块内全 -inf // （doc 掩码整块屏蔽）时 m_old=blk_m=-inf 会让 dm=−inf−−inf // =NaN 污染 l/O，见 fold v1 同款注释。 - const std::string init = - is_max ? "uintBitsToFloat(0xFF7FFFFFu)" : "0.0"; + const std::string init = is_max ? "uintBitsToFloat(0xFF7FFFFFu)" : "0.0"; const std::string eis = std::to_string(eidx[ins.a.idx]); L << " float v = " << init << ";\n"; L << " if (tid < gl_SubgroupSize) {\n"; L << " for (uint e = tid; e < valid; e += gl_SubgroupSize)\n"; L << " v = " - << (is_max ? ("max(v, sreg[" + eis + "][e])") - : ("(v + sreg[" + eis + "][e])")) + << (is_max ? ("max(v, sreg[" + eis + "][e])") : ("(v + sreg[" + eis + "][e])")) << ";\n"; L << " }\n"; L << " if (tid < gl_SubgroupSize) {\n"; @@ -1976,13 +1883,12 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec L << " {\n"; // max init = lowest()（0xFF7FFFFF），对齐 CPU 基准/防 −inf−−inf=NaN // ——与上面元素源路径同款理由，两处必须同改 - L << " float acc = " - << (is_max ? "uintBitsToFloat(0xFF7FFFFFu)" : "0.0") << ";\n"; + L << " float acc = " << (is_max ? "uintBitsToFloat(0xFF7FFFFFu)" : "0.0") + << ";\n"; L << " for (uint j = 0u; j < valid; ++j) {\n"; const std::string src = operand(ins.a, 'R', "j"); L << " acc = " - << (is_max ? ("max(acc, " + src + ")") : ("(acc + " + src + ")")) - << ";\n"; + << (is_max ? ("max(acc, " + src + ")") : ("(acc + " + src + ")")) << ";\n"; L << " }\n"; L << " r" << std::to_string(ins.dst) << " = acc;\n"; L << " }\n"; @@ -1990,7 +1896,7 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec } else if (is_elem[ins.dst]) { - chain.push_back(ins); // 元素链（遇标量/归约/尾冲） + chain.push_back(ins); // 元素链（遇标量/归约/尾冲） } else { @@ -2024,18 +1930,17 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec // 可达分片（与旧版一致 × VA_ACC），绝不对 16 槽全 tile 缩放 L << " if (vector_out > 0u && vector_out <= 256u) {\n"; for (std::uint32_t t = 0; t < VA_ACC; ++t) - L << " ob" << t << " *= r" - << static_cast(va.scale_reg) << ";\n"; + L << " ob" << t << " *= r" << static_cast(va.scale_reg) << ";\n"; L << " } else if (vector_out > 256u) {\n"; L << " for (uint dd = tid; dd < vector_out; dd += 256u)\n"; L << " for (uint an = 0u; an < " << VA_ACC << "u; ++an)\n"; - L << " o[(dd - tid) / 256u][an] *= r" - << static_cast(va.scale_reg) << ";\n"; + L << " o[(dd - tid) / 256u][an] *= r" << static_cast(va.scale_reg) + << ";\n"; L << " }\n"; } // GQA：vecacc 的 b_input（V_t）行块号同样按其视图的头分组比取 - const std::string b_row = !has_mm ? std::string("0u") - : gqa_blk_expr(spec, va.b_input, "(row / m_per)"); + const std::string b_row = + !has_mm ? std::string("0u") : gqa_blk_expr(spec, va.b_input, "(row / m_per)"); // vecacc：O += Σ_j w(j)·b(j, d) —— 4 路软件流水。 // 每拍同批发出 j..j+3 四个读使延迟重叠（GPU 无硬件预取，单发 // 串行读会暴露 64 级 L2 延迟链——探针量化其占 fold 开销 55%）； @@ -2053,21 +1958,21 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec L << " uint j = jlo;\n"; L << " for (; j + " << VA_ACC << "u <= jhi; j += " << VA_ACC << "u) {\n"; for (std::uint32_t t = 0; t < VA_ACC; ++t) - L << " const float v" << t << " = " << in_rd(va.b_input, - "(jr0 + j + " + std::to_string(t) + "u) * vector_out + cB") << ";\n"; + L << " const float v" << t << " = " + << in_rd(va.b_input, "(jr0 + j + " + std::to_string(t) + "u) * vector_out + cB") + << ";\n"; for (std::uint32_t t = 0; t < VA_ACC; ++t) - L << " ob" << t << " += sreg[" - << eidx[va.weight_reg] << "][j + " << t << "u] * v" << t << ";\n"; + L << " ob" << t << " += sreg[" << eidx[va.weight_reg] << "][j + " << t + << "u] * v" << t << ";\n"; L << " }\n"; // 尾段（余数 ≤ VA_ACC-1）：主循环出口 (j-jlo) ≡ 0 (mod VA_ACC) → // 按固定累加器序展开（标量无下标，零 dynamic-index） for (std::uint32_t t = 0; t < VA_ACC - 1u; ++t) { L << " if (j < jhi) {\n"; - L << " const float vj = " << in_rd(va.b_input, - "(jr0 + j) * vector_out + cB") << ";\n"; - L << " ob" << t << " += sreg[" << eidx[va.weight_reg] - << "][j] * vj;\n"; + L << " const float vj = " << in_rd(va.b_input, "(jr0 + j) * vector_out + cB") + << ";\n"; + L << " ob" << t << " += sreg[" << eidx[va.weight_reg] << "][j] * vj;\n"; L << " ++j;\n"; L << " }\n"; } @@ -2077,20 +1982,21 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec L << " uint j = 0u;\n"; L << " for (; j + " << VA_ACC << "u <= valid; j += " << VA_ACC << "u) {\n"; for (std::uint32_t t = 0; t < VA_ACC; ++t) - L << " const float v" << t << " = " << in_rd(va.b_input, - "(jr0 + j + " + std::to_string(t) + "u) * vector_out + dd") << ";\n"; + L << " const float v" << t << " = " + << in_rd(va.b_input, "(jr0 + j + " + std::to_string(t) + "u) * vector_out + dd") + << ";\n"; for (std::uint32_t t = 0; t < VA_ACC; ++t) - L << " o[(dd - tid) / 256u][" << t << "u] += sreg[" - << eidx[va.weight_reg] << "][j + " << t << "u] * v" << t << ";\n"; + L << " o[(dd - tid) / 256u][" << t << "u] += sreg[" << eidx[va.weight_reg] + << "][j + " << t << "u] * v" << t << ";\n"; L << " }\n"; L << " for (; j < valid; ++j)\n"; L << " o[(dd - tid) / 256u][j % " << VA_ACC << "u] += sreg[" - << eidx[va.weight_reg] << "][j] * " << in_rd(va.b_input, - "(jr0 + j) * vector_out + dd") << ";\n"; + << eidx[va.weight_reg] << "][j] * " << in_rd(va.b_input, "(jr0 + j) * vector_out + dd") + << ";\n"; L << " }\n"; L << " }\n"; } - L << " }\n"; // 块循环尾 + L << " }\n"; // 块循环尾 // ── vecacc 收尾（k0 循环后、finalize 前；每行一次）── // ① 崩叠：路径 B = 标量累加器 ob 求和（热路径零数组访问，见行首双形态 @@ -2102,16 +2008,21 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec if (f.vecacc) { // 归约槽：spart[jm][grp] 的线性重解释（tid = jm + grp*BLOCK） - constexpr std::uint32_t LG_SLOT = []() { + constexpr std::uint32_t LG_SLOT = []() + { std::uint32_t v = EXPR_FOLD_BLOCK, g = 0; - while (v > 1) { v >>= 1; ++g; } + while (v > 1) + { + v >>= 1; + ++g; + } return g; }(); const auto slot = [&](const std::string& idx) -> std::string { if (has_mm) - return "spart[(" + idx + ") & " + std::to_string(EXPR_FOLD_BLOCK - 1u) - + "u][(" + idx + ") >> " + std::to_string(LG_SLOT) + "u]"; + return "spart[(" + idx + ") & " + std::to_string(EXPR_FOLD_BLOCK - 1u) + "u][(" + + idx + ") >> " + std::to_string(LG_SLOT) + "u]"; return "sva[" + idx + "]"; }; L << "\n if (vector_out > 0u && vector_out <= 256u) {\n"; @@ -2155,21 +2066,20 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec L << " if (row_ok) {\n"; L << " if ((prec >> 16u & 1u) != 0u)\n"; L << " bout16[row * vector_out + dd] = " - << glsl_f16_store("r" + std::to_string(static_cast(f.finalize.back().dst))) - << ";\n"; + << glsl_f16_store("r" + std::to_string(static_cast(f.finalize.back().dst))) << ";\n"; L << " else\n"; - L << " bout[row * vector_out + dd] = r" - << static_cast(f.finalize.back().dst) << ";\n"; + L << " bout[row * vector_out + dd] = r" << static_cast(f.finalize.back().dst) + << ";\n"; L << " }\n"; } else L << " if (row_ok) bout[row * vector_out + dd] = " - << (out_f16 ? glsl_f16_store( - "r" + std::to_string(static_cast(f.finalize.back().dst))) - : ("r" + std::to_string(static_cast(f.finalize.back().dst)))) - << ";\n"; // 越界 ri 不写回 + << (out_f16 + ? glsl_f16_store("r" + std::to_string(static_cast(f.finalize.back().dst))) + : ("r" + std::to_string(static_cast(f.finalize.back().dst)))) + << ";\n"; // 越界 ri 不写回 L << " }\n"; - L << " }\n"; // NR 行循环尾 + L << " }\n"; // NR 行循环尾 L << "}\n"; return L.str(); } @@ -2210,9 +2120,8 @@ inline std::string generate_glsl_fold_v2(const std::string& name, const ExprSpec // // native16 恒走**标量 kernel**（vec4 的 f16vec4 向量化路径留作后续优化）—— // 对应变体的 FusedShader.vec_width 必须为 1（fused_generate.hpp 同源判定）。 -inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, - ExprPrecSig sig = 0, bool native16 = false, - bool batched = false, bool mm_fine = false, +inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, ExprPrecSig sig = 0, + bool native16 = false, bool batched = false, bool mm_fine = false, bool mm_ksplit = false) { // P3 批量变体（键 = key#b）：纯逐元素 + 仅 V0 全 f32。其余形态返回空串 @@ -2225,8 +2134,7 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, // 故 f16 批量变体只对 **extras 为空**的结构发射（实测优化器在 f16 下的 // 4 个结构全是单输出 → 覆盖到）；带 extras 的 f16 结构仍逐 kernel 降级。 // 运行期分派哨兵不叠加批量变体。 - if (batched && (expr_prec_sig_is_dispatch(sig) || native16 || - !expr_spec_batchable(spec) || + if (batched && (expr_prec_sig_is_dispatch(sig) || native16 || !expr_spec_batchable(spec) || (sig != 0 && !spec.extras.empty()))) return {}; if (mm_fine) @@ -2236,8 +2144,8 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, // PC 布局与 V0 完全一致（block 不进 PC）→ 同一固定头、同一绑定数， // 所以 sig 只是把"输入/输出读写的元素类型"静态化，不引入新槽。 // 运行期分派哨兵不支持（它多一个 `prec` 槽且绑定翻倍，网格变体不叠加）。 - if (expr_prec_sig_is_dispatch(sig) || native16 || batched || spec.fold || - !spec.matmul || expr_spec_reduce_axis(spec) != -1) + if (expr_prec_sig_is_dispatch(sig) || native16 || batched || spec.fold || !spec.matmul || + expr_spec_reduce_axis(spec) != -1) return {}; return generate_glsl_matmul(name, spec, sig, EXPR_MATMUL_BLOCK_FINE); } @@ -2248,54 +2156,44 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, //（K 切分只在细 tile 网格仍饥饿时派发）；PC 比 V0 多一槽 ks_span //（见 generate_glsl_matmul 的 k_split 分支）。sig 同上：0 = V0， // 具体签名 = 静态特化变体（**部分和缓冲恒 f32**，与 sig 输出精度无关）。 - if (expr_prec_sig_is_dispatch(sig) || native16 || batched || spec.fold || - !spec.matmul || expr_spec_reduce_axis(spec) != -1) + if (expr_prec_sig_is_dispatch(sig) || native16 || batched || spec.fold || !spec.matmul || + expr_spec_reduce_axis(spec) != -1) return {}; - return generate_glsl_matmul(name, spec, sig, - EXPR_MATMUL_BLOCK_FINE, /*k_split=*/true); + return generate_glsl_matmul(name, spec, sig, EXPR_MATMUL_BLOCK_FINE, /*k_split=*/true); } - if (native16 && !expr_prec_sig_native16(spec, sig)) - return {}; + if (native16 && !expr_prec_sig_native16(spec, sig)) return {}; if (sig != 0) { // 本函数只处理非归约形态（带归约轴的 spec 由 gen_fused 按 raxis 路由到 // generate_glsl_reduce）；带归约的 spec 误入本路径时防御性返回空串 // （上层跳过该变体，运行时回退边界 cast，正确性不受影响）。 - if (!spec.fold && expr_spec_reduce_axis(spec) != -1) - return {}; + if (!spec.fold && expr_spec_reduce_axis(spec) != -1) return {}; } // fold 段：双域形态（向量态/mm 段）走 v2，纯标量走 v1 if (spec.fold) return (spec.fold->vec_state_len > 0 || spec.fold->matmul) - ? generate_glsl_fold_v2(name, spec, sig) - : generate_glsl_fold(name, spec, sig); + ? generate_glsl_fold_v2(name, spec, sig) + : generate_glsl_fold(name, spec, sig); // 含前置 matmul 段：matmul + 尾逐元素链融合 shader - if (spec.matmul) - return generate_glsl_matmul(name, spec, sig); + if (spec.matmul) return generate_glsl_matmul(name, spec, sig); const std::size_t n_inputs = spec.views.size(); const bool disp = expr_prec_sig_is_dispatch(sig); // P2 多输出（extras，IR-C 写穿物化）：仅 V0（全 f32）生成——`#x` 分派/ // typed/native16 变体跳过（返回空串 = 上层 [skip]，该结构运行时只走 V0； // 多输出节点同精度融合由 fuse_expr_graph 的精度门禁保证）。 - if (!spec.extras.empty() && (disp || sig != 0 || native16)) - return {}; + if (!spec.extras.empty() && (disp || sig != 0 || native16)) return {}; std::ostringstream L; const auto in_type = [&](std::size_t i) -> const char* - { - return (sig != 0 && expr_prec_sig_in_f16(sig, i)) ? "float16_t" : "float"; - }; + { return (sig != 0 && expr_prec_sig_in_f16(sig, i)) ? "float16_t" : "float"; }; const bool out_f16 = sig != 0 && expr_prec_sig_out_f16(sig); const bool any_f16 = sig != 0; L << "// ── 自动生成（AOT 算子融合），请勿手动编辑 ──\n"; L << "// 表达式: " << name; - if (sig != 0 && !disp) - L << " [精度 " << expr_prec_sig_str(sig, n_inputs) << "]"; - else if (disp) - L << " [运行期精度分派]"; - else if (batched) - L << " [批量 #b：实例表 + buffer device address]"; + if (sig != 0 && !disp) L << " [精度 " << expr_prec_sig_str(sig, n_inputs) << "]"; + else if (disp) L << " [运行期精度分派]"; + else if (batched) L << " [批量 #b：实例表 + buffer device address]"; L << "\n"; L << "#version 450\n\n"; if (batched) @@ -2304,14 +2202,12 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, // （glslc 强制，缺则 "required extension not requested"）。 L << "#extension GL_EXT_buffer_reference : require\n" "#extension GL_EXT_buffer_reference_uvec2 : require\n\n"; - if (any_f16) - L << "#extension GL_EXT_shader_16bit_storage : require\n\n"; + if (any_f16) L << "#extension GL_EXT_shader_16bit_storage : require\n\n"; if (native16) // f16 算术扩展（float16_t 类型的 +,-,*,/,exp,log,... 与 f16vecN）； // 与 16bit_storage 同时声明（spike 验证两者共存可编译）。 L << "#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require\n\n"; - if (any_f16) - L << glsl_f16_rne_source(); + if (any_f16) L << glsl_f16_rne_source(); L << "layout(local_size_x = 256) in;\n\n"; // 逐输入/输出的存储类型 @@ -2331,7 +2227,8 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, // 版逐字一致（RowGather 的 b 标签槽同样覆盖）。 // `ib` 必须是模块级变量：gr_r 辅助函数也经宏访问实例表。 L << "// P3 批量派发：实例表槽位 = 输入…、输出…、(count,cols)；" - "stride S = " << (n_inputs + 2 + spec.extras.size()) << "\n"; + "stride S = " + << (n_inputs + 2 + spec.extras.size()) << "\n"; // 实例表的地址是类型无关的；元素类型由**注入签名**在编译期定死 //（f16 槽走 float16_t 的 buffer reference）。extras 只在 sig==0 //（V0）时存在 → 其引用类型恒 F32BufW（见上方资格拒绝）。 @@ -2352,46 +2249,42 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, L << "uint ib; // 模块级 per-invocation 私有（gr_r<> 辅助函数也用）\n"; for (std::size_t i = 0; i < n_inputs; ++i) L << "#define b" << i << " (" - << ((b_sign_f16 && expr_prec_sig_in_f16(sig, i)) ? "F16BufR" - : "F32BufR") + << ((b_sign_f16 && expr_prec_sig_in_f16(sig, i)) ? "F16BufR" : "F32BufR") << "(inst[ib + " << i << "u]).d)\n"; - L << "#define bout (" << (b_out_f16 ? "F16BufW" : "F32BufW") - << "(inst[ib + " << n_inputs << "u]).d)\n"; + L << "#define bout (" << (b_out_f16 ? "F16BufW" : "F32BufW") << "(inst[ib + " << n_inputs + << "u]).d)\n"; for (std::size_t k = 0; k < spec.extras.size(); ++k) - L << "#define bex" << k << " (F32BufW(inst[ib + " - << (n_inputs + 1 + k) << "u]).d)\n"; + L << "#define bex" << k << " (F32BufW(inst[ib + " << (n_inputs + 1 + k) << "u]).d)\n"; L << "\n"; } else { - // 非批量：per-binding 声明（发射文本与历史逐字一致） - for (std::size_t i = 0; i < n_inputs; ++i) - { - L << "layout(std430, binding = " << i << ") readonly buffer Buf" << i - << " { " << (disp ? "float" : in_type(i)) << " b" << i << "[]; };\n"; + // 非批量：per-binding 声明（发射文本与历史逐字一致） + for (std::size_t i = 0; i < n_inputs; ++i) + { + L << "layout(std430, binding = " << i << ") readonly buffer Buf" << i << " { " + << (disp ? "float" : in_type(i)) << " b" << i << "[]; };\n"; + if (disp) + L << "layout(std430, binding = " << (n_inputs + i) << ") readonly buffer Buf" << i + << "h { float16_t b" << i << "h[]; };\n"; + } if (disp) - L << "layout(std430, binding = " << (n_inputs + i) - << ") readonly buffer Buf" << i << "h { float16_t b" << i << "h[]; };\n"; - } - if (disp) - { - L << "layout(std430, binding = " << (2 * n_inputs) - << ") writeonly buffer BufOut { float bout[]; };\n"; - L << "layout(std430, binding = " << (2 * n_inputs + 1) - << ") writeonly buffer BufOuth { float16_t bout16[]; };\n\n"; - } - else - { - L << "layout(std430, binding = " << n_inputs - << ") writeonly buffer BufOut { " << (out_f16 ? "float16_t" : "float") - << " bout[]; };\n"; - // P2 多输出：额外输出缓冲（V0 专属，类型恒 float——sig==0 ⇒ out_f16=0） - for (std::size_t k = 0; k < spec.extras.size(); ++k) - L << "layout(std430, binding = " << (n_inputs + 1 + k) - << ") writeonly buffer BufOut" << (k + 1) - << " { float bex" << k << "[]; };\n"; - L << "\n"; - } + { + L << "layout(std430, binding = " << (2 * n_inputs) + << ") writeonly buffer BufOut { float bout[]; };\n"; + L << "layout(std430, binding = " << (2 * n_inputs + 1) + << ") writeonly buffer BufOuth { float16_t bout16[]; };\n\n"; + } + else + { + L << "layout(std430, binding = " << n_inputs << ") writeonly buffer BufOut { " + << (out_f16 ? "float16_t" : "float") << " bout[]; };\n"; + // P2 多输出：额外输出缓冲（V0 专属，类型恒 float——sig==0 ⇒ out_f16=0） + for (std::size_t k = 0; k < spec.extras.size(); ++k) + L << "layout(std430, binding = " << (n_inputs + 1 + k) + << ") writeonly buffer BufOut" << (k + 1) << " { float bex" << k << "[]; };\n"; + L << "\n"; + } } // push constants：count + cols（视图行/列）+ [prec] + 运行时视图参数 vp + 常量 @@ -2417,8 +2310,7 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, // （批量组键与 rows/cols 无关），故主体与 gr_r 辅助函数里的 `cols` // 一律取本实例值——宏放在 PC 块之后（PC 布局与 V0 逐字一致，批量路径 // 不读其 count/cols 两槽）、辅助函数之前（它们也用 `cols`）。 - L << "#define cols ((inst[ib + " - << (n_inputs + 1 + spec.extras.size()) << "u]).y)\n\n"; + L << "#define cols ((inst[ib + " << (n_inputs + 1 + spec.extras.size()) << "u]).y)\n\n"; } // 分组归约视图（MaxPool 等）的读取辅助函数：R 是运行期 vp 参数 → 循环。 @@ -2432,7 +2324,8 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, // 批量变体同理恒标量（实例表按元素下标派发，无 vec4 快路径）。 const bool vec4_ok = glsl_vec4_eligible(spec) && !native16 && !batched; // 索引操作数（Row/Col/Batch）也按需发射 row/col/batch 变量 - const auto instr_uses = [&](uint8_t kind) { + const auto instr_uses = [&](uint8_t kind) + { for (const auto& ins : spec.instrs) { const ExprOperand* ops[3] = {&ins.a, &ins.b, &ins.c}; @@ -2442,14 +2335,18 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, } return false; }; - const bool need_row = [&]{ + const bool need_row = [&] + { for (const auto& v : spec.views) if (glsl_view_uses_row(static_cast(v.kind))) return true; - return instr_uses(static_cast(ExprOperandKind::Row)); }(); - const bool need_col = [&]{ + return instr_uses(static_cast(ExprOperandKind::Row)); + }(); + const bool need_col = [&] + { for (const auto& v : spec.views) if (glsl_view_uses_col(static_cast(v.kind))) return true; - return instr_uses(static_cast(ExprOperandKind::Col)); }(); + return instr_uses(static_cast(ExprOperandKind::Col)); + }(); const bool need_batch = glsl_spec_uses_batch(spec); const std::uint32_t last_dst = spec.instrs.back().dst; @@ -2462,7 +2359,7 @@ 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): // Fanout 语义同 Reg return "r" + std::to_string(op.idx); case static_cast(ExprOperandKind::Const): // native16：常量按 float 推入 push constant，用点显式收 f16—— @@ -2492,10 +2389,8 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, // 发射一段标量指令链（输入读取 + 寄存器声明 + 指令展开，直线代码； // 寄存器先声明、后赋值——兼容 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 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 槽见 glsl_vp_slot @@ -2504,19 +2399,19 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, const std::uint32_t vp = glsl_vp_slot(spec, i); // native16：读取直接产出 float16_t（谓词保证全部输入 f16，故 // 无 float() 包装——全链 f16，无逐元素往返）；f32 算术变体保持原样。 - o << indent << "const " << (native16 ? "float16_t" : "float") - << " v" << i << " = "; + o << indent << "const " << (native16 ? "float16_t" : "float") << " v" << i << " = "; const bool f16_in = (sig != 0 && expr_prec_sig_in_f16(sig, i)); - if (f16_in && !native16) o << "float("; // f16 读 → 显式转 f32（算术仍在 f32） - glsl_view_read(o, spec.views[i], static_cast(i), - idx_var, row_var, col_var, vp, sig, native16); + if (f16_in && !native16) o << "float("; // f16 读 → 显式转 f32（算术仍在 f32） + glsl_view_read(o, spec.views[i], static_cast(i), idx_var, row_var, + col_var, vp, sig, native16); if (f16_in && !native16) o << ")"; o << ";\n"; } if (spec.num_regs > 0) { o << indent << (native16 ? "float16_t r0" : "float r0"); - for (std::uint32_t r = 1; r < spec.num_regs; ++r) o << ", r" << r; + for (std::uint32_t r = 1; r < spec.num_regs; ++r) + o << ", r" << r; o << ";\n"; } for (const auto& ins : spec.instrs) @@ -2525,14 +2420,13 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, const std::string dst = "r" + std::to_string(ins.dst); const std::string a = operand(ins.a); const detail::ExprOpInfo& info = detail::expr_op_info(op); - if (info.emit == detail::ExprEmitKind::None) - return; // 未知算子：让上层报错 + if (info.emit == detail::ExprEmitKind::None) return; // 未知算子：让上层报错 // 发射形态/符号查表；0/1 字面量按目标寄存器类型发射（f16 寄存器 // 不能隐式收窄 float）——由 native16 控制 o << indent << dst << " = " - << glsl_alu_expr(info, a, [&] { return operand(ins.b); }, - [&] { return operand(ins.c); }, - /*vec4=*/false, native16, /*wrap=*/false) + << glsl_alu_expr( + info, a, [&] { return operand(ins.b); }, [&] { return operand(ins.c); }, + /*vec4=*/false, native16, /*wrap=*/false) << ";\n"; } }; @@ -2560,22 +2454,16 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, } 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"; - if (need_batch) - L << " const uint batch = 0u;\n"; - emit_scalar_chain(L, " ", "i", - need_row ? "row" : "", need_col ? "col" : ""); + else if (need_row) L << " const uint row = i / cols;\n"; + else if (need_col) L << " const uint col = i % cols;\n"; + if (need_batch) L << " const uint batch = 0u;\n"; + emit_scalar_chain(L, " ", "i", need_row ? "row" : "", need_col ? "col" : ""); // native16：寄存器本就是 float16_t → 直写（免冗余 float16_t() 转换） - if (native16) - L << " bout[i] = r" << last_dst << ";\n"; + if (native16) L << " bout[i] = r" << last_dst << ";\n"; else if (disp) { L << " if ((prec >> 16u & 1u) != 0u) {\n"; - L << " bout16[i] = " << glsl_f16_store("r" + std::to_string(last_dst)) - << ";\n"; + L << " bout16[i] = " << glsl_f16_store("r" + std::to_string(last_dst)) << ";\n"; L << " } else {\n"; L << " bout[i] = r" << last_dst << ";\n"; L << " }\n"; @@ -2587,8 +2475,7 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, << ";\n"; // P2 多输出：额外输出寄存器同元素写各自的输出缓冲（V0 专属） for (std::size_t k = 0; k < spec.extras.size(); ++k) - L << " bex" << k << "[i] = r" - << static_cast(spec.extras[k]) << ";\n"; + L << " bex" << k << "[i] = r" << static_cast(spec.extras[k]) << ";\n"; L << "}\n"; return L.str(); } @@ -2603,10 +2490,8 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, 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"; + 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） @@ -2614,23 +2499,22 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, { const auto k = static_cast(spec.views[i].kind); const std::string cvt = (sig != 0 && expr_prec_sig_in_f16(sig, i)) ? "float" : ""; - const auto ld = [&](const std::string& idx) { + const auto ld = [&](const std::string& idx) + { const std::string bf = "b" + std::to_string(i) + "[" + idx + "]"; if (disp) - return "((prec >> " + std::to_string(i) + "u & 1u) != 0u ? float(b" - + std::to_string(i) + "h[" + idx + "]) : " + bf + ")"; + return "((prec >> " + std::to_string(i) + "u & 1u) != 0u ? float(b" + + std::to_string(i) + "h[" + idx + "]) : " + bf + ")"; return cvt.empty() ? bf : ("float(" + bf + ")"); }; if (k == ExprViewKind::Linear) - L << " const vec4 v" << i << " = vec4(" << ld("base") << ", " - << ld("base+1u") << ", " << ld("base+2u") << ", " << ld("base+3u") - << ");\n"; + L << " const vec4 v" << i << " = vec4(" << ld("base") << ", " << ld("base+1u") + << ", " << ld("base+2u") << ", " << ld("base+3u") << ");\n"; else if (k == ExprViewKind::RowBroadcast) L << " const vec4 v" << i << " = vec4(" << ld("row") << ");\n"; else // ColBroadcast - L << " const vec4 v" << i << " = vec4(" << ld("col0") << ", " - << ld("col0+1u") << ", " << ld("col0+2u") << ", " << ld("col0+3u") - << ");\n"; + L << " const vec4 v" << i << " = vec4(" << ld("col0") << ", " << ld("col0+1u") + << ", " << ld("col0+2u") << ", " << ld("col0+3u") << ");\n"; } // vec4 指令展开（GLSL 对 vec4 重载运算/内置函数；常量广播为 vec4） @@ -2653,7 +2537,8 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, if (spec.num_regs > 0) { L << " vec4 r0"; - for (std::uint32_t r = 1; r < spec.num_regs; ++r) L << ", r" << r; + for (std::uint32_t r = 1; r < spec.num_regs; ++r) + L << ", r" << r; L << ";\n"; } for (const auto& ins : spec.instrs) @@ -2662,22 +2547,23 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, const std::string dst = "r" + std::to_string(ins.dst); const std::string a = vec4_operand(ins.a); const detail::ExprOpInfo& info = detail::expr_op_info(op); - if (info.emit == detail::ExprEmitKind::None) - return std::string{}; // 未知算子：让上层报错 + if (info.emit == detail::ExprEmitKind::None) return std::string{}; // 未知算子：让上层报错 // 发射形态/符号查表；向量形态：比较用向量内建、Select 用 mix/notEqual L << " " << dst << " = " - << glsl_alu_expr(info, a, [&] { return vec4_operand(ins.b); }, - [&] { return vec4_operand(ins.c); }, - /*vec4=*/true, /*native16=*/false, /*wrap=*/false) + << glsl_alu_expr( + info, a, [&] { return vec4_operand(ins.b); }, [&] { return vec4_operand(ins.c); }, + /*vec4=*/true, /*native16=*/false, /*wrap=*/false) << ";\n"; } // 运行期精度分派：输出双视图 + uniform 存储分支（每元素只写一次， // 两侧视图按各自的元素宽度寻址，无别名风险）。 // 输出类型转换由本 lambda 统一负责（c16）——此前另有一个 st() 也包 // float16_t(...)，在"非分派 + f16 输出"上生成 float16_t(float16_t(x))。 - const auto emit_store4 = [&](const char* arr, bool c16, const char* ind) { + const auto emit_store4 = [&](const char* arr, bool c16, const char* ind) + { const std::string r = "r" + std::to_string(last_dst); - const auto val = [&](char comp) { + const auto val = [&](char comp) + { const std::string v = r + "." + comp; return c16 ? glsl_f16_store(v) : v; }; @@ -2694,8 +2580,7 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, emit_store4("bout", false, " "); L << " }\n"; } - else - emit_store4("bout", out_f16, " "); + else emit_store4("bout", out_f16, " "); // P2 多输出：vec4 快路径按分量写额外输出缓冲 for (std::size_t k = 0; k < spec.extras.size(); ++k) { @@ -2712,14 +2597,10 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, 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"; - if (need_batch) - L << " const uint batch = 0u;\n"; - emit_scalar_chain(L, " ", "e", - need_row ? "row" : "", need_col ? "col" : ""); + else if (need_row) L << " const uint row = e / cols;\n"; + else if (need_col) L << " const uint col = e % cols;\n"; + if (need_batch) L << " const uint batch = 0u;\n"; + emit_scalar_chain(L, " ", "e", need_row ? "row" : "", need_col ? "col" : ""); const std::string sv = "r" + std::to_string(last_dst); if (disp) { @@ -2729,12 +2610,10 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, L << " bout[e] = " << sv << ";\n"; L << " }\n"; } - else - L << " bout[e] = " << (out_f16 ? glsl_f16_store(sv) : sv) << ";\n"; + else L << " bout[e] = " << (out_f16 ? glsl_f16_store(sv) : sv) << ";\n"; // P2 多输出：标量回退路径同元素写额外输出缓冲（此处 rN 为标量） for (std::size_t k = 0; k < spec.extras.size(); ++k) - L << " bex" << k << "[e] = r" - << static_cast(spec.extras[k]) << ";\n"; + L << " bex" << k << "[e] = r" << static_cast(spec.extras[k]) << ";\n"; L << " }\n"; L << "}\n"; return L.str(); @@ -2754,12 +2633,11 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, // 最终值读取 s_red[slot][0]（完成屏障后全线程可见）。 // - 限制：归约槽 ≤ 8（共享内存 ≤ 8KB）；混合轴（行+列）不支持 → 返回 ""。 // ═══════════════════════════════════════════════════════════════════════════ -[[nodiscard]] inline std::string generate_glsl_reduce( - const std::string& name, const ExprSpec& spec, ExprPrecSig sig = 0) +[[nodiscard]] inline std::string generate_glsl_reduce(const std::string& name, const ExprSpec& spec, + ExprPrecSig sig = 0) { const int axis = expr_spec_reduce_axis(spec); - if (axis < 0) - return {}; // 无归约（走逐元素生成器）或混合轴（不支持） + if (axis < 0) return {}; // 无归约（走逐元素生成器）或混合轴（不支持） const bool is_row = (axis == 0); const std::size_t n_inputs = spec.views.size(); @@ -2773,26 +2651,21 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, 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; + if (!expr_view_is_reduce(static_cast(v.kind))) continue; slot_of_view[k] = n_slots++; - slot_is_max.push_back( - static_cast(v.kind) == ExprViewKind::RowReduceMax || - static_cast(v.kind) == ExprViewKind::ColReduceMax); + slot_is_max.push_back(static_cast(v.kind) == ExprViewKind::RowReduceMax || + static_cast(v.kind) == ExprViewKind::ColReduceMax); } for (const auto& ins : spec.instrs) { const ExprOp op = static_cast(ins.op); - if (!expr_op_is_reduce(op)) - continue; + if (!expr_op_is_reduce(op)) continue; slot_of_instr[ins.dst] = n_slots++; slot_is_max.push_back(op == ExprOp::RowMax || op == ExprOp::ColMax); } - if (n_slots == 0 || n_slots > 8) - return {}; // 无归约或槽过多（共享内存超限） + if (n_slots == 0 || n_slots > 8) return {}; // 无归约或槽过多（共享内存超限） - const bool last_is_reduce = - expr_op_is_reduce(static_cast(spec.instrs.back().op)); + const bool last_is_reduce = expr_op_is_reduce(static_cast(spec.instrs.back().op)); // ── 带类型变体（Phase 2 in-kernel f16）─────────────────────────────── // sig==0 → 全 f32：生成的就是 f32 变体源码。sig!=0 → 对应槽的 @@ -2801,44 +2674,39 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, const bool sign_f16 = (sig != 0); const bool disp = expr_prec_sig_is_dispatch(sig); const bool out_f16 = sign_f16 && expr_prec_sig_out_f16(sig); - const auto rd = [&](std::size_t k, const std::string& idx) -> std::string { + const auto rd = [&](std::size_t k, const std::string& idx) -> std::string + { const std::string raw = "b" + std::to_string(k) + "[" + idx + "]"; if (disp) - return "((prec >> " + std::to_string(k) + "u & 1u) != 0u ? float(b" - + std::to_string(k) + "h[" + idx + "]) : " + raw + ")"; + return "((prec >> " + std::to_string(k) + "u & 1u) != 0u ? float(b" + + std::to_string(k) + "h[" + idx + "]) : " + raw + ")"; return (sign_f16 && expr_prec_sig_in_f16(sig, k)) ? ("float(" + raw + ")") : raw; }; - const auto wr = [&](const std::string& v) -> std::string { - return out_f16 ? glsl_f16_store(v) : v; - }; + const auto wr = [&](const std::string& v) -> std::string + { return out_f16 ? glsl_f16_store(v) : v; }; std::ostringstream L; L << "// ── 自动生成（AOT 算子融合 · 归约 kernel），请勿手动编辑 ──\n"; L << "// 表达式: " << name; - if (sign_f16 && !disp) - L << " [精度 " << expr_prec_sig_str(sig, n_inputs) << "]"; - else if (disp) - L << " [运行期精度分派]"; + if (sign_f16 && !disp) L << " [精度 " << expr_prec_sig_str(sig, n_inputs) << "]"; + else if (disp) L << " [运行期精度分派]"; L << "\n"; L << "#version 450\n"; L << "#extension GL_KHR_shader_subgroup_basic : enable\n"; L << "#extension GL_KHR_shader_subgroup_arithmetic : enable\n"; - if (sign_f16) - L << "#extension GL_EXT_shader_16bit_storage : require\n"; + if (sign_f16) L << "#extension GL_EXT_shader_16bit_storage : require\n"; L << "\n"; - if (sign_f16) - L << glsl_f16_rne_source(); + if (sign_f16) L << glsl_f16_rne_source(); 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 - << " { " - << ((disp || (sign_f16 && expr_prec_sig_in_f16(sig, i))) - ? (disp ? "float" : "float16_t") : "float") + L << "layout(std430, binding = " << i << ") readonly buffer Buf" << i << " { " + << ((disp || (sign_f16 && expr_prec_sig_in_f16(sig, i))) ? (disp ? "float" : "float16_t") + : "float") << " b" << i << "[]; };\n"; if (disp) - L << "layout(std430, binding = " << (n_inputs + i) - << ") readonly buffer Buf" << i << "h { float16_t b" << i << "h[]; };\n"; + L << "layout(std430, binding = " << (n_inputs + i) << ") readonly buffer Buf" << i + << "h { float16_t b" << i << "h[]; };\n"; } if (disp) { @@ -2848,9 +2716,8 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, << ") writeonly buffer BufOuth { float16_t bout16[]; };\n\n"; } else - L << "layout(std430, binding = " << n_inputs - << ") writeonly buffer BufOut { " << (out_f16 ? "float16_t" : "float") - << " bout[]; };\n\n"; + L << "layout(std430, binding = " << n_inputs << ") writeonly buffer BufOut { " + << (out_f16 ? "float16_t" : "float") << " bout[]; };\n\n"; L << "layout(push_constant) uniform PC {\n"; L << " uint count;\n"; L << " uint cols;\n"; @@ -2865,8 +2732,7 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, L << " uint mm_batch;\n"; L << " uint mm_trans; // bit0=transA, bit1=transB（不进 key）\n"; } - if (disp) - L << " uint prec; // 运行期精度分派位（bit i=输入 i 为 f16，bit16=输出）\n"; + if (disp) L << " uint prec; // 运行期精度分派位（bit i=输入 i 为 f16，bit16=输出）\n"; const std::uint32_t n_vp = expr_spec_runtime_view_param_count(spec); for (std::uint32_t i = 0; i < n_vp; ++i) L << " uint vp" << i << ";\n"; @@ -2884,8 +2750,7 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, // 输出写回：非分派 = 单条赋值（wr 包 float16_t）；分派 = 输出双视图 // + uniform 分支（每个元素只写一次）。 - const auto emit_wr = [&](const char* ind, const std::string& lhs, - const std::string& v) + const auto emit_wr = [&](const char* ind, const std::string& lhs, const std::string& v) { if (!disp) { @@ -2905,15 +2770,13 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, const MatmulSpec* mm = spec.matmul ? &*spec.matmul : nullptr; const auto emit_mm_decl = [&]() { - if (!mm) - return; + if (!mm) return; const std::uint32_t a_slot = mm->a_input; const std::uint32_t b_slot = mm->b_input; L << " const uint row_in_batch = row % m_per;\n"; // 列归约：batch 无顶层分解（idx 是列 tile）→ 按当前 row 就地分解； // 行归约：顶层已有 batch（= idx/m_per），此处不重复声明 - if (!is_row) - L << " const uint batch = row / m_per;\n"; + if (!is_row) L << " const uint batch = row / m_per;\n"; // 转置布局（运行期 mm_trans，不进 key）：两种索引折算成 base + kk*stride， // 循环体内**无分支**（mm_trans 是 uniform push constant）。 // transA : A 存 (K,M) → idx = (batch*mm_k + kk)*m_per + row_in_batch @@ -2930,8 +2793,8 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, L << " const uint b_stride = tB ? 1u : cols;\n"; L << " float mm = 0.0;\n"; L << " for (uint kk = 0u; kk < mm_k; ++kk)\n"; - L << " mm += " << rd(a_slot, "a_base + kk*a_stride") - << " * " << rd(b_slot, "b_base + kk*b_stride") << ";\n"; + L << " mm += " << rd(a_slot, "a_base + kk*a_stride") << " * " + << rd(b_slot, "b_base + kk*b_stride") << ";\n"; }; L << "void main()\n{\n"; @@ -2949,8 +2812,7 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, if (mm) { L << " const uint m_per = rows / mm_batch;\n"; - if (is_row) - L << " const uint batch = idx / m_per;\n"; + if (is_row) L << " const uint batch = idx / m_per;\n"; } // 归约结果索引：行=工作组单行（s_red[slot][0]）；列 tile=每线程一列（s_red[slot][tid]） @@ -2972,7 +2834,7 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, case static_cast(ExprOperandKind::Reduce): return "s_red[" + std::to_string(slot_of_instr[op.idx]) + "][" + red_idx + "]"; case static_cast(ExprOperandKind::Matmul): - return "mm"; // 前置 matmul 段输出（emit_mm_decl 声明） + return "mm"; // 前置 matmul 段输出（emit_mm_decl 声明） // 索引操作数：Row 为 batch 内行号（row 变量是全局行 idx）； // batch 由 idx 分解（仅 matmul 段存在时） case static_cast(ExprOperandKind::Row): @@ -2980,8 +2842,7 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, case static_cast(ExprOperandKind::Col): return "float(col)"; case static_cast(ExprOperandKind::Batch): - if (!mm) - return "0.0"; + if (!mm) return "0.0"; // 行归约读顶层 batch（idx 分解）；列归约按当前 row 就地分解 return is_row ? "float(batch)" : "float(row / m_per)"; case static_cast(ExprOperandKind::Input): @@ -2990,20 +2851,16 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, const ExprViewKind vk = static_cast(v.kind); if (expr_view_is_reduce(vk)) return "s_red[" + std::to_string(slot_of_view[op.idx]) + "][" + red_idx + "]"; - if (vk == ExprViewKind::RowBroadcast) - return rd(op.idx, "row"); - if (vk == ExprViewKind::ColBroadcast) - return rd(op.idx, "col"); + if (vk == ExprViewKind::RowBroadcast) return rd(op.idx, "row"); + if (vk == ExprViewKind::ColBroadcast) return rd(op.idx, "col"); // Linear / RotateHalf / RowMod / RowGather / BatchMod：索引映射内联； // vp 槽见 glsl_vp_slot std::ostringstream os; - glsl_view_read(os, v, static_cast(op.idx), - "(row*cols + col)", "row", "col", - glsl_vp_slot(spec, op.idx), sig); + glsl_view_read(os, v, static_cast(op.idx), "(row*cols + col)", "row", + "col", glsl_vp_slot(spec, op.idx), sig); const std::string raw = os.str(); // f16 槽：视图读取返回 float16_t 左值 → 显式转 f32（算术仍在 f32） - return (sign_f16 && expr_prec_sig_in_f16(sig, op.idx)) - ? ("float(" + raw + ")") : raw; + return (sign_f16 && expr_prec_sig_in_f16(sig, op.idx)) ? ("float(" + raw + ")") : raw; } } }; @@ -3012,10 +2869,10 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, // 每个 emit_instrs 调用点所在作用域内须先调用一次。 const auto emit_reg_decl = [&]() { - if (spec.num_regs == 0) - return; + if (spec.num_regs == 0) return; L << " float r0"; - for (std::uint32_t r = 1; r < spec.num_regs; ++r) L << ", r" << r; + for (std::uint32_t r = 1; r < spec.num_regs; ++r) + L << ", r" << r; L << ";\n"; }; @@ -3028,27 +2885,23 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, { const ExprInstr& ins = spec.instrs[j]; const ExprOp op = static_cast(ins.op); - if (expr_op_is_reduce(op)) - continue; + if (expr_op_is_reduce(op)) continue; const std::string dst = "r" + std::to_string(ins.dst); const std::string a = operand(ins.a); const detail::ExprOpInfo& info = detail::expr_op_info(op); - if (info.emit == detail::ExprEmitKind::None) - continue; // 归约指令已在上方 continue 分流 + if (info.emit == detail::ExprEmitKind::None) continue; // 归约指令已在上方 continue 分流 // 发射形态/符号查表（expr_op_table） L << " " << dst << " = " - << glsl_alu_expr(info, a, [&] { return operand(ins.b); }, - [&] { return operand(ins.c); }, - /*vec4=*/false, /*native16=*/false, /*wrap=*/false) + << glsl_alu_expr( + info, a, [&] { return operand(ins.b); }, [&] { return operand(ins.c); }, + /*vec4=*/false, /*native16=*/false, /*wrap=*/false) << ";\n"; } }; // 组合表达式：sum → "a + b"（中缀），max → "max(a, b)"（函数） const auto combine = [](bool is_max, const std::string& a, const std::string& b) - { - return is_max ? "max(" + a + ", " + b + ")" : "(" + a + " + " + b + ")"; - }; + { return is_max ? "max(" + a + ", " + b + ")" : "(" + a + " + " + b + ")"; }; // 归约：warp shuffle 蝴蝶归约（subgroup）替代共享内存树形归约。 // 第 1 步：warp 内 subgroupAdd/subgroupMax → 每 warp 一个部分和（零共享/屏障） @@ -3060,7 +2913,7 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, const char* sub = is_max ? "subgroupMax" : "subgroupAdd"; // max init = lowest()（0xFF7FFFFF）对齐 CPU 基准，理由同归约视图处注释 const std::string idt = is_max ? "uintBitsToFloat(0xFF7FFFFFu)" : "0.0"; - const std::string vn = "v" + std::to_string(slot); // 按槽唯一命名，避免重定义 + 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"; @@ -3068,8 +2921,8 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, L << " if (gl_SubgroupInvocationID == 0u) " << s << "[gl_SubgroupID] = " << vn << ";\n"; L << " barrier();\n"; L << " if (gl_SubgroupID == 0u) {\n"; - L << " float " << wn << " = (gl_SubgroupInvocationID < gl_NumSubgroups) ? " - << s << "[gl_SubgroupInvocationID] : " << idt << ";\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"; @@ -3078,23 +2931,21 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, // 组合 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); - }; + const std::string& a2, const std::string& a3) + { return combine(is_max, combine(is_max, combine(is_max, a0, a1), a2), a3); }; if (is_row) { // ═══ 行归约（工作组单行、跨线程跨步 + 树形归约；合并访问）═══ - 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 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; + 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)]; // max init = lowest()（0xFF7FFFFF）对齐 CPU 基准（全库归约一律 @@ -3104,19 +2955,22 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, const std::string init = is_max ? "uintBitsToFloat(0xFF7FFFFFu)" : "0.0"; L << "\n // 归约视图 " << k << " (槽 " << slot << ")\n"; L << " {\n"; - L << " float acc0 = " << init << ", acc1 = " << init - << ", acc2 = " << init << ", acc3 = " << init << ";\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", rd(k, "idx*cols + i")) << ";\n"; - L << " acc1 = " << combine(is_max, "acc1", rd(k, "idx*cols + i+256u")) << ";\n"; - L << " acc2 = " << combine(is_max, "acc2", rd(k, "idx*cols + i+512u")) << ";\n"; - L << " acc3 = " << combine(is_max, "acc3", rd(k, "idx*cols + i+768u")) << ";\n"; + L << " acc1 = " << combine(is_max, "acc1", rd(k, "idx*cols + i+256u")) + << ";\n"; + L << " acc2 = " << combine(is_max, "acc2", rd(k, "idx*cols + i+512u")) + << ";\n"; + L << " acc3 = " << combine(is_max, "acc3", rd(k, "idx*cols + i+768u")) + << ";\n"; L << " }\n"; L << " for (; i < cols; i += 256u)\n"; L << " acc0 = " << combine(is_max, "acc0", rd(k, "idx*cols + i")) << ";\n"; - L << " s_red[" << slot << "][tid] = " - << combine4(is_max, "acc0", "acc1", "acc2", "acc3") << ";\n"; + L << " s_red[" << slot + << "][tid] = " << combine4(is_max, "acc0", "acc1", "acc2", "acc3") << ";\n"; L << " }\n"; emit_tree_reduce(slot, is_max); } @@ -3128,15 +2982,13 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, 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; + 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); + 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); // max init = lowest()（0xFF7FFFFF）对齐 CPU 基准（全库归约一律 // numeric_limits::lowest()）：-inf 作恒等元在全 -inf 输入下与 // 状态进位/减法组合会出 −inf−−inf=NaN，全库统一 lowest 根除 @@ -3144,8 +2996,8 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, const std::string init = is_max ? "uintBitsToFloat(0xFF7FFFFFu)" : "0.0"; L << "\n // 归约指令 " << ri << " (槽 " << slot << ")\n"; L << " {\n"; - L << " float acc0 = " << init << ", acc1 = " << init - << ", acc2 = " << init << ", acc3 = " << init << ";\n"; + L << " float acc0 = " << init << ", acc1 = " << init << ", acc2 = " << init + << ", acc3 = " << init << ";\n"; L << " uint c = tid;\n"; L << " for (; c + 3u*256u < cols; c += 4u*256u) {\n"; for (int k = 0; k < 4; ++k) @@ -3164,8 +3016,8 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, emit_instrs(0, ri); L << " acc0 = " << combine(is_max, "acc0", src) << "; }\n"; L << " }\n"; - L << " s_red[" << slot << "][tid] = " - << combine4(is_max, "acc0", "acc1", "acc2", "acc3") << ";\n"; + L << " s_red[" << slot + << "][tid] = " << combine4(is_max, "acc0", "acc1", "acc2", "acc3") << ";\n"; L << " }\n"; emit_tree_reduce(slot, is_max); } @@ -3222,17 +3074,16 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, const std::string s = "s_red[" + std::to_string(slot) + "]"; const std::string init = is_max ? "uintBitsToFloat(0xFF7FFFFFu)" : "0.0"; const std::string m = "m" + std::to_string(slot); - L << " barrier();\n"; // ① 各行块部分和写完 + L << " barrier();\n"; // ① 各行块部分和写完 L << " { float " << m << " = " << init << ";\n"; L << " for (uint w = 0u; w < 8u; ++w)\n"; - L << " " << m << " = " - << combine(is_max, m, s + "[w*32u + l]") << ";\n"; + L << " " << m << " = " << combine(is_max, m, s + "[w*32u + l]") << ";\n"; // ② 读/写分离：所有 lane 读完 8 项后才允许覆写自身槽位 // （写 s_red[tid=(wb*32+l)] 与他 lane 读 [w*32+l] 同址， // 无此 barrier 会读到已覆写值 → sum 双计） L << " barrier();\n"; L << " " << s << "[tid] = " << m << "; }\n"; - L << " barrier();\n"; // ③ 覆写完成，后续 pass/输出才可读 + L << " barrier();\n"; // ③ 覆写完成，后续 pass/输出才可读 }; L << " const uint l = tid & 31u;\n"; L << " const uint wb = tid >> 5u;\n"; @@ -3242,8 +3093,7 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, 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; + 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)]; // max init = lowest()（0xFF7FFFFF）对齐 CPU 基准（全库归约一律 @@ -3253,21 +3103,26 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, const std::string init = is_max ? "uintBitsToFloat(0xFF7FFFFFu)" : "0.0"; L << "\n // 归约视图 " << k << " (槽 " << slot << ")\n"; L << " {\n"; - L << " float acc0 = " << init << ", acc1 = " << init - << ", acc2 = " << init << ", acc3 = " << init << ";\n"; + L << " float acc0 = " << init << ", acc1 = " << init << ", acc2 = " << init + << ", acc3 = " << init << ";\n"; L << " if (col < cols) {\n"; L << " uint i = wb;\n"; L << " for (; i + 24u < rows; i += 32u) {\n"; - L << " acc0 = " << combine(is_max, "acc0", rd(k, "i*cols + col")) << ";\n"; - L << " acc1 = " << combine(is_max, "acc1", rd(k, "(i+8u)*cols + col")) << ";\n"; - L << " acc2 = " << combine(is_max, "acc2", rd(k, "(i+16u)*cols + col")) << ";\n"; - L << " acc3 = " << combine(is_max, "acc3", rd(k, "(i+24u)*cols + col")) << ";\n"; + L << " acc0 = " << combine(is_max, "acc0", rd(k, "i*cols + col")) + << ";\n"; + L << " acc1 = " << combine(is_max, "acc1", rd(k, "(i+8u)*cols + col")) + << ";\n"; + L << " acc2 = " << combine(is_max, "acc2", rd(k, "(i+16u)*cols + col")) + << ";\n"; + L << " acc3 = " << combine(is_max, "acc3", rd(k, "(i+24u)*cols + col")) + << ";\n"; L << " }\n"; L << " for (; i < rows; i += 8u)\n"; - L << " acc0 = " << combine(is_max, "acc0", rd(k, "i*cols + col")) << ";\n"; + L << " acc0 = " << combine(is_max, "acc0", rd(k, "i*cols + col")) + << ";\n"; L << " }\n"; - L << " s_red[" << slot << "][tid] = " - << combine4(is_max, "acc0", "acc1", "acc2", "acc3") << ";\n"; + L << " s_red[" << slot + << "][tid] = " << combine4(is_max, "acc0", "acc1", "acc2", "acc3") << ";\n"; L << " }\n"; emit_col_merge(slot, is_max); } @@ -3277,15 +3132,13 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, 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; + 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); + 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); // max init = lowest()（0xFF7FFFFF）对齐 CPU 基准（全库归约一律 // numeric_limits::lowest()）：-inf 作恒等元在全 -inf 输入下与 // 状态进位/减法组合会出 −inf−−inf=NaN，全库统一 lowest 根除 @@ -3293,15 +3146,14 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, const std::string init = is_max ? "uintBitsToFloat(0xFF7FFFFFu)" : "0.0"; L << "\n // 归约指令 " << ri << " (槽 " << slot << ")\n"; L << " {\n"; - L << " float acc0 = " << init << ", acc1 = " << init - << ", acc2 = " << init << ", acc3 = " << init << ";\n"; + L << " float acc0 = " << init << ", acc1 = " << init << ", acc2 = " << init + << ", acc3 = " << init << ";\n"; L << " if (col < cols) {\n"; L << " uint i = wb;\n"; L << " for (; i + 24u < rows; i += 32u) {\n"; for (int k = 0; k < 4; ++k) { - L << " { const uint row = i + " << std::to_string(8 * k) - << "u; "; + L << " { const uint row = i + " << std::to_string(8 * k) << "u; "; emit_reg_decl(); emit_instrs(0, ri); L << " acc" << k << " = " << combine(is_max, "acc" + std::to_string(k), src) @@ -3315,8 +3167,8 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, L << " acc0 = " << combine(is_max, "acc0", src) << "; }\n"; L << " }\n"; L << " }\n"; - L << " s_red[" << slot << "][tid] = " - << combine4(is_max, "acc0", "acc1", "acc2", "acc3") << ";\n"; + L << " s_red[" << slot + << "][tid] = " << combine4(is_max, "acc0", "acc1", "acc2", "acc3") << ";\n"; L << " }\n"; emit_col_merge(slot, is_max); } @@ -3330,11 +3182,14 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, { L << " if (col < cols) {\n"; emit_wr(" ", "col", - "s_red[" + std::to_string(slot_of_instr[spec.instrs.back().dst]) + "][tid]"); + "s_red[" + std::to_string(slot_of_instr[spec.instrs.back().dst]) + + "][tid]"); L << " }\n"; } else - L << " if (col < cols) bout[col] = " << wr("s_red[" + std::to_string(slot_of_instr[spec.instrs.back().dst]) + "][tid]") << ";\n"; + L << " if (col < cols) bout[col] = " + << wr("s_red[" + std::to_string(slot_of_instr[spec.instrs.back().dst]) + "][tid]") + << ";\n"; } else { @@ -3376,60 +3231,57 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec, class GlslEmitter final : public ExprEmitter { public: - [[nodiscard]] std::string_view name() const noexcept override - { return "glsl"; } + [[nodiscard]] std::string_view name() const noexcept override { return "glsl"; } - [[nodiscard]] std::string generate( - const std::string& name_, const ExprSpec& spec, - ExprPrecSig sig = 0, bool native16 = false) override + [[nodiscard]] std::string generate(const std::string& name_, const ExprSpec& spec, + ExprPrecSig sig = 0, bool native16 = false) override { return nn::generate_glsl(name_, spec, sig, native16); } - [[nodiscard]] std::string generate_reduce( - const std::string& name_, const ExprSpec& spec, - ExprPrecSig sig = 0) override + [[nodiscard]] std::string generate_reduce(const std::string& name_, const ExprSpec& spec, + ExprPrecSig sig = 0) override { return nn::generate_glsl_reduce(name_, spec, sig); } // P3 批量变体（键 = key#b）：纯逐元素 + V0 全 f32 的多实例批量派发形态 // （实例表 + buffer device address）。不支持的形态返回空串 → 上层跳过 // 该变体，运行期逐 kernel 降级（见 expr_spec_batchable）。 - [[nodiscard]] std::string generate_batched( - const std::string& name_, const ExprSpec& spec, - ExprPrecSig sig = 0) const override - { return nn::generate_glsl(name_, spec, sig, /*native16=*/false, - /*batched=*/true); } + [[nodiscard]] std::string generate_batched(const std::string& name_, const ExprSpec& spec, + ExprPrecSig sig = 0) const override + { + return nn::generate_glsl(name_, spec, sig, /*native16=*/false, + /*batched=*/true); + } // matmul 细 tile 变体（键 = key#f，P0 wgrad 网格饥饿）：输出块 16×16、 // 每线程 2×2 输出（R=2）的同结构 shader。不支持的形态（非纯 matmul） // 返回空串 → 上层跳过该变体，运行期留在 64-tile 路径（降级，不是错误）。 - [[nodiscard]] std::string generate_matmul_fine( - const std::string& name_, const ExprSpec& spec, - ExprPrecSig sig = 0) const override - { return nn::generate_glsl(name_, spec, sig, /*native16=*/false, - /*batched=*/false, /*mm_fine=*/true); } + [[nodiscard]] std::string generate_matmul_fine(const std::string& name_, const ExprSpec& spec, + ExprPrecSig sig = 0) const override + { + return nn::generate_glsl(name_, spec, sig, /*native16=*/false, + /*batched=*/false, /*mm_fine=*/true); + } // matmul K 切分变体（键 = key#ks）：细 tile 网格上把 K 切成 S 段，每 WG // 一段写部分和（不跑尾链）。不支持的形态（非纯 matmul/非 V0）返回空串 // → 上层跳过该变体，运行期留在细 tile 路径（降级，不是错误）。 - [[nodiscard]] std::string generate_matmul_ksplit( - const std::string& name_, const ExprSpec& spec, - ExprPrecSig sig = 0) const override - { return nn::generate_glsl(name_, spec, sig, /*native16=*/false, - /*batched=*/false, /*mm_fine=*/false, - /*mm_ksplit=*/true); } + [[nodiscard]] std::string generate_matmul_ksplit(const std::string& name_, const ExprSpec& spec, + ExprPrecSig sig = 0) const override + { + return nn::generate_glsl(name_, spec, sig, /*native16=*/false, + /*batched=*/false, /*mm_fine=*/false, + /*mm_ksplit=*/true); + } // K 切分伴生归约（键 = key#kr）：跨段升序求和 + 同一条尾链。 // 资格判据在 generate_glsl_matmul_ksplit_reduce 内，与 #ks 同源。 - [[nodiscard]] std::string generate_matmul_ksplit_reduce( - const std::string& name_, const ExprSpec& spec, - ExprPrecSig sig = 0) const override + [[nodiscard]] std::string generate_matmul_ksplit_reduce(const std::string& name_, + const ExprSpec& spec, + ExprPrecSig sig = 0) const override { return nn::generate_glsl_matmul_ksplit_reduce(name_, spec, sig); } }; // 登记到 emitter 注册表（静态初始化；重复包含无害——同名拒绝覆盖） -inline const bool kGlslEmitterRegistered = - emitter_registry::register_backend("glsl", - []() -> std::unique_ptr { - return std::make_unique(); - }); +inline const bool kGlslEmitterRegistered = emitter_registry::register_backend( + "glsl", []() -> std::unique_ptr { return std::make_unique(); }); } // namespace nn diff --git a/include/neuralnet.cpp/expr_graph.hpp b/include/neuralnet.cpp/expr_graph.hpp index 61e1087f..5c07d5b5 100644 --- a/include/neuralnet.cpp/expr_graph.hpp +++ b/include/neuralnet.cpp/expr_graph.hpp @@ -64,11 +64,11 @@ namespace nn // ── 图 IR 节点 ──────────────────────────────────────────────────────────── struct ExprGraphNode { - ExprSpec spec; // canonical IR（录制时已 canonicalize） + 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)） - Precision prec = Precision::F32; // 求值精度（恢复时新增：融合成员须一致） + bool vector_out = false; // 输出为归约向量（(rows,1)/(1,cols)） + Precision prec = Precision::F32; // 求值精度（恢复时新增：融合成员须一致） // RV-03：本节点的**写目标存储标识**（Tensor::storage_id()）。0 = 本节点输出 // 是一块新占位存储（dsl::compute 建的新张量）；非 0 = 目标传递节点 // （dsl::compute_into / eval_expr_into 写进既有张量）。用于把"同目标的多次写" @@ -101,28 +101,24 @@ struct ExprGraph // 录制：把一个表达式加入图，返回节点下标。 // 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, - Precision prec = Precision::F32, - std::uintptr_t dst_id = 0) + [[nodiscard]] int add_node(const ExprSpec& raw_spec, std::span inputs, + std::size_t rows, std::size_t cols, bool vector_out, + Precision prec = Precision::F32, std::uintptr_t dst_id = 0) { 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.prec = prec; - nd.dst_id = dst_id; // RV-03：目标传递节点的写目标（0 = 新占位存储） + nd.spec = spec; + nd.rows = rows; + nd.cols = cols; + nd.vector_out = vector_out; + nd.prec = prec; + nd.dst_id = dst_id; // RV-03：目标传递节点的写目标（0 = 新占位存储） 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 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()) { @@ -132,8 +128,7 @@ struct ExprGraph else { nd.dep_of_input.push_back(-1); - nd.input_tensors.push_back( - (k < inputs.size()) ? inputs[k] : Tensor{}); + nd.input_tensors.push_back((k < inputs.size()) ? inputs[k] : Tensor{}); } } nodes.push_back(std::move(nd)); @@ -149,23 +144,23 @@ struct ExprGraph // 一个 kernel = 一个（可能为多个节点拼接的）复合 ExprSpec + 外部输入来源。 struct FusedKernelInput { - int node = -1; // >= 0：该输入是图节点 node 的输出（作为外部 buffer 绑定） - Tensor external; // node < 0：Layer 外部输入（运行时张量） + int node = -1; // >= 0：该输入是图节点 node 的输出（作为外部 buffer 绑定） + Tensor external; // node < 0：Layer 外部输入（运行时张量） // P2-12（图级缓存跨 step 复用）：node < 0 时记录该外部输入的来源， // 以便从缓存计划反查当前图对应节点的输入槽（input_tensors[ext_slot]）。 - int ext_node = -1; // 来源原图节点 index - int ext_slot = -1; // 来源节点输入槽 index + int ext_node = -1; // 来源原图节点 index + int ext_slot = -1; // 来源节点输入槽 index }; struct FusedKernel { - ExprSpec spec; // 复合 spec（canonical） + ExprSpec spec; // 复合 spec（canonical） std::vector inputs; // 与 spec.views 一一对应 std::size_t rows = 0; std::size_t cols = 0; - bool vector_out = false; - Precision prec = Precision::F32; // 成员共同精度（恢复时新增） - int tail = -1; // 末尾节点（输出节点）在图中下标 + bool vector_out = false; + Precision prec = Precision::F32; // 成员共同精度（恢复时新增） + int tail = -1; // 末尾节点（输出节点）在图中下标 std::vector members; // kernel 包含的节点（顺序 = 融合顺序） }; @@ -220,16 +215,14 @@ namespace fused const std::uintptr_t id = t.storage_id(); if (id == 0) continue; // 无存储（不该发生） const auto w = last_writer.find(id); - if (w != last_writer.end() && w->second < i) - out.emplace_back(w->second, i); + if (w != last_writer.end() && w->second < i) out.emplace_back(w->second, i); readers[id].push_back(i); } // ② 写：与上一次写（WAW）、以及上次写之后的所有读者（WAR）建边 const std::uintptr_t d = n.dst_id; if (d == 0) continue; // 新占位存储，无写序隐患 const auto w = last_writer.find(d); - if (w != last_writer.end() && w->second < i) - out.emplace_back(w->second, i); + if (w != last_writer.end() && w->second < i) out.emplace_back(w->second, i); const auto r = readers.find(d); if (r != readers.end()) { @@ -246,8 +239,8 @@ namespace fused // RV-03 兜底校验：给定**执行序**下的 kernel 列表，检查每条顺序隐患边都保序。 // 违反 = 静默错值的隐患已经成真 → fail-fast（绝不静默算错）。 -[[nodiscard]] inline Result validate_dst_order( - const ExprGraph& g, std::span ordered_kernels) +[[nodiscard]] inline Result validate_dst_order(const ExprGraph& g, + std::span ordered_kernels) { const std::vector> edges = dst_order_constraints(g); if (edges.empty()) return {}; @@ -257,8 +250,7 @@ namespace fused for (std::size_t m = 0; m < ordered_kernels[k].members.size(); ++m) { const int node = ordered_kernels[k].members[m]; - if (node < 0 || static_cast(node) >= g.nodes.size()) - continue; + if (node < 0 || static_cast(node) >= g.nodes.size()) continue; kernel_pos[static_cast(node)] = static_cast(k); member_pos[static_cast(node)] = static_cast(m); } @@ -269,12 +261,11 @@ namespace fused if (pi < 0 || pj < 0) continue; // 未被任何 kernel 覆盖（不该发生） if (pi < pj) continue; // i 的 kernel 先执行 → 保序 // 同一 kernel 内：成员序必须保持 i 在 j 之前 - if (pi == pj && member_pos[static_cast(i)] < - member_pos[static_cast(j)]) + if (pi == pj && + member_pos[static_cast(i)] < member_pos[static_cast(j)]) continue; - NN_FAIL("IR-C 顺序隐患被破坏（同目标写序/WAR 要求节点 i 先于 j 执行）：i=" - + std::to_string(i) + " j=" + std::to_string(j) - + "（分组或重排拆掉了显式顺序边）"); + NN_FAIL("IR-C 顺序隐患被破坏（同目标写序/WAR 要求节点 i 先于 j 执行）：i=" + + std::to_string(i) + " j=" + std::to_string(j) + "（分组或重排拆掉了显式顺序边）"); } return {}; } @@ -285,15 +276,13 @@ namespace fused [[nodiscard]] inline std::vector fuse_expr_graph(const ExprGraph& g) { std::vector kernels; - if (g.nodes.empty()) - return 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)); + if (dep >= 0) consumers[static_cast(dep)].push_back(static_cast(i)); // RV-03：顺序隐患边（同目标写序 / WAR）——有隐患关系的两个节点**不得**合进 // 同一 kernel：同一 dispatch 内的多次写，其先后不再由 kernel 执行序决定。 @@ -302,8 +291,7 @@ namespace fused { if (a == b) return false; if (a > b) std::swap(a, b); - return std::binary_search(order_edges.begin(), order_edges.end(), - std::make_pair(a, b)); + return std::binary_search(order_edges.begin(), order_edges.end(), std::make_pair(a, b)); }; // ── 尝试把节点 bi 并入当前 kernel（cur） ─────────────────────────── @@ -349,8 +337,7 @@ namespace fused other_slots.push_back(static_cast(k)); } } - if (tail_slots.empty()) - return false; + if (tail_slots.empty()) return false; // 3) 拼接（确定性：tail 的 views/consts/instrs 在前，B 的在后） // A 纯 matmul（空链，S5）时无输出寄存器：B 对 A 输出的引用替换为 @@ -365,8 +352,7 @@ namespace fused // 寄存器空洞（A 的寄存器 0..num_regs-1 之后空出 instrs.size()-num_regs 个号）， // 浪费寄存器并可能使本可融合的表达式超出 EXPR_MAX_REGS=16 而被放弃。 for (const auto& ins : Bs.instrs) - if (static_cast(ins.dst) + A.num_regs >= 256) - return false; + if (static_cast(ins.dst) + A.num_regs >= 256) return false; std::vector views = A.views; std::vector b_view_map(Bs.views.size(), -1); @@ -402,20 +388,17 @@ namespace fused 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)); + const auto it = + std::find(tail_slots.begin(), tail_slots.end(), static_cast(k)); if (it != tail_slots.end()) { - if (a_pure_matmul) - return expr::matmul_op(); // 引用前置 matmul 段输出 + if (a_pure_matmul) return expr::matmul_op(); // 引用前置 matmul 段输出 return expr::reg(static_cast(out_reg_A)); } - return expr::input(static_cast( - b_view_map[k])); + 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)); + return expr::cst(static_cast(o.idx + b_const_base)); if (o.kind == static_cast(ExprOperandKind::RParam)) return {o.kind, static_cast(o.idx + b_rp_base)}; if (o.kind == static_cast(ExprOperandKind::Reg) || @@ -424,8 +407,7 @@ namespace fused return {o.kind, static_cast(o.idx + reg_base)}; return o; }; - const std::size_t nops = - expr_instr_num_operands(static_cast(ins.op)); + 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); @@ -433,12 +415,12 @@ namespace fused } ExprSpec fused; - fused.views = std::move(views); - fused.consts = std::move(consts); + fused.views = std::move(views); + fused.consts = std::move(consts); fused.rparams = std::move(rparams); - fused.instrs = std::move(instrs); + fused.instrs = std::move(instrs); fused.num_regs = static_cast(reg_base + Bs.num_regs); - fused.matmul = A.matmul; // 保留前置 matmul 段（S5：matmul 参与图融合） + fused.matmul = A.matmul; // 保留前置 matmul 段（S5：matmul 参与图融合） fused = canonicalize_expr_spec(fused); if (auto v = validate_expr_spec(fused, fused.views.size()); !v) return false; // 超限/非法 → 放弃融合（保守） @@ -449,19 +431,19 @@ namespace fused { 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], - /*ext_node=*/bi, /*ext_slot=*/static_cast(kk)}); + new_inputs.push_back(FusedKernelInput{/*node=*/-1, /*external=*/B.input_tensors[kk], + /*ext_node=*/bi, + /*ext_slot=*/static_cast(kk)}); else - new_inputs.push_back(FusedKernelInput{ - /*node=*/B.dep_of_input[kk], /*external=*/Tensor{}}); + 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.prec = B.prec; - cur.tail = bi; + cur.spec = std::move(fused); + cur.inputs = std::move(new_inputs); + cur.rows = B.rows; + cur.cols = B.cols; + cur.prec = B.prec; + cur.tail = bi; cur.members.push_back(bi); return true; }; @@ -474,9 +456,8 @@ namespace fused const auto p2_eligible = [&](std::size_t i) -> bool { const auto& N = g.nodes[i]; - return !N.vector_out && N.prec == Precision::F32 && - !N.spec.instrs.empty() && !N.spec.matmul && !N.spec.fold && - expr_spec_reduce_axis(N.spec) == -1; + return !N.vector_out && N.prec == Precision::F32 && !N.spec.instrs.empty() && + !N.spec.matmul && !N.spec.fold && expr_spec_reduce_axis(N.spec) == -1; }; std::vector uf_parent(g.nodes.size()); for (std::size_t i = 0; i < uf_parent.size(); ++i) @@ -486,32 +467,27 @@ namespace fused while (uf_parent[static_cast(x)] != x) { uf_parent[static_cast(x)] = - uf_parent[static_cast( - uf_parent[static_cast(x)])]; + uf_parent[static_cast(uf_parent[static_cast(x)])]; x = uf_parent[static_cast(x)]; } return x; }; for (std::size_t i = 0; i < g.nodes.size(); ++i) { - if (!p2_eligible(i)) - continue; + if (!p2_eligible(i)) continue; for (const int d : g.nodes[i].dep_of_input) { - if (d < 0 || !p2_eligible(static_cast(d))) - continue; + if (d < 0 || !p2_eligible(static_cast(d))) continue; const auto& A = g.nodes[static_cast(d)]; const auto& B = g.nodes[i]; - if (A.rows != B.rows || A.cols != B.cols || A.prec != B.prec) - continue; + if (A.rows != B.rows || A.cols != B.cols || A.prec != B.prec) continue; const int ra = uf_find(d); const int rb = uf_find(static_cast(i)); if (ra != rb) { // RV-03：先检查顺序隐患——有隐患的节点不得进入同一分量 // （同一 kernel 内的写序不再由 kernel 执行序决定） - if (has_hazard(d, static_cast(i))) - continue; + if (has_hazard(d, static_cast(i))) continue; uf_parent[static_cast(std::max(ra, rb))] = std::min(ra, rb); // 小根为根：确定性 } @@ -529,30 +505,29 @@ namespace fused { const auto& N = g.nodes[i]; FusedKernel k; - k.spec = N.spec; - k.rows = N.rows; - k.cols = N.cols; + k.spec = N.spec; + k.rows = N.rows; + k.cols = N.cols; k.vector_out = N.vector_out; - k.prec = N.prec; - k.tail = static_cast(i); + k.prec = N.prec; + 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 (kk >= N.dep_of_input.size() || N.dep_of_input[kk] < 0) k.inputs.push_back(FusedKernelInput{ /*node=*/-1, - /*external=*/(kk < N.input_tensors.size()) ? N.input_tensors[kk] - : Tensor{}, + /*external=*/(kk < N.input_tensors.size()) ? N.input_tensors[kk] : Tensor{}, /*ext_node=*/static_cast(i), /*ext_slot=*/static_cast(kk)}); else - k.inputs.push_back(FusedKernelInput{ - /*node=*/N.dep_of_input[kk], /*external=*/Tensor{}}); + k.inputs.push_back( + FusedKernelInput{/*node=*/N.dep_of_input[kk], /*external=*/Tensor{}}); } kernels.push_back(std::move(k)); k_members.push_back({static_cast(i)}); k_out_reg.push_back(N.spec.instrs.empty() - ? std::vector{} - : std::vector{N.spec.instrs.back().dst}); + ? std::vector{} + : std::vector{N.spec.instrs.back().dst}); k_is_comp.push_back(comp ? 1 : 0); return kernels.size() - 1; }; @@ -569,16 +544,14 @@ namespace fused const auto& outs0 = k_out_reg[ki]; const ExprGraphNode& B = g.nodes[bi]; const ExprSpec& Bs = B.spec; - if (Bs.instrs.empty() || Bs.matmul || Bs.fold) - return false; + if (Bs.instrs.empty() || Bs.matmul || Bs.fold) return false; // RV-03：同目标写序 / WAR 隐患的两端不合组（见 has_hazard 的说明） for (const int m : mem0) if (has_hazard(m, static_cast(bi))) return false; const auto member_idx = [&](int dep) -> int { for (std::size_t m = 0; m < mem0.size(); ++m) - if (mem0[m] == dep) - return static_cast(m); + if (mem0[m] == dep) return static_cast(m); return -1; }; @@ -605,27 +578,19 @@ namespace fused if (dep < 0) inputs.push_back(FusedKernelInput{ /*node=*/-1, - /*external=*/(k < B.input_tensors.size()) ? B.input_tensors[k] - : Tensor{}, + /*external=*/(k < B.input_tensors.size()) ? B.input_tensors[k] : Tensor{}, /*ext_node=*/static_cast(bi), /*ext_slot=*/static_cast(k)}); - else - inputs.push_back(FusedKernelInput{ - /*node=*/dep, /*external=*/Tensor{}}); + else inputs.push_back(FusedKernelInput{/*node=*/dep, /*external=*/Tensor{}}); } } // 预算预检（保守：按拼接后原始计数；上限 = validate_expr_spec 同源） - if (cur.spec.instrs.size() + Bs.instrs.size() > EXPR_MAX_INSTRS) - return false; - if (cur.spec.num_regs + Bs.num_regs > 255) - return false; - if (cur.spec.consts.size() + Bs.consts.size() > EXPR_MAX_CONSTS) - return false; - if (cur.spec.rparams.size() + Bs.rparams.size() > EXPR_MAX_CONSTS) - return false; - if (views.size() > EXPR_MAX_INPUTS) - return false; + if (cur.spec.instrs.size() + Bs.instrs.size() > EXPR_MAX_INSTRS) return false; + if (cur.spec.num_regs + Bs.num_regs > 255) return false; + if (cur.spec.consts.size() + Bs.consts.size() > EXPR_MAX_CONSTS) return false; + if (cur.spec.rparams.size() + Bs.rparams.size() > EXPR_MAX_CONSTS) return false; + if (views.size() > EXPR_MAX_INPUTS) return false; std::vector consts = cur.spec.consts; const std::size_t b_const_base = consts.size(); @@ -654,8 +619,7 @@ namespace fused return expr::input(static_cast(slot_map[k])); } if (o.kind == static_cast(ExprOperandKind::Const)) - return expr::cst(static_cast( - o.idx + b_const_base)); + return expr::cst(static_cast(o.idx + b_const_base)); if (o.kind == static_cast(ExprOperandKind::RParam)) return {o.kind, static_cast(o.idx + b_rp_base)}; if (o.kind == static_cast(ExprOperandKind::Reg) || @@ -664,8 +628,7 @@ namespace fused return {o.kind, static_cast(o.idx + reg_base)}; return o; }; - const std::size_t nops = - expr_instr_num_operands(static_cast(ins.op)); + const std::size_t nops = expr_instr_num_operands(static_cast(ins.op)); ni.a = remap(ins.a); if (nops >= 2) ni.b = remap(ins.b); if (nops >= 3) ni.c = remap(ins.c); @@ -673,27 +636,27 @@ namespace fused } ExprSpec fused; - fused.views = std::move(views); - fused.consts = std::move(consts); + fused.views = std::move(views); + fused.consts = std::move(consts); fused.rparams = std::move(rparams); - fused.instrs = std::move(instrs); + fused.instrs = std::move(instrs); fused.num_regs = static_cast(reg_base + Bs.num_regs); outs.push_back(static_cast(Bs.instrs.back().dst + reg_base)); fused.extras.assign(outs.begin(), outs.end() - 1); if (auto v = validate_expr_spec(fused, fused.views.size()); !v) return false; // 超限/非法 → 放弃并入（保守） - cur.spec = std::move(fused); - cur.inputs = std::move(inputs); - cur.rows = B.rows; - cur.cols = B.cols; - cur.prec = B.prec; - cur.tail = static_cast(bi); + cur.spec = std::move(fused); + cur.inputs = std::move(inputs); + cur.rows = B.rows; + cur.cols = B.cols; + cur.prec = B.prec; + cur.tail = static_cast(bi); cur.members.push_back(static_cast(bi)); mem.push_back(static_cast(bi)); // 簿记同步（供后续成员的依赖查找） - kernels[ki] = std::move(cur); - k_members[ki] = std::move(mem); - k_out_reg[ki] = std::move(outs); + kernels[ki] = std::move(cur); + k_members[ki] = std::move(mem); + k_out_reg[ki] = std::move(outs); return true; }; @@ -704,8 +667,8 @@ namespace fused // 按 P1 链语义。分量 kernel 与 P1 kernel 互不混拼。 for (std::size_t i = 0; i < g.nodes.size(); ++i) { - if (!kernels.empty() && kernels.back().spec.matmul && - !k_is_comp.back() && try_append(kernels.back(), static_cast(i))) + if (!kernels.empty() && kernels.back().spec.matmul && !k_is_comp.back() && + try_append(kernels.back(), static_cast(i))) { k_members.back().push_back(static_cast(i)); k_out_reg.back().clear(); // P1 单输出（不写穿） @@ -721,8 +684,7 @@ namespace fused comp_kernel.emplace(root, open_kernel(i, /*comp=*/true)); continue; } - if (comp_append(it->second, i)) - continue; + if (comp_append(it->second, i)) continue; // 预算不足：分裂——同分量余下成员进新 kernel（跨 kernel 依赖经 // 写穿输出读取），映射顶替为新 kernel it->second = open_kernel(i, /*comp=*/true); @@ -757,15 +719,15 @@ namespace fused // 缓存的 kernel 计划（不含运行时张量） struct FusedKernelPlan { - ExprSpec spec; // 复合 spec（canonical） - std::vector input_node; // 每输入槽：来源节点 idx（>=0）或 -1（外部） - std::vector input_ext_node; // 外部时：来源原图节点 idx - std::vector input_ext_slot; // 外部时：来源节点输入槽 index - std::size_t rows = 0, cols = 0; - bool vector_out = false; - Precision prec = Precision::F32; - int tail = -1; - std::vector members; + ExprSpec spec; // 复合 spec（canonical） + std::vector input_node; // 每输入槽：来源节点 idx（>=0）或 -1（外部） + std::vector input_ext_node; // 外部时：来源原图节点 idx + std::vector input_ext_slot; // 外部时：来源节点输入槽 index + std::size_t rows = 0, cols = 0; + bool vector_out = false; + Precision prec = Precision::F32; + int tail = -1; + std::vector members; }; // 整图结构 key（FNV-1a）：覆盖节点 spec key、形状、vector_out、精度、依赖拓扑。 @@ -807,8 +769,8 @@ struct FusedKernelPlan int grp = 0; if (n.dst_id != 0) { - const auto [it, inserted] = dst_group.emplace( - n.dst_id, static_cast(dst_group.size()) + 1); + const auto [it, inserted] = + dst_group.emplace(n.dst_id, static_cast(dst_group.size()) + 1); (void)inserted; grp = it->second; } @@ -826,13 +788,13 @@ struct FusedKernelPlan [[nodiscard]] inline FusedKernelPlan plan_from_kernel(const FusedKernel& k) { FusedKernelPlan p; - p.spec = k.spec; - p.rows = k.rows; - p.cols = k.cols; - p.vector_out = k.vector_out; - p.prec = k.prec; - p.tail = k.tail; - p.members = k.members; + p.spec = k.spec; + p.rows = k.rows; + p.cols = k.cols; + p.vector_out = k.vector_out; + p.prec = k.prec; + p.tail = k.tail; + p.members = k.members; p.input_node.reserve(k.inputs.size()); p.input_ext_node.reserve(k.inputs.size()); p.input_ext_slot.reserve(k.inputs.size()); @@ -849,13 +811,13 @@ struct FusedKernelPlan [[nodiscard]] inline FusedKernel instantiate_plan(const FusedKernelPlan& p, const ExprGraph& g) { FusedKernel k; - k.spec = p.spec; - k.rows = p.rows; - k.cols = p.cols; - k.vector_out = p.vector_out; - k.prec = p.prec; - k.tail = p.tail; - k.members = p.members; + k.spec = p.spec; + k.rows = p.rows; + k.cols = p.cols; + k.vector_out = p.vector_out; + k.prec = p.prec; + k.tail = p.tail; + k.members = p.members; k.inputs.reserve(p.input_node.size()); for (std::size_t i = 0; i < p.input_node.size(); ++i) { @@ -866,9 +828,9 @@ struct FusedKernelPlan in.ext_node = p.input_ext_node[i]; in.ext_slot = p.input_ext_slot[i]; in.external = (in.ext_node >= 0 && in.ext_slot >= 0) - ? g.nodes[static_cast(in.ext_node)] - .input_tensors[static_cast(in.ext_slot)] - : Tensor{}; + ? g.nodes[static_cast(in.ext_node)] + .input_tensors[static_cast(in.ext_slot)] + : Tensor{}; } k.inputs.push_back(std::move(in)); } @@ -896,12 +858,10 @@ struct FusedKernelPlan namespace fused { -[[nodiscard]] inline std::unordered_map>& +[[nodiscard]] inline std::unordered_map>& graph_plan_cache() { - static thread_local std::unordered_map> c; + static thread_local std::unordered_map> c; return c; } inline constexpr std::size_t GRAPH_PLAN_CACHE_MAX = 512; diff --git a/include/neuralnet.cpp/expr_op_table.hpp b/include/neuralnet.cpp/expr_op_table.hpp index c7d4b53b..1cad6137 100644 --- a/include/neuralnet.cpp/expr_op_table.hpp +++ b/include/neuralnet.cpp/expr_op_table.hpp @@ -42,79 +42,76 @@ enum class ExprEmitKind : std::uint8_t }; // CPU 标量语义函数指针（比较类返回 1.0/0.0 的 Scalar，与写回寄存器的语义一致） -using ExprUnaryFn = Scalar (*)(Scalar) noexcept; +using ExprUnaryFn = Scalar (*)(Scalar) noexcept; using ExprBinaryFn = Scalar (*)(Scalar, Scalar) noexcept; struct ExprOpInfo { std::uint8_t arity; // 实际使用操作数个数（== expr_instr_num_operands） - bool is_reduce; // dst 是隐式"每行/每列一个标量"的归约向量 - bool reduces_cols; // 归约方向：true=按列 → (1,cols)，false=按行 → (rows,1) - bool foldable; // 常量折叠（仅 IEEE 754 基本运算；不含超越函数/归约） - ExprUnaryFn unary; // arity==1 且非归约：unary(a) + bool is_reduce; // dst 是隐式"每行/每列一个标量"的归约向量 + bool reduces_cols; // 归约方向：true=按列 → (1,cols)，false=按行 → (rows,1) + bool foldable; // 常量折叠（仅 IEEE 754 基本运算；不含超越函数/归约） + ExprUnaryFn unary; // arity==1 且非归约：unary(a) ExprBinaryFn binary; // arity==2：binary(a,b) ExprEmitKind emit; // GLSL 发射形态 - const char* sym; // 中缀符号 / 函数名 - const char* sym_vec; // Compare 的 vec4 内建名（否则 ""） + const char* sym; // 中缀符号 / 函数名 + const char* sym_vec; // Compare 的 vec4 内建名（否则 ""） }; // 比较算子：`ops::Lt` 等的 apply 返回 bool（供条件表达式），寄存器写回要 Scalar // 1.0/0.0——此处做唯一一次桥接。 namespace cmp { -[[nodiscard]] inline Scalar lt(Scalar a, Scalar b) noexcept { return ops::Lt::apply(a, b) ? Scalar{1} : Scalar{0}; } -[[nodiscard]] inline Scalar le(Scalar a, Scalar b) noexcept { return ops::Le::apply(a, b) ? Scalar{1} : Scalar{0}; } -[[nodiscard]] inline Scalar gt(Scalar a, Scalar b) noexcept { return ops::Gt::apply(a, b) ? Scalar{1} : Scalar{0}; } -[[nodiscard]] inline Scalar ge(Scalar a, Scalar b) noexcept { return ops::Ge::apply(a, b) ? Scalar{1} : Scalar{0}; } -[[nodiscard]] inline Scalar eq(Scalar a, Scalar b) noexcept { return ops::Eq::apply(a, b) ? Scalar{1} : Scalar{0}; } -[[nodiscard]] inline Scalar ne(Scalar a, Scalar b) noexcept { return ops::Ne::apply(a, b) ? Scalar{1} : Scalar{0}; } +[[nodiscard]] inline Scalar lt(Scalar a, Scalar b) noexcept +{ return ops::Lt::apply(a, b) ? Scalar{1} : Scalar{0}; } +[[nodiscard]] inline Scalar le(Scalar a, Scalar b) noexcept +{ return ops::Le::apply(a, b) ? Scalar{1} : Scalar{0}; } +[[nodiscard]] inline Scalar gt(Scalar a, Scalar b) noexcept +{ return ops::Gt::apply(a, b) ? Scalar{1} : Scalar{0}; } +[[nodiscard]] inline Scalar ge(Scalar a, Scalar b) noexcept +{ return ops::Ge::apply(a, b) ? Scalar{1} : Scalar{0}; } +[[nodiscard]] inline Scalar eq(Scalar a, Scalar b) noexcept +{ return ops::Eq::apply(a, b) ? Scalar{1} : Scalar{0}; } +[[nodiscard]] inline Scalar ne(Scalar a, Scalar b) noexcept +{ return ops::Ne::apply(a, b) ? Scalar{1} : Scalar{0}; } } // namespace cmp // 行构造器（字段全列——避开 -Wmissing-designated-field-initializers） namespace row { -[[nodiscard]] consteval ExprOpInfo unary(ExprUnaryFn f, ExprEmitKind e, - const char* sym, bool fold = false) noexcept -{ - return ExprOpInfo{1, false, false, fold, f, nullptr, e, sym, ""}; -} -[[nodiscard]] consteval ExprOpInfo binary(ExprBinaryFn f, ExprEmitKind e, - const char* sym, const char* sym_vec = "", - bool fold = true) noexcept -{ - return ExprOpInfo{2, false, false, fold, nullptr, f, e, sym, sym_vec}; -} +[[nodiscard]] consteval ExprOpInfo unary(ExprUnaryFn f, ExprEmitKind e, const char* sym, + bool fold = false) noexcept +{ return ExprOpInfo{1, false, false, fold, f, nullptr, e, sym, ""}; } +[[nodiscard]] consteval ExprOpInfo binary(ExprBinaryFn f, ExprEmitKind e, const char* sym, + const char* sym_vec = "", bool fold = true) noexcept +{ return ExprOpInfo{2, false, false, fold, nullptr, f, e, sym, sym_vec}; } [[nodiscard]] consteval ExprOpInfo select_op() noexcept -{ - return ExprOpInfo{3, false, false, true, nullptr, nullptr, ExprEmitKind::Select, "", ""}; -} +{ return ExprOpInfo{3, false, false, true, nullptr, nullptr, ExprEmitKind::Select, "", ""}; } [[nodiscard]] consteval ExprOpInfo reduce(bool cols) noexcept -{ - return ExprOpInfo{1, true, cols, false, nullptr, nullptr, ExprEmitKind::None, "", ""}; -} +{ return ExprOpInfo{1, true, cols, false, nullptr, nullptr, ExprEmitKind::None, "", ""}; } } // namespace row // 表按下标 == static_cast(ExprOp) 直接索引（顺序不得改动） inline constexpr ExprOpInfo kExprOpTable[] = { - /* Neg */ row::unary(&ops::Neg::apply, ExprEmitKind::Neg, "-", /*fold=*/true), - /* Exp */ row::unary(&ops::Exp::apply, ExprEmitKind::Func1, "exp"), - /* Log */ row::unary(&ops::Log::apply, ExprEmitKind::Func1, "log"), - /* Sqrt */ row::unary(&ops::Sqrt::apply, ExprEmitKind::Func1, "sqrt"), + /* Neg */ row::unary(&ops::Neg::apply, ExprEmitKind::Neg, "-", /*fold=*/true), + /* Exp */ row::unary(&ops::Exp::apply, ExprEmitKind::Func1, "exp"), + /* Log */ row::unary(&ops::Log::apply, ExprEmitKind::Func1, "log"), + /* Sqrt */ row::unary(&ops::Sqrt::apply, ExprEmitKind::Func1, "sqrt"), /* Rsqrt */ row::unary(&ops::Rsqrt::apply, ExprEmitKind::Func1, "inversesqrt"), - /* Abs */ row::unary(&ops::Abs::apply, ExprEmitKind::Func1, "abs", /*fold=*/true), - /* Tanh */ row::unary(&ops::Tanh::apply, ExprEmitKind::Func1, "tanh"), - /* Add */ row::binary(&ops::Add::apply, ExprEmitKind::Infix, "+"), - /* Sub */ row::binary(&ops::Sub::apply, ExprEmitKind::Infix, "-"), - /* Mul */ row::binary(&ops::Mul::apply, ExprEmitKind::Infix, "*"), - /* Div */ row::binary(&ops::Div::apply, ExprEmitKind::Infix, "/"), - /* Max */ row::binary(&ops::Max::apply, ExprEmitKind::Func2, "max"), - /* Min */ row::binary(&ops::Min::apply, ExprEmitKind::Func2, "min"), - /* Lt */ row::binary(&cmp::lt, ExprEmitKind::Compare, "<", "lessThan"), - /* Le */ row::binary(&cmp::le, ExprEmitKind::Compare, "<=", "lessThanEqual"), - /* Gt */ row::binary(&cmp::gt, ExprEmitKind::Compare, ">", "greaterThan"), - /* Ge */ row::binary(&cmp::ge, ExprEmitKind::Compare, ">=", "greaterThanEqual"), - /* Eq */ row::binary(&cmp::eq, ExprEmitKind::Compare, "==", "equal"), - /* Ne */ row::binary(&cmp::ne, ExprEmitKind::Compare, "!=", "notEqual"), + /* Abs */ row::unary(&ops::Abs::apply, ExprEmitKind::Func1, "abs", /*fold=*/true), + /* Tanh */ row::unary(&ops::Tanh::apply, ExprEmitKind::Func1, "tanh"), + /* Add */ row::binary(&ops::Add::apply, ExprEmitKind::Infix, "+"), + /* Sub */ row::binary(&ops::Sub::apply, ExprEmitKind::Infix, "-"), + /* Mul */ row::binary(&ops::Mul::apply, ExprEmitKind::Infix, "*"), + /* Div */ row::binary(&ops::Div::apply, ExprEmitKind::Infix, "/"), + /* Max */ row::binary(&ops::Max::apply, ExprEmitKind::Func2, "max"), + /* Min */ row::binary(&ops::Min::apply, ExprEmitKind::Func2, "min"), + /* Lt */ row::binary(&cmp::lt, ExprEmitKind::Compare, "<", "lessThan"), + /* Le */ row::binary(&cmp::le, ExprEmitKind::Compare, "<=", "lessThanEqual"), + /* Gt */ row::binary(&cmp::gt, ExprEmitKind::Compare, ">", "greaterThan"), + /* Ge */ row::binary(&cmp::ge, ExprEmitKind::Compare, ">=", "greaterThanEqual"), + /* Eq */ row::binary(&cmp::eq, ExprEmitKind::Compare, "==", "equal"), + /* Ne */ row::binary(&cmp::ne, ExprEmitKind::Compare, "!=", "notEqual"), /* Select*/ row::select_op(), /* ColSum*/ row::reduce(/*cols=*/true), /* ColMax*/ row::reduce(/*cols=*/true), @@ -122,12 +119,11 @@ inline constexpr ExprOpInfo kExprOpTable[] = { /* RowMax*/ row::reduce(/*cols=*/false), }; -[[nodiscard]] inline constexpr std::size_t expr_op_count() noexcept { return std::size(kExprOpTable); } +[[nodiscard]] inline constexpr std::size_t expr_op_count() noexcept +{ return std::size(kExprOpTable); } [[nodiscard]] inline constexpr const ExprOpInfo& expr_op_info(ExprOp op) noexcept -{ - return kExprOpTable[static_cast(op)]; -} +{ return kExprOpTable[static_cast(op)]; } // CPU 标量求值（一元 / 二元 / Select 三元）；**归约指令不得走本函数**（调用方分流）。 [[nodiscard]] inline Scalar expr_alu_scalar(ExprOp op, Scalar a, Scalar b, Scalar c) noexcept @@ -144,9 +140,9 @@ inline constexpr ExprOpInfo kExprOpTable[] = { for (std::size_t i = 0; i < expr_op_count(); ++i) { const ExprOp op = static_cast(i); - if (kExprOpTable[i].arity != expr_instr_num_operands(op)) return false; - if (kExprOpTable[i].is_reduce != expr_op_is_reduce(op)) return false; - if (kExprOpTable[i].reduces_cols != expr_op_reduces_cols(op)) return false; + if (kExprOpTable[i].arity != expr_instr_num_operands(op)) return false; + if (kExprOpTable[i].is_reduce != expr_op_is_reduce(op)) return false; + if (kExprOpTable[i].reduces_cols != expr_op_reduces_cols(op)) return false; } return true; } diff --git a/include/neuralnet.cpp/expr_opt.hpp b/include/neuralnet.cpp/expr_opt.hpp index 1888ac42..e4c59e39 100644 --- a/include/neuralnet.cpp/expr_opt.hpp +++ b/include/neuralnet.cpp/expr_opt.hpp @@ -41,19 +41,15 @@ namespace nn // ── 常量折叠求值（仅 IEEE 754 基本运算；不折叠超越函数/归约）──────────── // 返回 false 表示该算子不参与折叠（调用方保留原指令）。 -[[nodiscard]] inline bool expr_eval_const_fold(ExprOp op, Scalar a, Scalar b, - Scalar c, Scalar& out_v) +[[nodiscard]] inline bool expr_eval_const_fold(ExprOp op, Scalar a, Scalar b, Scalar c, + Scalar& out_v) { // 可折叠性 + 标量语义均来自 expr_op_table（超越函数 / 归约 foldable=false） const detail::ExprOpInfo& info = detail::expr_op_info(op); - if (!info.foldable) - return false; - if (info.arity == 1) - out_v = info.unary(a); - else if (info.arity == 2) - out_v = info.binary(a, b); - else - out_v = (a != Scalar{0}) ? b : c; // Select + if (!info.foldable) return false; + if (info.arity == 1) out_v = info.unary(a); + else if (info.arity == 2) out_v = info.binary(a, b); + else out_v = (a != Scalar{0}) ? b : c; // Select return true; } @@ -67,16 +63,14 @@ namespace nn // 注意：0-x→neg(x) 需生成新指令，不在此处理（调用方负责）。 // 需要向常量池写入时（x*0→0）经 consts 就地追加（确定性：保持出现序）。 // ═══════════════════════════════════════════════════════════════════════════ -[[nodiscard]] inline std::optional expr_try_algebraic_simplify( - ExprOp op, const ExprOperand& a, const ExprOperand& b, const ExprOperand& c, - std::vector& consts, - const std::vector& def_op, const std::vector& def_a) +[[nodiscard]] inline std::optional +expr_try_algebraic_simplify(ExprOp op, const ExprOperand& a, const ExprOperand& b, + const ExprOperand& c, std::vector& consts, + const std::vector& def_op, + const std::vector& def_a) { const auto is_const = [&](const ExprOperand& o, Scalar v) - { - return o.kind == static_cast(ExprOperandKind::Const) && - consts[o.idx] == v; - }; + { return o.kind == static_cast(ExprOperandKind::Const) && consts[o.idx] == v; }; const auto add_const_op = [&](Scalar v) -> ExprOperand { for (std::size_t j = 0; j < consts.size(); ++j) @@ -147,7 +141,7 @@ namespace nn [[nodiscard]] inline ExprSpec fold_constants_and_algebra(const ExprSpec& spec) { ExprSpec out; - out.views = spec.views; // 不变量：views 不变 + out.views = spec.views; // 不变量：views 不变 out.matmul = spec.matmul; // 不变量：matmul 段不变（pass 只优化 instrs/consts） out.rparams = spec.rparams; // RParam 值序列：非结构、不进 key，原样透传 // （值不参与常量折叠/代数化简，跨 pass 保持顺序） @@ -158,7 +152,11 @@ namespace nn { int hit = -1; for (std::size_t j = 0; j < out.consts.size(); ++j) - if (out.consts[j] == spec.consts[i]) { hit = static_cast(j); break; } + if (out.consts[j] == spec.consts[i]) + { + hit = static_cast(j); + break; + } if (hit < 0) { hit = static_cast(out.consts.size()); @@ -176,7 +174,7 @@ namespace nn // replace[r]：寄存器 r 化简后等价的操作数（非空则引用处替换） std::vector> replace(spec.num_regs); // 定义 r 的指令（化简后）的 op 与 a 操作数（用于 neg(neg(x)) 化简） - std::vector def_op(spec.num_regs, ExprOp::Add); + std::vector def_op(spec.num_regs, ExprOp::Add); std::vector def_a(spec.num_regs, ExprOperand{}); const auto remap_op = [&](const ExprOperand& op) -> ExprOperand @@ -191,11 +189,15 @@ namespace nn return op; }; - const auto emit = [&](std::uint8_t opc, std::uint8_t dst, - const ExprOperand& a, const ExprOperand& b, - const ExprOperand& c) -> void + const auto emit = [&](std::uint8_t opc, std::uint8_t dst, const ExprOperand& a, + const ExprOperand& b, const ExprOperand& c) -> void { - ExprInstr ni; ni.op = opc; ni.dst = dst; ni.a = a; ni.b = b; ni.c = c; + ExprInstr ni; + ni.op = opc; + ni.dst = dst; + ni.a = a; + ni.b = b; + ni.c = c; out.instrs.push_back(ni); def_op[dst] = static_cast(opc); def_a[dst] = a; @@ -210,7 +212,11 @@ namespace nn // P2 多输出（extras）：额外输出的定义指令同输出指令处理（只重映射 // 操作数，不折叠 / 不 alias）——值必须落在真实寄存器供输出存储。 for (const auto e : spec.extras) - if (e == ins.dst) { is_output = true; break; } + if (e == ins.dst) + { + is_output = true; + break; + } // 只重映射实际使用的操作数（未使用的 b/c 是默认哨兵 {0,0}， // 不得当 Reg(0) 重编号——见 docs/internal/03-ir-optimization.md 关键坑 ②） @@ -239,7 +245,10 @@ namespace nn const ExprOperand* ops[3] = {&a, &b, &c}; for (std::size_t oi = 0; oi < nops; ++oi) if (ops[oi]->kind != static_cast(ExprOperandKind::Const)) - { all_const = false; break; } + { + all_const = false; + break; + } if (all_const) { Scalar v = Scalar{0}; @@ -249,8 +258,12 @@ namespace nn // 常量入池（去重） std::optional cv; for (std::size_t j = 0; j < out.consts.size(); ++j) - if (out.consts[j] == v) { cv = {static_cast(ExprOperandKind::Const), - static_cast(j)}; break; } + if (out.consts[j] == v) + { + cv = {static_cast(ExprOperandKind::Const), + static_cast(j)}; + break; + } if (!cv) { const std::uint8_t ci = static_cast(out.consts.size()); @@ -270,8 +283,7 @@ namespace nn emit(static_cast(ExprOp::Neg), ins.dst, b, {}, {}); continue; } - if (auto repl = expr_try_algebraic_simplify(op, a, b, c, out.consts, - def_op, def_a)) + if (auto repl = expr_try_algebraic_simplify(op, a, b, c, out.consts, def_op, def_a)) { replace[dst] = *repl; continue; @@ -294,8 +306,7 @@ namespace nn // ═══════════════════════════════════════════════════════════════════════════ [[nodiscard]] inline ExprSpec dead_code_elimination(const ExprSpec& spec) { - if (spec.instrs.empty()) - return spec; + if (spec.instrs.empty()) return spec; const std::size_t n = spec.instrs.size(); std::vector live_reg(spec.num_regs, 0); std::vector live_instr(n, 0); @@ -306,10 +317,8 @@ namespace nn for (std::size_t i = n; i-- > 0;) { const ExprInstr& ins = spec.instrs[i]; - if (ins.dst < spec.num_regs && live_reg[ins.dst]) - live_instr[i] = 1; - if (!live_instr[i]) - continue; + if (ins.dst < spec.num_regs && live_reg[ins.dst]) live_instr[i] = 1; + if (!live_instr[i]) continue; const ExprOp op = static_cast(ins.op); const ExprOperand* ops[3] = {&ins.a, &ins.b, &ins.c}; const std::size_t nops = expr_instr_num_operands(op); @@ -325,16 +334,15 @@ namespace nn } ExprSpec out; - out.views = spec.views; - out.consts = spec.consts; + out.views = spec.views; + out.consts = spec.consts; out.rparams = spec.rparams; // RParam 值序列：非结构、跨 pass 原样透传 out.num_regs = spec.num_regs; - out.matmul = spec.matmul; - out.extras = spec.extras; // P2 多输出：跨 pass 原样透传（寄存器号未变） + out.matmul = spec.matmul; + out.extras = spec.extras; // P2 多输出：跨 pass 原样透传（寄存器号未变） out.instrs.reserve(n); for (std::size_t i = 0; i < n; ++i) - if (live_instr[i]) - out.instrs.push_back(spec.instrs[i]); + if (live_instr[i]) out.instrs.push_back(spec.instrs[i]); return out; } @@ -347,7 +355,7 @@ namespace nn [[nodiscard]] inline ExprSpec renumber_registers(const ExprSpec& spec) { ExprSpec out; - out.views = spec.views; + out.views = spec.views; out.rparams = spec.rparams; // RParam 值序列：跨 pass 原样透传 out.matmul = spec.matmul; @@ -397,8 +405,7 @@ namespace nn return o; }; // 只重映射实际使用的操作数（未用 b/c 是默认哨兵，不得当 Reg(0) 重编号） - const std::size_t nops = - expr_instr_num_operands(static_cast(in.op)); + const std::size_t nops = expr_instr_num_operands(static_cast(in.op)); ni.a = remap(in.a); if (nops >= 2) ni.b = remap(in.b); if (nops >= 3) ni.c = remap(in.c); @@ -407,8 +414,8 @@ namespace nn // P2 多输出：extras 随新寄存器号重映射 out.extras.reserve(spec.extras.size()); for (const auto e : spec.extras) - out.extras.push_back((e < spec.num_regs && new_reg[e] >= 0) - ? static_cast(new_reg[e]) : e); + out.extras.push_back( + (e < spec.num_regs && new_reg[e] >= 0) ? static_cast(new_reg[e]) : e); out.num_regs = next; return out; } @@ -424,7 +431,7 @@ namespace nn [[nodiscard]] inline ExprSpec common_subexpression_elimination(const ExprSpec& spec) { ExprSpec out; - out.views = spec.views; + out.views = spec.views; out.consts = spec.consts; out.rparams = spec.rparams; // RParam 值序列：跨 pass 原样透传 out.matmul = spec.matmul; @@ -454,19 +461,18 @@ namespace nn return o; }; // 只重映射实际使用的操作数（未用 b/c 是默认哨兵；key 中保持哨兵不变） - const std::size_t nops = - expr_instr_num_operands(static_cast(in.op)); + const std::size_t nops = expr_instr_num_operands(static_cast(in.op)); const ExprOperand a = remap(in.a); const ExprOperand b = (nops >= 2) ? remap(in.b) : in.b; const ExprOperand c = (nops >= 3) ? remap(in.c) : in.c; std::uint64_t key = static_cast(in.op); key |= static_cast(a.kind) << 8; - key |= static_cast(a.idx) << 16; + key |= static_cast(a.idx) << 16; key |= static_cast(b.kind) << 24; - key |= static_cast(b.idx) << 32; + key |= static_cast(b.idx) << 32; key |= static_cast(c.kind) << 40; - key |= static_cast(c.idx) << 48; + key |= static_cast(c.idx) << 48; const auto it = cse.find(key); if (it != cse.end()) @@ -477,14 +483,19 @@ namespace nn const std::uint8_t nd = next++; cse.emplace(key, nd); dst_map[in.dst] = nd; - ExprInstr ni; ni.op = in.op; ni.dst = nd; ni.a = a; ni.b = b; ni.c = c; + ExprInstr ni; + ni.op = in.op; + ni.dst = nd; + ni.a = a; + ni.b = b; + ni.c = c; out.instrs.push_back(ni); } // P2 多输出：extras 随 CSE 复用映射重指向（同一值合并后指向规范 dst） out.extras.reserve(spec.extras.size()); for (const auto e : spec.extras) - out.extras.push_back((e < spec.num_regs && dst_map[e] >= 0) - ? static_cast(dst_map[e]) : e); + out.extras.push_back( + (e < spec.num_regs && dst_map[e] >= 0) ? static_cast(dst_map[e]) : e); out.num_regs = next; return out; } @@ -505,8 +516,7 @@ namespace nn // ═══════════════════════════════════════════════════════════════════════════ [[nodiscard]] inline ExprSpec allocate_registers_liveness(const ExprSpec& spec) { - if (spec.instrs.empty()) - return spec; + if (spec.instrs.empty()) return spec; const std::size_t n = spec.instrs.size(); // def_at[reg] = 定义该寄存器的指令下标 @@ -528,15 +538,13 @@ namespace nn o.kind == static_cast(ExprOperandKind::Fanout) || o.kind == static_cast(ExprOperandKind::Reduce)) && o.idx < spec.num_regs && def_at[o.idx] >= 0) - last_use[def_at[o.idx]] = - std::max(last_use[def_at[o.idx]], static_cast(i)); + last_use[def_at[o.idx]] = std::max(last_use[def_at[o.idx]], static_cast(i)); } } last_use[n - 1] = static_cast(n - 1); // 输出寄存器保护 // P2 多输出：extras 的定义值同样存活到 kernel 末（防被 liveness 复用） for (const auto e : spec.extras) - if (e < spec.num_regs && def_at[e] >= 0) - last_use[def_at[e]] = static_cast(n - 1); + if (e < spec.num_regs && def_at[e] >= 0) last_use[def_at[e]] = static_cast(n - 1); // 分配（确定性贪心；逐元素 liveness 复用，归约独立区段） std::vector reg_of(n, -1); @@ -545,8 +553,7 @@ namespace nn for (std::size_t i = 0; i < n; ++i) { const ExprOp op = static_cast(spec.instrs[i].op); - if (expr_op_is_reduce(op)) - continue; // 归约指令第二遍独立分配 + if (expr_op_is_reduce(op)) continue; // 归约指令第二遍独立分配 // 释放所有 last_use < i 的逐元素寄存器 for (std::size_t r = 0; r < release_pos.size(); ++r) if (release_pos[r] >= 0 && release_pos[r] < static_cast(i) && !free_mark[r]) @@ -554,7 +561,12 @@ namespace nn // 取最小可用 int reg = -1; for (std::size_t r = 0; r < release_pos.size(); ++r) - if (free_mark[r]) { reg = static_cast(r); free_mark[r] = 0; break; } + if (free_mark[r]) + { + reg = static_cast(r); + free_mark[r] = 0; + break; + } if (reg < 0) { reg = static_cast(release_pos.size()); @@ -569,13 +581,12 @@ namespace nn for (std::size_t i = 0; i < n; ++i) { const ExprOp op = static_cast(spec.instrs[i].op); - if (expr_op_is_reduce(op)) - reg_of[i] = static_cast(next_red++); + if (expr_op_is_reduce(op)) reg_of[i] = static_cast(next_red++); } const std::uint32_t num_regs = next_red; ExprSpec out; - out.views = spec.views; + out.views = spec.views; out.consts = spec.consts; out.rparams = spec.rparams; // RParam 值序列：跨 pass 原样透传 out.matmul = spec.matmul; @@ -583,7 +594,8 @@ namespace nn out.extras.reserve(spec.extras.size()); for (const auto e : spec.extras) out.extras.push_back((e < spec.num_regs && def_at[e] >= 0) - ? static_cast(reg_of[def_at[e]]) : e); + ? static_cast(reg_of[def_at[e]]) + : e); out.instrs.reserve(n); for (std::size_t i = 0; i < n; ++i) { @@ -599,8 +611,7 @@ namespace nn return o; }; // 只重映射实际使用的操作数（未用 b/c 是默认哨兵，不得当 Reg(0) 重编号） - const std::size_t nops = - expr_instr_num_operands(static_cast(spec.instrs[i].op)); + const std::size_t nops = expr_instr_num_operands(static_cast(spec.instrs[i].op)); ni.a = remap(ni.a); if (nops >= 2) ni.b = remap(ni.b); if (nops >= 3) ni.c = remap(ni.c); @@ -626,8 +637,7 @@ namespace nn // （canonical 不变式：expr_spec_key(spec) ≡ expr_spec_key(canonical(spec))， // 对 fold spec 退化为恒等——双域 fold 若要在 body 上跑 DCE/CSE，须先给 // 各 pass 补 fold 字段透传再放开，勿直接删本行）。 - if (spec.fold) - return spec; + if (spec.fold) return spec; ExprSpec s = fold_constants_and_algebra(spec); s = dead_code_elimination(s); s = renumber_registers(s); @@ -637,4 +647,3 @@ namespace nn } } // namespace nn - diff --git a/include/neuralnet.cpp/expr_registry.hpp b/include/neuralnet.cpp/expr_registry.hpp index 557b9fb6..6f093fd4 100644 --- a/include/neuralnet.cpp/expr_registry.hpp +++ b/include/neuralnet.cpp/expr_registry.hpp @@ -47,13 +47,13 @@ struct ExprRegistry { struct VariantRef { - std::string key; // 结构 key（canonical 形态） - ExprPrecSig sig; // 观察到的精度签名（≠0、≠ EXPR_PREC_SIG_DISPATCH） + std::string key; // 结构 key（canonical 形态） + ExprPrecSig sig; // 观察到的精度签名（≠0、≠ EXPR_PREC_SIG_DISPATCH） }; - std::vector specs; + std::vector specs; std::unordered_set keys; - std::vector variants; + std::vector variants; std::unordered_set variant_keys; void add(const ExprSpec& s) @@ -63,8 +63,7 @@ struct ExprRegistry // 同一规则见 expr_opt.hpp 的 expr_spec_lookup_key）。 const ExprSpec canon = canonicalize_expr_spec(s); const std::string k = expr_spec_key(canon); - if (keys.insert(k).second) - specs.push_back(canon); + if (keys.insert(k).second) specs.push_back(canon); } // 记录"某结构在某签名下真的被调用过"（扫描期 = dry-run / 模型 pass / @@ -72,13 +71,11 @@ struct ExprRegistry // 不记：前者与 V0 共用条目，后者不是签名。 void add_variant(const ExprSpec& s, ExprPrecSig sig) { - if (sig == 0u || expr_prec_sig_is_dispatch(sig)) - return; + if (sig == 0u || expr_prec_sig_is_dispatch(sig)) return; const ExprSpec canon = canonicalize_expr_spec(s); const std::string k = expr_spec_key(canon); const std::string vk = expr_prec_sig_key(k, sig); - if (variant_keys.insert(vk).second) - variants.push_back(VariantRef{k, sig}); + if (variant_keys.insert(vk).second) variants.push_back(VariantRef{k, sig}); } }; @@ -90,4 +87,3 @@ struct ExprRegistry } } // namespace nn::fused - diff --git a/include/neuralnet.cpp/expr_spec.hpp b/include/neuralnet.cpp/expr_spec.hpp index 33dc242d..83d4269c 100644 --- a/include/neuralnet.cpp/expr_spec.hpp +++ b/include/neuralnet.cpp/expr_spec.hpp @@ -32,27 +32,27 @@ namespace nn enum class ExprOp : uint8_t { // 一元 - Neg = 0, // -x - Exp = 1, // e^x - Log = 2, // ln(x) - Sqrt = 3, // √x + Neg = 0, // -x + Exp = 1, // e^x + Log = 2, // ln(x) + Sqrt = 3, // √x Rsqrt = 4, // 1/√x - Abs = 5, // |x| - Tanh = 6, // tanh(x) + Abs = 5, // |x| + Tanh = 6, // tanh(x) // 二元 - Add = 7, // a + b - Sub = 8, // a - b - Mul = 9, // a * b - Div = 10, // a / b - Max = 11, // max(a, b) - Min = 12, // min(a, b) + Add = 7, // a + b + Sub = 8, // a - b + Mul = 9, // a * b + Div = 10, // a / b + Max = 11, // max(a, b) + Min = 12, // min(a, b) // 比较（输出 1.0 / 0.0 到寄存器，供 Select 使用） - Lt = 13, // a < b - Le = 14, // a <= b - Gt = 15, // a > b - Ge = 16, // a >= b - Eq = 17, // a == b - Ne = 18, // a != b + Lt = 13, // a < b + Le = 14, // a <= b + Gt = 15, // a > b + Ge = 16, // a >= b + Eq = 17, // a == b + Ne = 18, // a != b // 选择：dst = (a != 0) ? b : c Select = 19, // ── 新增：归约指令（dst 为隐式"每行/每列一个标量"的归约向量， @@ -67,23 +67,29 @@ enum class ExprOp : uint8_t // 归约指令：dst 是"每行/每列一个标量"的隐式向量，供 Reduce 操作数广播引用。 [[nodiscard]] inline constexpr bool expr_op_is_reduce(ExprOp op) noexcept { - return op == ExprOp::ColSum || op == ExprOp::ColMax || - op == ExprOp::RowSum || op == ExprOp::RowMax; + return op == ExprOp::ColSum || op == ExprOp::ColMax || op == ExprOp::RowSum || + op == ExprOp::RowMax; } // 该归约指令产出的向量方向：true=按列归约 → (1, cols)，false=按行归约 → (rows, 1) [[nodiscard]] inline constexpr bool expr_op_reduces_cols(ExprOp op) noexcept -{ - return op == ExprOp::ColSum || op == ExprOp::ColMax; -} +{ return op == ExprOp::ColSum || op == ExprOp::ColMax; } // 指令实际使用的操作数个数（未用的 b/c 保持默认 {0,0}，不可当作 Reg(0) 引用） // constexpr：供 expr_op_table.hpp 的算子元数据一致性编译期断言复用（RV-07） [[nodiscard]] inline constexpr std::size_t expr_instr_num_operands(ExprOp op) noexcept { switch (op) { - case ExprOp::Neg: case ExprOp::Exp: case ExprOp::Log: - case ExprOp::Sqrt: case ExprOp::Rsqrt: case ExprOp::Abs: case ExprOp::Tanh: - case ExprOp::ColSum: case ExprOp::ColMax: case ExprOp::RowSum: case ExprOp::RowMax: + case ExprOp::Neg: + case ExprOp::Exp: + case ExprOp::Log: + case ExprOp::Sqrt: + case ExprOp::Rsqrt: + case ExprOp::Abs: + case ExprOp::Tanh: + case ExprOp::ColSum: + case ExprOp::ColMax: + case ExprOp::RowSum: + case ExprOp::RowMax: return 1; case ExprOp::Select: return 3; @@ -95,9 +101,9 @@ enum class ExprOp : uint8_t // ── 操作数 ───────────────────────────────────────────────────────────────── enum class ExprOperandKind : uint8_t { - Reg = 0, // 引用前驱指令的输出寄存器 - Input = 1, // 引用第 idx 个输入 Tensor（按 views[idx] 视图访问） - Const = 2, // 引用常量池第 idx 项 + Reg = 0, // 引用前驱指令的输出寄存器 + Input = 1, // 引用第 idx 个输入 Tensor（按 views[idx] 视图访问） + Const = 2, // 引用常量池第 idx 项 Fanout = 3, // 引用前驱结果寄存器（语义同 Reg，显式表达 fanout 语义） // ── 新增：引用一个"行/列归约结果"，自动按行或按列广播 ── Reduce = 4, // 引用某归约指令 dst（idx 指向归约指令的目标寄存器号） @@ -113,8 +119,8 @@ enum class ExprOperandKind : uint8_t // Batch = 批次下标（batched matmul 段网格的 batch 维；无 matmul 段 = 0） // 用于上三角掩码 select(Col > Row, -inf, 0)、按 batch 取斜率的线性偏置 // -slope*(Col-Row) 等位置相关的键域改写。 - Row = 6, - Col = 7, + Row = 6, + Col = 7, Batch = 8, // 运行时标量参数（RParam，形状无关融合的标量推广）── // 引用 spec.rparams 第 idx 项。与视图参数（RowMod 周期 / RotateHalf @@ -137,7 +143,7 @@ enum class ExprOperandKind : uint8_t struct ExprOperand { uint8_t kind = 0; // ExprOperandKind - uint8_t idx = 0; // 寄存器号 / 输入下标 / 常量池下标 + uint8_t idx = 0; // 寄存器号 / 输入下标 / 常量池下标 friend bool operator==(const ExprOperand&, const ExprOperand&) = default; }; @@ -148,11 +154,11 @@ struct ExprOperand // Tensor 的关键；未来 VM 中视图即廉价索引映射，无需物化。 enum class ExprViewKind : uint8_t { - Linear = 0, // data[r*cols + c]（默认） - RotateHalf = 1, // 按 param(block_rows) 分块，块内前后半行交换； + Linear = 0, // data[r*cols + c]（默认） + RotateHalf = 1, // 按 param(block_rows) 分块，块内前后半行交换； // negate_first_half=true 时，目标行 rl inits; // 标量态初值，size == num_state（进 key） std::vector body; // 键域块指令序列 std::vector finalize; // 向量域收尾指令（veclen=0 时单列执行） // ── fold v2 双域扩展（默认值 = 纯 fold v1 行为）── - std::uint32_t vec_state_len = 0; // 行向量态长度（0=无；输出列数） - std::optional matmul; // 键域内层收缩段（块局部，N=k 全轴） - std::optional vecacc; // 行向量态块更新 + std::uint32_t vec_state_len = 0; // 行向量态长度（0=无；输出列数） + std::optional matmul; // 键域内层收缩段（块局部，N=k 全轴） + std::optional vecacc; // 行向量态块更新 // 行界整块跳过 / tri_skip（上三角类掩码的 codegen 恒等优化；仅注意力 // fold 构造非 Plain 置位；进 key——结构/codegen 分歧点）：生成器把块内 // valid 钳到 min(BLOCK, fold_k-k0, qt+1-k0)，qt = row%m_per（行在 // batched 网格内的位置），k0>qt 的整块空转。被跳过的 j 恰为链内 select // 屏蔽项（-inf/0）→ max 加 -inf、sum 加 0、w=0 时 +0·V=+0 均为恒等 // → 与全量计算逐位一致。CPU 不钳（全量算，等价性同上）。 - bool tri_skip = false; + bool tri_skip = false; friend bool operator==(const FoldSpec&, const FoldSpec&) = default; }; @@ -359,64 +361,58 @@ struct FoldSpec // 源必须全为行标量——否则同一状态被各 kb/各线程覆盖（未定义）。静态单遍 // 传递即可判定（指令序内 src 类别已知）。 // 返回 reg_elem[r]（1=元素类）；状态收到元素源 → unexpected。 -[[nodiscard]] inline Result> expr_fold_classify( - const FoldSpec& f, std::uint32_t num_regs) +[[nodiscard]] inline Result> expr_fold_classify(const FoldSpec& f, + std::uint32_t num_regs) { std::vector is_elem(num_regs, 0); std::vector is_reduce_dst(num_regs, 0); const auto src_elem = [&](const ExprOperand& op) -> int { const auto k = static_cast(op.kind); - if (k == ExprOperandKind::Input) - return 1; + if (k == ExprOperandKind::Input) return 1; // fold v2：依赖网格列（=块内 kb）的都算元素类——Matmul 段（N=全轴列）、 // Col（全局列索引）。Row / Batch 派生自行号（行固定 → 行标量）。 // BatchMod/BatchCol 是**视图 kind**（经 Input 操作数访问，Input 已=1， // 保守偏元素类：掩码链只进权重/归约源、不直接写状态，无碍）。 // VecState 不允许出现在 body（validate 拒）。 - if (k == ExprOperandKind::Matmul || k == ExprOperandKind::Col) - return 1; + if (k == ExprOperandKind::Matmul || k == ExprOperandKind::Col) return 1; if (k == ExprOperandKind::Row || k == ExprOperandKind::Batch || k == ExprOperandKind::Const || k == ExprOperandKind::RParam) return 0; - if (k == ExprOperandKind::VecState) - return -1; + if (k == ExprOperandKind::VecState) return -1; if (k == ExprOperandKind::Reg || k == ExprOperandKind::Fanout) { if (op.idx >= num_regs) return -1; if (is_reduce_dst[op.idx]) return 0; // 块归约 dst = 行标量 return is_elem[op.idx]; } - if (k == ExprOperandKind::Reduce) - return (op.idx < num_regs) ? 0 : -1; // 归约结果 = 行标量 - return -1; // 未知 kind → 域错误 + if (k == ExprOperandKind::Reduce) return (op.idx < num_regs) ? 0 : -1; // 归约结果 = 行标量 + return -1; // 未知 kind → 域错误 }; for (const auto& ins : f.body) { const ExprOp op = static_cast(ins.op); - const std::size_t nops = expr_op_is_reduce(op) ? 1 - : expr_instr_num_operands(op); + const std::size_t nops = expr_op_is_reduce(op) ? 1 : expr_instr_num_operands(op); const ExprOperand* ops[3] = {&ins.a, &ins.b, &ins.c}; int e = 0; for (std::size_t oi = 0; oi < nops; ++oi) { const int se = src_elem(*ops[oi]); - if (se < 0) - NN_FAIL("expr_fold_classify: operand domain error"); + if (se < 0) NN_FAIL("expr_fold_classify: operand domain error"); if (se > e) e = se; } - if (ins.dst >= num_regs) - NN_FAIL("expr_fold_classify: dst out of range"); + if (ins.dst >= num_regs) NN_FAIL("expr_fold_classify: dst out of range"); if (expr_op_is_reduce(op)) { is_reduce_dst[ins.dst] = 1; - is_elem[ins.dst] = 0; // 归约输出恒为行标量 + is_elem[ins.dst] = 0; // 归约输出恒为行标量 } else { if (ins.dst < f.num_state && e != 0) - NN_FAIL("expr_fold_classify: state written from element-class source (kb-dependent overwrite)"); - if (e != 0) is_elem[ins.dst] = 1; // 类别单调升为元素（跨指令保守） + NN_FAIL("expr_fold_classify: state written from element-class source (kb-dependent " + "overwrite)"); + if (e != 0) is_elem[ins.dst] = 1; // 类别单调升为元素（跨指令保守） } } return is_elem; @@ -441,14 +437,14 @@ struct FoldSpec // - 指令 ≤ 64，寄存器 ≤ 16，输入 ≤ 8，常量 ≤ 16 struct ExprSpec { - std::vector instrs; - std::vector views; // 与 inputs 一一对应 - std::vector consts; - std::vector rparams; // 运行时标量参数（不进 key，运行时按实际值填充） - std::uint32_t num_regs = 0; - std::optional matmul; // 前置 matmul 段（可选；缺省=无） - std::optional fold; // 分块状态归约段（可选；缺省=无） - std::vector extras; // P2 额外输出寄存器（槽 1..k；见上） + std::vector instrs; + std::vector views; // 与 inputs 一一对应 + std::vector consts; + std::vector rparams; // 运行时标量参数（不进 key，运行时按实际值填充） + std::uint32_t num_regs = 0; + std::optional matmul; // 前置 matmul 段（可选；缺省=无） + std::optional fold; // 分块状态归约段（可选；缺省=无） + std::vector extras; // P2 额外输出寄存器（槽 1..k；见上） }; // 多输出上限（Adam 金刚石 = 3 extras：m/v/delta，主输出 p） @@ -464,15 +460,14 @@ inline constexpr std::size_t EXPR_MAX_EXTRA_OUT = 4; // k/batch 不进 key，作为 push constant 运行时填充（与 RowMod/RotateHalf // 同一处理）。**双域 fold 自带 matmul 段同语义**（spec.fold // 的段优先于顶层——fold spec 顶层恒无 matmul，两处不冲突）。 -[[nodiscard]] inline std::optional expr_spec_runtime_matmul_k( - const ExprSpec& s) noexcept +[[nodiscard]] inline std::optional +expr_spec_runtime_matmul_k(const ExprSpec& s) noexcept { if (s.matmul) return s.matmul->k; if (s.fold && s.fold->matmul) return s.fold->matmul->k; return std::nullopt; } -[[nodiscard]] inline std::uint32_t expr_spec_runtime_matmul_batch( - const ExprSpec& s) noexcept +[[nodiscard]] inline std::uint32_t expr_spec_runtime_matmul_batch(const ExprSpec& s) noexcept { if (s.matmul) return s.matmul->batch; if (s.fold && s.fold->matmul) return s.fold->matmul->batch; @@ -483,18 +478,15 @@ inline constexpr std::size_t EXPR_MAX_EXTRA_OUT = 4; // （见 generate_glsl_matmul / generate_glsl_reduce 的 PC `mm_trans`）。 // ⚠ **双域 fold 自带的 matmul 段不在此列**：fold 的 transA/transB 仍进 key // （生成期定死，见 expr_spec_key 的 fold 分支），此处只描述顶层 matmul 段。 -[[nodiscard]] inline std::uint32_t expr_spec_runtime_matmul_trans( - const ExprSpec& s) noexcept +[[nodiscard]] inline std::uint32_t expr_spec_runtime_matmul_trans(const ExprSpec& s) noexcept { if (!s.matmul) return 0u; return (s.matmul->transA ? 1u : 0u) | (s.matmul->transB ? 2u : 0u); } // fold 收缩轴长度（形状参数，不进 key → push constant 运行时填充） -[[nodiscard]] inline std::optional expr_spec_runtime_fold_k( - const ExprSpec& s) noexcept -{ - return s.fold ? std::optional{s.fold->k} : std::nullopt; -} +[[nodiscard]] inline std::optional +expr_spec_runtime_fold_k(const ExprSpec& s) noexcept +{ return s.fold ? std::optional{s.fold->k} : std::nullopt; } // batched 网格的"每 batch 行数"（M）：rows 为输出总行（batch*M） [[nodiscard]] inline std::size_t expr_spec_rows_per_batch(const ExprSpec& s, std::size_t rows) noexcept @@ -519,21 +511,15 @@ inline constexpr std::size_t EXPR_MAX_EXTRA_OUT = 4; // 分组归约的 R（= 池化窗口面积）同理：运行期循环取代编译期展开链。 return k == ExprViewKind::RowMod || k == ExprViewKind::RotateHalf || k == ExprViewKind::BatchMod || k == ExprViewKind::BatchCol || - k == ExprViewKind::HeadGroup || - k == ExprViewKind::RowAccess || - k == ExprViewKind::GroupedReduceSum || - k == ExprViewKind::GroupedReduceMax || + k == ExprViewKind::HeadGroup || k == ExprViewKind::RowAccess || + k == ExprViewKind::GroupedReduceSum || k == ExprViewKind::GroupedReduceMax || k == ExprViewKind::Transpose; } // 该视图消耗的运行时视图参数槽位数（RowAccess 用 offset+mod 两个 vp 槽） -[[nodiscard]] inline constexpr std::uint32_t expr_view_runtime_param_slots( - ExprViewKind k) noexcept -{ - return (k == ExprViewKind::RowAccess) ? 2u : 1u; -} +[[nodiscard]] inline constexpr std::uint32_t expr_view_runtime_param_slots(ExprViewKind k) noexcept +{ return (k == ExprViewKind::RowAccess) ? 2u : 1u; } // 该 spec 的运行时视图参数个数（= 融合 shader 的 push constant vp 槽位数） -[[nodiscard]] inline std::uint32_t expr_spec_runtime_view_param_count( - const ExprSpec& s) noexcept +[[nodiscard]] inline std::uint32_t expr_spec_runtime_view_param_count(const ExprSpec& s) noexcept { std::uint32_t n = 0; for (const auto& v : s.views) @@ -544,30 +530,24 @@ inline constexpr std::size_t EXPR_MAX_EXTRA_OUT = 4; // ── GQA：该输入的"头分组比"（docs/internal/22 §3.2 P2）──────────────── // HeadGroup 视图的 param 即该操作数的头分组比 n_rep（1 = 恒等/非 GQA）。 // GPU 生成器与 CPU 参考都从这里取——单一真源，避免两侧口径漂移。 -[[nodiscard]] inline std::uint32_t gqa_ratio_of(const ExprSpec& s, - std::size_t view_idx) noexcept +[[nodiscard]] inline std::uint32_t gqa_ratio_of(const ExprSpec& s, std::size_t view_idx) noexcept { - if (view_idx >= s.views.size()) - return 1; - if (static_cast(s.views[view_idx].kind) != ExprViewKind::HeadGroup) - return 1; + if (view_idx >= s.views.size()) return 1; + if (static_cast(s.views[view_idx].kind) != ExprViewKind::HeadGroup) return 1; return s.views[view_idx].param != 0 ? s.views[view_idx].param : 1; } // 按视图顺序提取运行时视图参数（RowMod 周期 / RotateHalf 块大小 / RowAccess // 的 mod+offset），运行时 eval_expr 用它填充融合 shader 的 push constant vp 槽。 -[[nodiscard]] inline std::vector expr_spec_runtime_view_params( - const ExprSpec& s) +[[nodiscard]] inline std::vector expr_spec_runtime_view_params(const ExprSpec& s) { std::vector out; for (const auto& v : s.views) { const auto k = static_cast(v.kind); - if (!expr_view_has_runtime_param(k)) - continue; - out.push_back(v.param); // RowMod/RotateHalf: mod/block；RowAccess: mod - if (k == ExprViewKind::RowAccess) - out.push_back(v.param2); // RowAccess 额外 offset 槽 + if (!expr_view_has_runtime_param(k)) continue; + out.push_back(v.param); // RowMod/RotateHalf: mod/block；RowAccess: mod + if (k == ExprViewKind::RowAccess) out.push_back(v.param2); // RowAccess 额外 offset 槽 } return out; } @@ -579,16 +559,14 @@ inline constexpr std::size_t EXPR_MAX_EXTRA_OUT = 4; int axis = -1; for (const auto& v : s.views) { - if (!expr_view_is_reduce(static_cast(v.kind))) - continue; + if (!expr_view_is_reduce(static_cast(v.kind))) continue; const int a = expr_view_reduces_rows(static_cast(v.kind)) ? 0 : 1; if (axis == -1) axis = a; else if (axis != a) return -2; } for (const auto& in : s.instrs) { - if (!expr_op_is_reduce(static_cast(in.op))) - continue; + if (!expr_op_is_reduce(static_cast(in.op))) continue; const int a = expr_op_reduces_cols(static_cast(in.op)) ? 1 : 0; if (axis == -1) axis = a; else if (axis != a) return -2; @@ -602,10 +580,7 @@ inline constexpr std::size_t EXPR_MAX_EXTRA_OUT = 4; // 复合结构——optimizer 主力）；只发 V0 全 f32（f16/混精度走既有逐 kernel 路径）。 // **生成器（fused_generate）与运行期分组（execute_fused_graph）同源用本函数**。 [[nodiscard]] inline bool expr_spec_batchable(const ExprSpec& s) -{ - return expr_spec_reduce_axis(s) == -1 && !s.fold && !s.matmul && - !s.instrs.empty(); -} +{ return expr_spec_reduce_axis(s) == -1 && !s.fold && !s.matmul && !s.instrs.empty(); } // ── 表达式规格相等比较（GPU AOT 匹配用）──────────────────────────────── // 两个 ExprSpec 相等 ⟺ 指令序列、输入视图、常量池、运行时参数、寄存器数、 @@ -613,20 +588,14 @@ inline constexpr std::size_t EXPR_MAX_EXTRA_OUT = 4; // 用于运行时 eval_expr 判断"该表达式是否有预生成融合 shader"。 [[nodiscard]] inline bool expr_spec_equal(const ExprSpec& a, const ExprSpec& b) { - return a.instrs == b.instrs && a.views == b.views && - a.consts == b.consts && a.rparams == b.rparams && - a.num_regs == b.num_regs && - a.matmul == b.matmul && - a.fold == b.fold && - a.extras == b.extras; + return a.instrs == b.instrs && a.views == b.views && a.consts == b.consts && + a.rparams == b.rparams && a.num_regs == b.num_regs && a.matmul == b.matmul && + a.fold == b.fold && a.extras == b.extras; } // 该 spec 的运行时标量参数个数（= 融合 shader 的 push constant 浮点 p 槽位数） -[[nodiscard]] inline std::uint32_t expr_spec_runtime_param_count( - const ExprSpec& s) noexcept -{ - return static_cast(s.rparams.size()); -} +[[nodiscard]] inline std::uint32_t expr_spec_runtime_param_count(const ExprSpec& s) noexcept +{ return static_cast(s.rparams.size()); } // ── 规范结构 key（AOT 收集/匹配的单一依据）────────────────────────────── // 把 ExprSpec 的**结构**（指令、视图、常量池**个数**、寄存器数；不含输入张量、 @@ -638,7 +607,7 @@ inline constexpr std::size_t EXPR_MAX_EXTRA_OUT = 4; // 逐字段字节级串接（不用 memcpy 整个 struct，避免 padding/平台差异）。 [[nodiscard]] inline std::string expr_spec_key(const ExprSpec& s) { - std::uint64_t h = 0xcbf29ce484222325ull; // FNV-1a 64 offset basis + std::uint64_t h = 0xcbf29ce484222325ull; // FNV-1a 64 offset basis const auto feed = [&h](const void* p, std::size_t n) { const auto* b = static_cast(p); @@ -653,13 +622,14 @@ inline constexpr std::size_t EXPR_MAX_EXTRA_OUT = 4; feed_u32(s.num_regs); // P2 多输出（extras）：寄存器号是结构（决定 shader 的额外输出存储）→ 进 key feed_u32(static_cast(s.extras.size())); - for (const auto& e : s.extras) feed(&e, 1); + for (const auto& e : s.extras) + feed(&e, 1); feed_u32(static_cast(s.instrs.size())); for (const auto& in : s.instrs) { feed(&in.op, 1); feed(&in.dst, 1); - feed(&in.a, sizeof(in.a)); // ExprOperand: 2 字节 POD + feed(&in.a, sizeof(in.a)); // ExprOperand: 2 字节 POD feed(&in.b, sizeof(in.b)); feed(&in.c, sizeof(in.c)); } @@ -750,8 +720,7 @@ inline constexpr std::size_t EXPR_MAX_EXTRA_OUT = 4; } char buf[17]; - std::snprintf(buf, sizeof(buf), "%016llx", - static_cast(h)); + std::snprintf(buf, sizeof(buf), "%016llx", static_cast(h)); return std::string(buf); } @@ -766,16 +735,13 @@ inline constexpr std::size_t EXPR_MAX_EXTRA_OUT = 4; // bit 0..15 第 i 个输入是 F16（i = 输入槽位，与 views 顺序一一对应） // bit 16 输出是 F16 // 全 0 = 全 f32：key 不加后缀、registry 不生成额外变体。 -inline constexpr std::uint32_t EXPR_PREC_SIG_OUT_BIT = 16u; +inline constexpr std::uint32_t EXPR_PREC_SIG_OUT_BIT = 16u; inline constexpr std::uint32_t EXPR_PREC_SIG_INPUT_MASK = 0xFFFFu; using ExprPrecSig = std::uint32_t; -[[nodiscard]] inline constexpr ExprPrecSig expr_prec_sig_make( - std::uint32_t input_bits, bool out_f16) noexcept -{ - return (out_f16 ? (1u << EXPR_PREC_SIG_OUT_BIT) : 0u) - | (input_bits & EXPR_PREC_SIG_INPUT_MASK); -} +[[nodiscard]] inline constexpr ExprPrecSig expr_prec_sig_make(std::uint32_t input_bits, + bool out_f16) noexcept +{ return (out_f16 ? (1u << EXPR_PREC_SIG_OUT_BIT) : 0u) | (input_bits & EXPR_PREC_SIG_INPUT_MASK); } // 是否全 f32（sig == 0） [[nodiscard]] inline constexpr bool expr_prec_sig_is_f32(ExprPrecSig s) noexcept @@ -849,12 +815,8 @@ inline constexpr const char* EXPR_MM_KSPLIT_REDUCE_SUFFIX = "#kr"; // ⚠ in_f16 必须先与 EXPR_PREC_SIG_INPUT_MASK 相与：EXPR_PREC_SIG_DISPATCH // （0x00020000）是**哨兵**而非位图，不对齐输入字段就会把 bit 17 读成 // "第 17 个输入是 f16"（EXPR_MAX_INPUTS=16 时侥幸为假，改上限即踩雷）。 -[[nodiscard]] inline constexpr bool expr_prec_sig_in_f16(ExprPrecSig s, - std::size_t i) noexcept -{ - return i < EXPR_PREC_SIG_OUT_BIT && - ((s & EXPR_PREC_SIG_INPUT_MASK) >> i & 1u) != 0u; -} +[[nodiscard]] inline constexpr bool expr_prec_sig_in_f16(ExprPrecSig s, std::size_t i) noexcept +{ return i < EXPR_PREC_SIG_OUT_BIT && ((s & EXPR_PREC_SIG_INPUT_MASK) >> i & 1u) != 0u; } [[nodiscard]] inline constexpr bool expr_prec_sig_out_f16(ExprPrecSig s) noexcept { return ((s >> EXPR_PREC_SIG_OUT_BIT) & 1u) != 0u; } @@ -865,25 +827,21 @@ inline constexpr const char* EXPR_MM_KSPLIT_REDUCE_SUFFIX = "#kr"; [[nodiscard]] inline constexpr bool expr_variant_out_f16(ExprPrecSig variant_sig, ExprPrecSig runtime_sig) noexcept { - return expr_prec_sig_is_dispatch(variant_sig) - ? expr_prec_sig_out_f16(runtime_sig) - : expr_prec_sig_out_f16(variant_sig); + return expr_prec_sig_is_dispatch(variant_sig) ? expr_prec_sig_out_f16(runtime_sig) + : expr_prec_sig_out_f16(variant_sig); } // 变体索引 key：全 f32 → 结构 key 本身；否则加 "#xxxx" -[[nodiscard]] inline std::string expr_prec_sig_key(const std::string& spec_key, - ExprPrecSig sig) +[[nodiscard]] inline std::string expr_prec_sig_key(const std::string& spec_key, ExprPrecSig sig) { - if (sig == 0u) - return spec_key; + if (sig == 0u) return spec_key; char buf[16]; std::snprintf(buf, sizeof(buf), "#%04x", static_cast(sig)); return spec_key + buf; } // ── 诊断用：把签名渲染成 "in=[f16,f32,...] out=f16" -[[nodiscard]] inline std::string expr_prec_sig_str(ExprPrecSig sig, - std::size_t num_inputs) +[[nodiscard]] inline std::string expr_prec_sig_str(ExprPrecSig sig, std::size_t num_inputs) { std::string s = "in=["; for (std::size_t i = 0; i < num_inputs; ++i) @@ -903,14 +861,11 @@ inline constexpr const char* EXPR_MM_KSPLIT_REDUCE_SUFFIX = "#kr"; // 检查：结构内每个输入都是 f16 且输出 f16。含 fold / matmul / 归约 / rparams // 的结构同样适用（这正是 `#a` 覆盖不到、而 `#s` 要覆盖的部分）——数值语义与 // `#x` 的 f16 分支逐字相同，所以没有任何额外的结构禁忌。 -[[nodiscard]] inline bool expr_prec_sig_all_f16(const ExprSpec& spec, - ExprPrecSig sig) noexcept +[[nodiscard]] inline bool expr_prec_sig_all_f16(const ExprSpec& spec, ExprPrecSig sig) noexcept { - if (sig == 0u || expr_prec_sig_is_dispatch(sig) || !expr_prec_sig_out_f16(sig)) - return false; + if (sig == 0u || expr_prec_sig_is_dispatch(sig) || !expr_prec_sig_out_f16(sig)) return false; for (std::size_t i = 0; i < spec.views.size(); ++i) - if (!expr_prec_sig_in_f16(sig, i)) - return false; + if (!expr_prec_sig_in_f16(sig, i)) return false; return true; } @@ -942,31 +897,24 @@ inline constexpr const char* EXPR_MM_KSPLIT_REDUCE_SUFFIX = "#kr"; // 设备要求：shaderFloat16（VulkanDevice 已查询 + 启用，NN_VULKAN_NO_16BIT_ALU // 逃生阀可关）。不支持 → 后端跳过该变体 pipeline → supports_expr_precision_ // variant=false → 基类 NVI 入口回退边界 cast（正确性不变，拿不到原生收益）。 -[[nodiscard]] inline bool expr_prec_sig_native16(const ExprSpec& spec, - ExprPrecSig sig) +[[nodiscard]] inline bool expr_prec_sig_native16(const ExprSpec& spec, ExprPrecSig sig) { - if (sig == 0u || !expr_prec_sig_out_f16(sig)) - return false; - if (spec.fold || spec.matmul) - return false; - if (expr_spec_reduce_axis(spec) != -1) - return false; - if (!spec.rparams.empty()) - return false; + if (sig == 0u || !expr_prec_sig_out_f16(sig)) return false; + if (spec.fold || spec.matmul) return false; + if (expr_spec_reduce_axis(spec) != -1) return false; + if (!spec.rparams.empty()) return false; for (std::size_t i = 0; i < spec.views.size(); ++i) - if (!expr_prec_sig_in_f16(sig, i)) - return false; + if (!expr_prec_sig_in_f16(sig, i)) return false; for (const auto& c : spec.consts) { const float a = std::fabs(c); - if (a != 0.0f && a < 5.9e-8f) // f16 最小次正规边界下方 → 会变 0 + if (a != 0.0f && a < 5.9e-8f) // f16 最小次正规边界下方 → 会变 0 return false; } for (const auto& ins : spec.instrs) { const ExprOperand* ops[3] = {&ins.a, &ins.b, &ins.c}; - const std::size_t nops = - expr_instr_num_operands(static_cast(ins.op)); + const std::size_t nops = expr_instr_num_operands(static_cast(ins.op)); for (std::size_t o = 0; o < nops; ++o) { const auto k = ops[o]->kind; @@ -980,18 +928,20 @@ inline constexpr const char* EXPR_MM_KSPLIT_REDUCE_SUFFIX = "#kr"; } // ── 上限（GPU 资源 / 校验共用）─────────────────────────────────────────── -inline constexpr std::size_t EXPR_MAX_INPUTS = 16; // 树型 DSL 重复叶子上限（绑定按 spec 实际 views 动态创建，非固定） +inline constexpr std::size_t EXPR_MAX_INPUTS = + 16; // 树型 DSL 重复叶子上限（绑定按 spec 实际 views 动态创建，非固定） inline constexpr std::size_t EXPR_MAX_CONSTS = 16; -inline constexpr std::size_t EXPR_MAX_REGS = 32; // 全融合分支 select 表达式（如 SwiGLU backward）所需；shader 仅声明实际 num_regs +inline constexpr std::size_t EXPR_MAX_REGS = + 32; // 全融合分支 select 表达式（如 SwiGLU backward）所需；shader 仅声明实际 num_regs inline constexpr std::size_t EXPR_MAX_INSTRS = 64; // fold 段上限（body/finalize 独立计数，不占顶层 instrs 预算—— // 两者在 GPU 上分别生成块循环体与收尾段，资源约束与主链解耦） -inline constexpr std::size_t FOLD_MAX_BODY = 48; +inline constexpr std::size_t FOLD_MAX_BODY = 48; inline constexpr std::size_t FOLD_MAX_FINALIZE = 16; -inline constexpr std::size_t FOLD_MAX_STATE = 8; -inline constexpr std::size_t FOLD_MAX_VEC = 1024; // 行向量态长度上限（=输出列数） -inline constexpr std::uint32_t FOLD_MAX_MMK = 1024; // fold mm 段内层 k 上限 - // （A 操作数行预载 shared Ash[1024] 的编译期尺寸——d_k 超限在 validate 静态拒） +inline constexpr std::size_t FOLD_MAX_STATE = 8; +inline constexpr std::size_t FOLD_MAX_VEC = 1024; // 行向量态长度上限（=输出列数） +inline constexpr std::uint32_t FOLD_MAX_MMK = 1024; // fold mm 段内层 k 上限 +// （A 操作数行预载 shared Ash[1024] 的编译期尺寸——d_k 超限在 validate 静态拒） // fold 块大小：CPU 执行器与 GPU 生成器**共用**的常量（不进 key——分块是 // 实现细节，但两侧必须同值以对齐分块边界与 max 类逐位；sum 类 GPU subgroup // 蝶形结合序异于 CPU 串行 → 对拍仍走小容差（1e-4~1e-6），并非全逐位）。 @@ -1016,8 +966,8 @@ inline constexpr std::uint32_t EXPR_FOLD_ROWS_PER_WG = 2; // - K 方向每块 32（BK=32，barrier 减半） // 后端 dispatch 按输出块缩放：(ceil(cols/BLOCK), ceil(rows/BLOCK), 1)， // 与生成器必须保持一致。 -inline constexpr std::uint32_t EXPR_MATMUL_TILE = 16; -inline constexpr std::uint32_t EXPR_MATMUL_BLOCK = 64; +inline constexpr std::uint32_t EXPR_MATMUL_TILE = 16; +inline constexpr std::uint32_t EXPR_MATMUL_BLOCK = 64; // ── matmul 细 tile 变体（P0 wgrad 网格饥饿，键后缀 EXPR_MM_FINE_SUFFIX）── // 反传 wgrad 形状（M,N = 64..256 × K=16384）在 64-tile 下只派 // ceil(N/64)×ceil(M/64) = 1~4 个工作组，40HX（34 SM）大面积空转 @@ -1031,7 +981,7 @@ inline constexpr std::uint32_t EXPR_MATMUL_BLOCK = 64; // warp 工作组，169 GFLOPS 比不修还差；R=1（16×16 线程）305~439 也不如 // R=2 预期。逐输出元素的 K 累加序与 64-tile 一致（逐 k 标量升序 FMA） // → 结果逐位一致，字节锚不移动。 -inline constexpr std::uint32_t EXPR_MATMUL_TILE_FINE = 8; +inline constexpr std::uint32_t EXPR_MATMUL_TILE_FINE = 8; inline constexpr std::uint32_t EXPR_MATMUL_BLOCK_FINE = 16; // 饥饿判据：64-tile 网格的工作组数（含 batch 维）低于该阈值 → 派细 tile // 变体。32 ≈ 40HX 的 SM 数（34）——低于它基本必有一波 SM 空转；高于它 @@ -1050,30 +1000,24 @@ inline constexpr std::uint32_t EXPR_MATMUL_FINE_MIN_WGS = 32; // 运行期可用 NN_MM_KSPLIT_S= 强制 S（A/B 调参）、NN_MM_NO_KSPLIT=1 // 关闭（回落细 tile 路径）。 inline constexpr std::uint32_t EXPR_MATMUL_KSPLIT_MIN_WGS = 256; -inline constexpr std::uint32_t EXPR_MATMUL_KSPLIT_MAX_S = 8; +inline constexpr std::uint32_t EXPR_MATMUL_KSPLIT_MAX_S = 8; // ── 校验 ────────────────────────────────────────────────────────────────── // 返回 unexpected(Error) 描述首个非法点。Layer 侧可在提交前调用以尽早报错。 -[[nodiscard]] inline Result validate_expr_spec(const ExprSpec& spec, - std::size_t num_inputs) +[[nodiscard]] inline Result validate_expr_spec(const ExprSpec& spec, std::size_t num_inputs) { // matmul 段或 fold 段存在时允许空指令表（输出 = 段结果本身； // fold 的 finalize 序列即其尾链） if (spec.instrs.empty() && !spec.matmul && !spec.fold) NN_FAIL("validate_expr_spec: empty instruction list"); - if (spec.instrs.size() > EXPR_MAX_INSTRS) - NN_FAIL("validate_expr_spec: too many instructions"); - if (spec.num_regs > EXPR_MAX_REGS) - NN_FAIL("validate_expr_spec: too many registers"); - if (num_inputs > EXPR_MAX_INPUTS) - NN_FAIL("validate_expr_spec: too many inputs"); - if (spec.consts.size() > EXPR_MAX_CONSTS) - NN_FAIL("validate_expr_spec: too many constants"); + if (spec.instrs.size() > EXPR_MAX_INSTRS) NN_FAIL("validate_expr_spec: too many instructions"); + if (spec.num_regs > EXPR_MAX_REGS) NN_FAIL("validate_expr_spec: too many registers"); + if (num_inputs > EXPR_MAX_INPUTS) NN_FAIL("validate_expr_spec: too many inputs"); + if (spec.consts.size() > EXPR_MAX_CONSTS) NN_FAIL("validate_expr_spec: too many constants"); // 运行时标量参数同样受限（PC 浮点槽位有限） if (spec.rparams.size() > EXPR_MAX_CONSTS) NN_FAIL("validate_expr_spec: too many runtime params"); - if (spec.views.size() != num_inputs) - NN_FAIL("validate_expr_spec: views count != inputs count"); + if (spec.views.size() != num_inputs) NN_FAIL("validate_expr_spec: views count != inputs count"); // P2 多输出（extras，IR-C 写穿物化）：仅纯逐元素 spec；额外输出寄存器须 // 被某指令定义、槽互异、且不与主输出（instrs.back().dst）重复。 if (!spec.extras.empty()) @@ -1095,9 +1039,12 @@ inline constexpr std::uint32_t EXPR_MATMUL_KSPLIT_MAX_S = 8; NN_FAIL("validate_expr_spec: duplicate extra output reg"); bool defined = false; for (const auto& ins : spec.instrs) - if (ins.dst == spec.extras[k]) { defined = true; break; } - if (!defined) - NN_FAIL("validate_expr_spec: extra output reg never defined"); + if (ins.dst == spec.extras[k]) + { + defined = true; + break; + } + if (!defined) NN_FAIL("validate_expr_spec: extra output reg never defined"); } } // matmul 段：A/B 输入下标必须在输入范围内（形状由引擎按实际张量推导） @@ -1105,15 +1052,14 @@ inline constexpr std::uint32_t EXPR_MATMUL_KSPLIT_MAX_S = 8; { if (spec.matmul->a_input >= num_inputs || spec.matmul->b_input >= num_inputs) NN_FAIL("validate_expr_spec: matmul input out of range"); - if (spec.matmul->batch == 0) - NN_FAIL("validate_expr_spec: matmul batch must be > 0"); + if (spec.matmul->batch == 0) NN_FAIL("validate_expr_spec: matmul batch must be > 0"); } // 分组归约视图：组长度 R 必须 > 0（R=0 时 GLSL 展开循环无下界、CPU 除零） for (const auto& v : spec.views) { const auto vk = static_cast(v.kind); - if ((vk == ExprViewKind::GroupedReduceSum || vk == ExprViewKind::GroupedReduceMax) - && v.param == 0) + if ((vk == ExprViewKind::GroupedReduceSum || vk == ExprViewKind::GroupedReduceMax) && + v.param == 0) NN_FAIL("validate_expr_spec: grouped reduce R must be > 0"); } @@ -1128,14 +1074,12 @@ inline constexpr std::uint32_t EXPR_MATMUL_KSPLIT_MAX_S = 8; if (f.num_state == 0 || f.num_state > FOLD_MAX_STATE || f.inits.size() != static_cast(f.num_state)) NN_FAIL("validate_expr_spec: fold state count/init mismatch"); - if (f.k == 0) - NN_FAIL("validate_expr_spec: fold k must be > 0"); + if (f.k == 0) NN_FAIL("validate_expr_spec: fold k must be > 0"); if (f.body.empty() || f.body.size() > FOLD_MAX_BODY) NN_FAIL("validate_expr_spec: fold body size out of range"); if (f.finalize.empty() || f.finalize.size() > FOLD_MAX_FINALIZE) NN_FAIL("validate_expr_spec: fold finalize size out of range"); - if (spec.num_regs < f.num_state) - NN_FAIL("validate_expr_spec: fold num_regs < num_state"); + if (spec.num_regs < f.num_state) NN_FAIL("validate_expr_spec: fold num_regs < num_state"); // 双域字段 if (f.vec_state_len > FOLD_MAX_VEC) NN_FAIL("validate_expr_spec: fold vec_state_len out of range"); @@ -1145,8 +1089,7 @@ inline constexpr std::uint32_t EXPR_MATMUL_KSPLIT_MAX_S = 8; { if (f.matmul->a_input >= num_inputs || f.matmul->b_input >= num_inputs) NN_FAIL("validate_expr_spec: fold matmul input out of range"); - if (f.matmul->batch == 0) - NN_FAIL("validate_expr_spec: fold matmul batch must be > 0"); + if (f.matmul->batch == 0) NN_FAIL("validate_expr_spec: fold matmul batch must be > 0"); if (f.matmul->k > FOLD_MAX_MMK) NN_FAIL("validate_expr_spec: fold matmul k exceeds shared preload cap (1024)"); } @@ -1176,8 +1119,10 @@ inline constexpr std::uint32_t EXPR_MATMUL_KSPLIT_MAX_S = 8; if (k == ExprOperandKind::Matmul && !has_mm) NN_FAIL("validate_expr_spec: fold: Matmul operand without fold matmul segment"); if ((k == ExprOperandKind::Row || k == ExprOperandKind::Col || - k == ExprOperandKind::Batch) && !has_mm) - NN_FAIL("validate_expr_spec: fold: index operands require fold matmul segment (batch 网格语义)"); + k == ExprOperandKind::Batch) && + !has_mm) + NN_FAIL("validate_expr_spec: fold: index operands require fold matmul segment " + "(batch 网格语义)"); // VecState 的 body/finalize 区分不在这里做：本 lambda 两处共用—— // body 侧由 expr_fold_classify 拒（src_elem(VecState) = -1 域错误）， // finalize 侧由下方主链做域校验（veclen/单槽）。 @@ -1195,12 +1140,11 @@ inline constexpr std::uint32_t EXPR_MATMUL_KSPLIT_MAX_S = 8; { const ExprOp op = static_cast(ins.op); if (expr_op_is_reduce(op) && expr_op_reduces_cols(op)) - NN_FAIL("validate_expr_spec: fold body: col-reduce not allowed (block axis is the row-reduce axis)"); - if (ins.dst >= spec.num_regs) - NN_FAIL("validate_expr_spec: fold body dst out of range"); + NN_FAIL("validate_expr_spec: fold body: col-reduce not allowed (block axis is the " + "row-reduce axis)"); + if (ins.dst >= spec.num_regs) NN_FAIL("validate_expr_spec: fold body dst out of range"); const ExprOperand* ops[3] = {&ins.a, &ins.b, &ins.c}; - const std::size_t nops = expr_op_is_reduce(op) ? 1 - : expr_instr_num_operands(op); + const std::size_t nops = expr_op_is_reduce(op) ? 1 : expr_instr_num_operands(op); for (std::size_t oi = 0; oi < nops; ++oi) { const ExprOperand& opnd = *ops[oi]; @@ -1211,16 +1155,19 @@ inline constexpr std::uint32_t EXPR_MATMUL_KSPLIT_MAX_S = 8; if (opnd.idx >= spec.num_regs || opnd.idx == ins.dst) NN_FAIL("validate_expr_spec: fold body reduce ref invalid"); if (!body_reduce_dst[opnd.idx]) - NN_FAIL("validate_expr_spec: fold body Reduce operand must reference a prior block-reduce instr"); + NN_FAIL("validate_expr_spec: fold body Reduce operand must reference a " + "prior block-reduce instr"); } else if (k == ExprOperandKind::Reg || k == ExprOperandKind::Fanout) { if (opnd.idx >= spec.num_regs) NN_FAIL("validate_expr_spec: fold body reg out of range"); if (body_reduce_dst[opnd.idx]) - NN_FAIL("validate_expr_spec: fold body reg ref to block-reduce dst (use Reduce operand)"); + NN_FAIL("validate_expr_spec: fold body reg ref to block-reduce dst (use " + "Reduce operand)"); if (!defined[opnd.idx]) - NN_FAIL("validate_expr_spec: fold body use before def (or block-local reg read)"); + NN_FAIL("validate_expr_spec: fold body use before def (or block-local reg " + "read)"); } else if (k == ExprOperandKind::Input) { @@ -1230,15 +1177,14 @@ inline constexpr std::uint32_t EXPR_MATMUL_KSPLIT_MAX_S = 8; if (expr_view_is_reduce(vk)) NN_FAIL("validate_expr_spec: fold body: reduce-view input not allowed"); if (vk == ExprViewKind::Linear || vk == ExprViewKind::HeadGroup || - vk == ExprViewKind::RowMod || - vk == ExprViewKind::RowBroadcast || + vk == ExprViewKind::RowMod || vk == ExprViewKind::RowBroadcast || vk == ExprViewKind::Transpose) - ; // 标量域恒允许；RowBroadcast = 行参数向量（b[row]， - // 无列依赖——doc_col 等行级掩码参数恒可用）； - // Transpose = 转置读（data[col*param+row]，通用索引） - else if ((vk == ExprViewKind::BatchMod || - vk == ExprViewKind::BatchCol) && has_mm) - ; // 双域：掩码钩子消费（ALiBi 斜率 / doc_ids），需 batch 网格 + ; // 标量域恒允许；RowBroadcast = 行参数向量（b[row]， + // 无列依赖——doc_col 等行级掩码参数恒可用）； + // Transpose = 转置读（data[col*param+row]，通用索引） + else if ((vk == ExprViewKind::BatchMod || vk == ExprViewKind::BatchCol) && + has_mm) + ; // 双域：掩码钩子消费（ALiBi 斜率 / doc_ids），需 batch 网格 else NN_FAIL("validate_expr_spec: fold body: view kind outside supported scope"); } @@ -1253,15 +1199,14 @@ inline constexpr std::uint32_t EXPR_MATMUL_KSPLIT_MAX_S = 8; NN_FAIL("validate_expr_spec: fold body rparam out of range"); } } - if (expr_op_is_reduce(op)) - body_reduce_dst[ins.dst] = 1; + if (expr_op_is_reduce(op)) body_reduce_dst[ins.dst] = 1; defined[ins.dst] = 1; } // finalize：只读状态/常量/rparam/Row——禁归约、禁 Input、禁块内临时 // （链内临时 = 各指令 dst，def-before-use 可读，见下） - std::vector fin_def(spec.num_regs, 0); // 本列已写（def-before-use） - std::vector fin_write_max(spec.num_regs, -1);// 任一写的位置（污染判定） + std::vector fin_def(spec.num_regs, 0); // 本列已写（def-before-use） + std::vector fin_write_max(spec.num_regs, -1); // 任一写的位置（污染判定） for (std::size_t fi = 0; fi < f.finalize.size(); ++fi) if (f.finalize[fi].dst < spec.num_regs) fin_write_max[f.finalize[fi].dst] = static_cast(fi); @@ -1287,10 +1232,12 @@ inline constexpr std::uint32_t EXPR_MATMUL_KSPLIT_MAX_S = 8; if (k == ExprOperandKind::VecState) { if (f.vec_state_len == 0) - NN_FAIL("validate_expr_spec: fold finalize: VecState without vec_state_len"); + NN_FAIL( + "validate_expr_spec: fold finalize: VecState without vec_state_len"); if (opnd.idx != 0) - NN_FAIL("validate_expr_spec: fold finalize: VecState multi-slot unsupported"); - continue; // 单槽域校验通过 + NN_FAIL( + "validate_expr_spec: fold finalize: VecState multi-slot unsupported"); + continue; // 单槽域校验通过 } if (k == ExprOperandKind::Reg || k == ExprOperandKind::Fanout) { @@ -1306,10 +1253,12 @@ inline constexpr std::uint32_t EXPR_MATMUL_KSPLIT_MAX_S = 8; // 全错值）。单列（=0）无跨迭代，fold v1 的状态自复制 // finalize（同指令读写状态）合法。 if (opnd.idx >= f.num_state && !fin_def[opnd.idx]) - NN_FAIL("validate_expr_spec: fold finalize reads undefined chain temp (def-before-use)"); - if (opnd.idx < f.num_state && !fin_def[opnd.idx] && - f.vec_state_len > 0 && fin_write_max[opnd.idx] >= 0) - NN_FAIL("validate_expr_spec: fold finalize: written reg also read " "(multi-column loop would read its own prior output)"); + NN_FAIL("validate_expr_spec: fold finalize reads undefined chain temp " + "(def-before-use)"); + if (opnd.idx < f.num_state && !fin_def[opnd.idx] && f.vec_state_len > 0 && + fin_write_max[opnd.idx] >= 0) + NN_FAIL("validate_expr_spec: fold finalize: written reg also read " + "(multi-column loop would read its own prior output)"); } else if (k == ExprOperandKind::Const && opnd.idx >= spec.consts.size()) NN_FAIL("validate_expr_spec: fold finalize const out of range"); @@ -1337,21 +1286,23 @@ inline constexpr std::uint32_t EXPR_MATMUL_KSPLIT_MAX_S = 8; } for (const auto& ins : spec.instrs) { - if (ins.dst >= spec.num_regs) - NN_FAIL("validate_expr_spec: dst reg out of range"); + if (ins.dst >= spec.num_regs) NN_FAIL("validate_expr_spec: dst reg out of range"); const ExprOperand* ops[3] = {&ins.a, &ins.b, &ins.c}; const std::size_t nops = expr_instr_num_operands(static_cast(ins.op)); for (std::size_t oi = 0; oi < nops; ++oi) { const ExprOperand& op = *ops[oi]; if ((op.kind == static_cast(ExprOperandKind::Reg) || - op.kind == static_cast(ExprOperandKind::Fanout)) && op.idx >= spec.num_regs) + op.kind == static_cast(ExprOperandKind::Fanout)) && + op.idx >= spec.num_regs) NN_FAIL("validate_expr_spec: src reg out of range"); if (op.kind == static_cast(ExprOperandKind::Input) && op.idx >= num_inputs) NN_FAIL("validate_expr_spec: input index out of range"); - if (op.kind == static_cast(ExprOperandKind::Const) && op.idx >= spec.consts.size()) + if (op.kind == static_cast(ExprOperandKind::Const) && + op.idx >= spec.consts.size()) NN_FAIL("validate_expr_spec: const index out of range"); - if (op.kind == static_cast(ExprOperandKind::RParam) && op.idx >= spec.rparams.size()) + if (op.kind == static_cast(ExprOperandKind::RParam) && + op.idx >= spec.rparams.size()) NN_FAIL("validate_expr_spec: rparam index out of range"); if (op.kind == static_cast(ExprOperandKind::Matmul) && !spec.matmul) NN_FAIL("validate_expr_spec: Matmul operand without matmul segment"); @@ -1365,8 +1316,8 @@ inline constexpr std::uint32_t EXPR_MATMUL_KSPLIT_MAX_S = 8; // 按行/列广播访问）——拒绝，避免引擎侧产生歧义。 // 3. 归约指令只需源操作数 a，b/c 忽略。 { - std::vector reduce_dst(spec.num_regs, 0); // dst -> 是否为归约结果 - std::vector elem_dst(spec.num_regs, 0); // dst -> 是否为逐元素结果 + std::vector reduce_dst(spec.num_regs, 0); // dst -> 是否为归约结果 + std::vector elem_dst(spec.num_regs, 0); // dst -> 是否为逐元素结果 for (const auto& ins : spec.instrs) { const ExprOp op = static_cast(ins.op); @@ -1374,8 +1325,7 @@ inline constexpr std::uint32_t EXPR_MATMUL_KSPLIT_MAX_S = 8; (is_reduce ? reduce_dst : elem_dst)[ins.dst] = 1; const ExprOperand* ops[3] = {&ins.a, &ins.b, &ins.c}; - const std::size_t nops = is_reduce ? 1 - : expr_instr_num_operands(op); + const std::size_t nops = is_reduce ? 1 : expr_instr_num_operands(op); for (std::size_t oi = 0; oi < nops; ++oi) { const ExprOperand& opnd = *ops[oi]; @@ -1383,10 +1333,10 @@ inline constexpr std::uint32_t EXPR_MATMUL_KSPLIT_MAX_S = 8; { if (opnd.idx >= spec.num_regs) NN_FAIL("validate_expr_spec: reduce ref out of range"); - if (opnd.idx == ins.dst) - NN_FAIL("validate_expr_spec: reduce self-reference"); + if (opnd.idx == ins.dst) NN_FAIL("validate_expr_spec: reduce self-reference"); if (!reduce_dst[opnd.idx] || elem_dst[opnd.idx]) - NN_FAIL("validate_expr_spec: Reduce operand must reference a prior reduce instruction"); + NN_FAIL("validate_expr_spec: Reduce operand must reference a prior reduce " + "instruction"); } else if ((opnd.kind == static_cast(ExprOperandKind::Reg) || opnd.kind == static_cast(ExprOperandKind::Fanout)) && @@ -1403,47 +1353,65 @@ inline constexpr std::uint32_t EXPR_MATMUL_KSPLIT_MAX_S = 8; // ── 便捷构造（避免 Layer 写魔法数）────────────────────────────────────── namespace expr { - inline constexpr ExprOperand reg(std::uint8_t r) { return {0, r}; } - inline constexpr ExprOperand input(std::uint8_t i) { return {1, i}; } - inline constexpr ExprOperand cst(std::uint8_t c) { return {2, c}; } - inline constexpr ExprOperand fanout(std::uint8_t r){ return {3, r}; } - inline constexpr ExprOperand reduce(std::uint8_t r){ return {4, r}; } // 引用归约指令 dst - inline constexpr ExprOperand matmul_op() { return {5, 0}; } // 引用前置 matmul 段输出 - inline constexpr ExprOperand row() { return {6, 0}; } // 当前行号（batch 内） - inline constexpr ExprOperand col() { return {7, 0}; } // 当前列号 - inline constexpr ExprOperand batch() { return {8, 0}; } // 当前批次下标 - // 运行时标量参数（RParam）：运行时按实际 spec.rparams[idx] 填充 - inline constexpr ExprOperand rval(std::uint8_t r) { return {9, r}; } - // fold 行向量态：按向量域当前输出列读单槽向量态 - inline constexpr ExprOperand vec_state(std::uint8_t slot = 0) { return {10, slot}; } - inline constexpr ExprView linear() { return {0, 0, 0}; } - inline constexpr ExprView rotate_half(std::uint32_t block_rows, bool negate_first_half = true) - { return {1, negate_first_half ? std::uint8_t{1} : std::uint8_t{0}, block_rows}; } - inline constexpr ExprView row_mod(std::uint32_t modulo) - { return {2, 0, modulo}; } - // RowAccess(offset, mod)：source_row = offset + (r % mod)；param=mod, param2=offset - inline constexpr ExprView row_access(std::uint32_t offset, std::uint32_t modulo) - { return {12, 0, modulo, offset}; } - inline constexpr ExprView col_reduce_sum() { return {3, 0, 0}; } - inline constexpr ExprView col_reduce_max() { return {4, 0, 0}; } - inline constexpr ExprView row_reduce_sum() { return {5, 0, 0}; } - inline constexpr ExprView row_reduce_max() { return {6, 0, 0}; } - inline constexpr ExprView row_broadcast() { return {7, 0, 0}; } - inline constexpr ExprView col_broadcast() { return {8, 0, 0}; } - // 标签行收集（param = 标签输入槽）与按批次索引/切片 - inline constexpr ExprView row_gather(std::uint8_t label_slot) - { return {9, 0, label_slot}; } - inline constexpr ExprView batch_mod(std::uint32_t modulo) - { return {10, 0, modulo}; } - inline constexpr ExprView batch_col(std::uint32_t per_batch_cols) - { return {11, 0, per_batch_cols}; } - // 转置读（param = 存储行距；运行期 vp 槽、不进 key） - // ⚠ 视图种类**一律用符号名**，不要在这里写数字字面量：本文件合并 - // （2026-10-11）时 HeadGroup / Transpose 的枚举值各占 15 → Transpose 顺延 - // 为 16，写死的 `{15, …}` 会让 trans() 静默构造出 HeadGroup 视图 - // （症状 = CE 的 lse fold 在 CPU 端报 "input shape mismatch"）。 - inline constexpr ExprView trans(std::uint32_t row_stride) - { return {static_cast(ExprViewKind::Transpose), 0, row_stride}; } +inline constexpr ExprOperand reg(std::uint8_t r) +{ return {0, r}; } +inline constexpr ExprOperand input(std::uint8_t i) +{ return {1, i}; } +inline constexpr ExprOperand cst(std::uint8_t c) +{ return {2, c}; } +inline constexpr ExprOperand fanout(std::uint8_t r) +{ return {3, r}; } +inline constexpr ExprOperand reduce(std::uint8_t r) +{ return {4, r}; } // 引用归约指令 dst +inline constexpr ExprOperand matmul_op() +{ return {5, 0}; } // 引用前置 matmul 段输出 +inline constexpr ExprOperand row() +{ return {6, 0}; } // 当前行号（batch 内） +inline constexpr ExprOperand col() +{ return {7, 0}; } // 当前列号 +inline constexpr ExprOperand batch() +{ return {8, 0}; } // 当前批次下标 +// 运行时标量参数（RParam）：运行时按实际 spec.rparams[idx] 填充 +inline constexpr ExprOperand rval(std::uint8_t r) +{ return {9, r}; } +// fold 行向量态：按向量域当前输出列读单槽向量态 +inline constexpr ExprOperand vec_state(std::uint8_t slot = 0) +{ return {10, slot}; } +inline constexpr ExprView linear() +{ return {0, 0, 0}; } +inline constexpr ExprView rotate_half(std::uint32_t block_rows, bool negate_first_half = true) +{ return {1, negate_first_half ? std::uint8_t{1} : std::uint8_t{0}, block_rows}; } +inline constexpr ExprView row_mod(std::uint32_t modulo) +{ return {2, 0, modulo}; } +// RowAccess(offset, mod)：source_row = offset + (r % mod)；param=mod, param2=offset +inline constexpr ExprView row_access(std::uint32_t offset, std::uint32_t modulo) +{ return {12, 0, modulo, offset}; } +inline constexpr ExprView col_reduce_sum() +{ return {3, 0, 0}; } +inline constexpr ExprView col_reduce_max() +{ return {4, 0, 0}; } +inline constexpr ExprView row_reduce_sum() +{ return {5, 0, 0}; } +inline constexpr ExprView row_reduce_max() +{ return {6, 0, 0}; } +inline constexpr ExprView row_broadcast() +{ return {7, 0, 0}; } +inline constexpr ExprView col_broadcast() +{ return {8, 0, 0}; } +// 标签行收集（param = 标签输入槽）与按批次索引/切片 +inline constexpr ExprView row_gather(std::uint8_t label_slot) +{ return {9, 0, label_slot}; } +inline constexpr ExprView batch_mod(std::uint32_t modulo) +{ return {10, 0, modulo}; } +inline constexpr ExprView batch_col(std::uint32_t per_batch_cols) +{ return {11, 0, per_batch_cols}; } +// 转置读（param = 存储行距；运行期 vp 槽、不进 key） +// ⚠ 视图种类**一律用符号名**，不要在这里写数字字面量：本文件合并 +// （2026-10-11）时 HeadGroup / Transpose 的枚举值各占 15 → Transpose 顺延 +// 为 16，写死的 `{15, …}` 会让 trans() 静默构造出 HeadGroup 视图 +// （症状 = CE 的 lse fold 在 CPU 端报 "input shape mismatch"）。 +inline constexpr ExprView trans(std::uint32_t row_stride) +{ return {static_cast(ExprViewKind::Transpose), 0, row_stride}; } } // namespace expr } // namespace nn diff --git a/include/neuralnet.cpp/lr_scheduler.hpp b/include/neuralnet.cpp/lr_scheduler.hpp index c1c5d5e3..69fa3d8b 100644 --- a/include/neuralnet.cpp/lr_scheduler.hpp +++ b/include/neuralnet.cpp/lr_scheduler.hpp @@ -128,7 +128,9 @@ struct LrScheduleConfig int seg, off, this_len; if (seg_len <= 0) { - seg = 0; off = 0; this_len = anneal_span > 0 ? anneal_span : 1; + seg = 0; + off = 0; + this_len = anneal_span > 0 ? anneal_span : 1; } else if (rel >= anneal_span) { @@ -160,8 +162,8 @@ struct LrScheduleConfig } const Scalar progress = static_cast(off) / static_cast(this_len); const Scalar p = progress < 0 ? Scalar{0} : (progress > 1 ? Scalar{1} : progress); - const Scalar annealed = lr_min + 0.5f * (max_lr - lr_min) * - (1.0f + std::cos(3.14159265358979323846f * p)); + const Scalar annealed = + lr_min + 0.5f * (max_lr - lr_min) * (1.0f + std::cos(3.14159265358979323846f * p)); // 最后一段不重启（停在 min_lr）；其余段末重启回 base_lr（下一段起点） if (seg + 1 >= segs) return annealed; // 段内最后一个 epoch（off == this_len-1）即"段末"，已 annealed 到接近 diff --git a/include/neuralnet.cpp/model_container.hpp b/include/neuralnet.cpp/model_container.hpp index f3f6cee1..b48f9a6d 100644 --- a/include/neuralnet.cpp/model_container.hpp +++ b/include/neuralnet.cpp/model_container.hpp @@ -53,8 +53,7 @@ class Model public: Model() = default; explicit Model(ComputeEngine& engine) : engine_(engine.weak_from_this().lock()) {} - explicit Model(std::shared_ptr engine) - : engine_(std::move(engine)) {} + explicit Model(std::shared_ptr engine) : engine_(std::move(engine)) {} // 不可拷贝（unique_ptr 语义），可移动 Model(const Model&) = delete; @@ -68,20 +67,17 @@ class Model // 构造函数 `Model(ComputeEngine&)`（或移动后沿用既有绑定）。 // 是否已绑定引擎（不触发 fail-fast 的只读查询；TrainingSession::create // 等"编排层"在构造前用它做 Result 错误而不是 abort）。 - [[nodiscard]] bool engine_bound() const noexcept - { - return static_cast(engine_); - } + [[nodiscard]] bool engine_bound() const noexcept { return static_cast(engine_); } - [[nodiscard]] ComputeEngine& engine( - std::source_location loc = std::source_location::current()) const + [[nodiscard]] ComputeEngine& + engine(std::source_location loc = std::source_location::current()) const { if (!engine_) { const std::string msg = "[model-lifetime] Model: 未绑定引擎（构造时引擎必须由 shared_ptr 管理）"; - detail::fatal_at("NN_FATAL", loc.file_name(), - static_cast(loc.line()), msg.c_str()); + detail::fatal_at("NN_FATAL", loc.file_name(), static_cast(loc.line()), + msg.c_str()); } return *engine_; } @@ -92,8 +88,7 @@ class Model // 注意：MultiHeadAttention 的构造函数是 protected（设计为基类， // 由 TransformerEncoderLayer / CausalSelfAttention 内部组合）， // 因此不直接通过 add 添加。 - template - Result add(Args&&... args) + template Result add(Args&&... args) { auto layer = std::make_unique(std::forward(args)...); if (default_precision_.has_value()) @@ -130,9 +125,7 @@ class Model // ── 默认精度配置（§9.2）：须在 add()/add_layer() 之前调用 ───────────── void set_default_precision_profile(const PrecisionProfile& profile) - { - default_precision_ = profile; - } + { default_precision_ = profile; } [[nodiscard]] const PrecisionProfile& default_precision_profile() const noexcept { static const PrecisionProfile f32{}; @@ -162,16 +155,13 @@ class Model void set_output_schema(TensorSchema schema) { output_schema_ = std::move(schema); - if (spec_) - spec_->output_schema = output_schema_; + if (spec_) spec_->output_schema = output_schema_; } [[nodiscard]] Result output_schema() const { - if (output_schema_) - return *output_schema_; - if (layers_.empty()) - NN_FAIL("Model::output_schema: no layers"); + if (output_schema_) return *output_schema_; + if (layers_.empty()) NN_FAIL("Model::output_schema: no layers"); return layers_.back()->output_schema(); } @@ -238,16 +228,13 @@ class Model // ── 前向传播：Tensor → Tensor（全程不离开 engine 设备） ─────────────── [[nodiscard]] Result forward(const Tensor& input) { - if (layers_.empty()) - NN_FAIL("Model::forward: no layers"); - if (!engine_) - NN_FAIL("Model::forward: engine 未绑定（构造时引擎必须由 shared_ptr 管理）"); + if (layers_.empty()) NN_FAIL("Model::forward: no layers"); + if (!engine_) NN_FAIL("Model::forward: engine 未绑定（构造时引擎必须由 shared_ptr 管理）"); // GPU 单次推理自动 batch 作用域（issue #28）：外层未持有 batch 时 // 临时开一个，多 Conv 模型不再因逐原语路径不回收帧内分配而 // VK_ERROR_OUT_OF_DEVICE_MEMORY。外层训练循环已持有 batch → 透传。 InferenceScope scope(*engine_); - if (!scope.ok()) - NN_FAIL("Model::forward: " + scope.error().message); + if (!scope.ok()) NN_FAIL("Model::forward: " + scope.error().message); Tensor x = input; for (std::size_t li = 0; li < layers_.size(); ++li) { @@ -262,14 +249,11 @@ class Model // ── 反向传播：Tensor → Tensor（从最后一层到第一层） ─────────────────── [[nodiscard]] Result backward(const Tensor& grad_output) { - if (layers_.empty()) - NN_FAIL("Model::backward: no layers"); - if (!engine_) - NN_FAIL("Model::backward: engine 未绑定（构造时引擎必须由 shared_ptr 管理）"); + if (layers_.empty()) NN_FAIL("Model::backward: no layers"); + if (!engine_) NN_FAIL("Model::backward: engine 未绑定（构造时引擎必须由 shared_ptr 管理）"); // 同 forward（issue #28）：外层未持有 batch 时临时开一个。 InferenceScope scope(*engine_); - if (!scope.ok()) - NN_FAIL("Model::backward: " + scope.error().message); + if (!scope.ok()) NN_FAIL("Model::backward: " + scope.error().message); Tensor g = grad_output; for (std::size_t i = layers_.size(); i-- > 0;) { @@ -284,10 +268,8 @@ class Model // ── 具名 forward/backward：中间层仍走单 Tensor，最后一层可展开多端口 ── [[nodiscard]] Result forward_named(const Tensor& input) { - if (layers_.empty()) - NN_FAIL("Model::forward_named: no layers"); - if (!engine_) - NN_FAIL("Model::forward_named: engine 未绑定"); + if (layers_.empty()) NN_FAIL("Model::forward_named: no layers"); + if (!engine_) NN_FAIL("Model::forward_named: engine 未绑定"); NN_TRY(schema, output_schema()); Tensor x = input; @@ -297,22 +279,17 @@ class Model NN_TRY(out, layers_[li]->forward(x)); x = std::move(*out); } - NN_PROF_SCOPE(nn::prof::interned("fwd_named.layer", - layers_.size() - 1)); + NN_PROF_SCOPE(nn::prof::interned("fwd_named.layer", layers_.size() - 1)); return layers_.back()->forward_named(x, *schema); } - [[nodiscard]] Result backward_named( - const GradientBundle& gradients) + [[nodiscard]] Result backward_named(const GradientBundle& gradients) { - if (layers_.empty()) - NN_FAIL("Model::backward_named: no layers"); - if (!engine_) - NN_FAIL("Model::backward_named: engine 未绑定"); + if (layers_.empty()) NN_FAIL("Model::backward_named: no layers"); + if (!engine_) NN_FAIL("Model::backward_named: engine 未绑定"); NN_TRY(schema, output_schema()); - NN_TRY(last_grad, - layers_.back()->backward_named(gradients, *schema)); + NN_TRY(last_grad, layers_.back()->backward_named(gradients, *schema)); Tensor g = std::move(*last_grad); for (std::size_t i = layers_.size() - 1; i-- > 0;) { diff --git a/include/neuralnet.cpp/model_serialization.hpp b/include/neuralnet.cpp/model_serialization.hpp index 160e6361..9164dff0 100644 --- a/include/neuralnet.cpp/model_serialization.hpp +++ b/include/neuralnet.cpp/model_serialization.hpp @@ -67,7 +67,7 @@ namespace nn // 仍需能被推理入口直接读（裁决第 3 条的"测试推理"用途），故不做独立格式。 // ═══════════════════════════════════════════════════════════════════════════ -inline constexpr uint32_t MODEL_MAGIC = 0x4E4E4E4E; // "NNNN" +inline constexpr uint32_t MODEL_MAGIC = 0x4E4E4E4E; // "NNNN" // v7（docs/internal/23）：**文件角色**（role 字节：Model=推理 / Checkpoint=训练） // + **参数级 kind**（[kind][precision][rows][cols][aux_len][payload_len][aux][payload]）。 // · 角色决定"写哪些参数"：`.nnmodel` 写推理所需全部参数（三值层写打包形态）； @@ -81,7 +81,7 @@ inline constexpr uint32_t MODEL_MAGIC = 0x4E4E4E4E; // "NNNN" // · **整文件校验和 + 尾部完整性标记**：`[sha256 32B]["NNEND" 5B]`（摘要覆盖 // 其前的全部字节，边写边算：见 `detail::HashingStreamBuf`）；截断/损坏的 // 文件在读入前即明确报错（`verify_model_integrity`），不再静默载入错误权重。 -inline constexpr uint32_t MODEL_VERSION = 8; +inline constexpr uint32_t MODEL_VERSION = 8; // ── 尾部完整性标记（v8；见上）──────────────────────────────────────────── inline constexpr std::size_t MODEL_DIGEST_BYTES = 32; // sha256 @@ -109,26 +109,24 @@ namespace detail template requires std::is_trivially_copyable_v -[[nodiscard]] inline Result write_bytes(std::ostream &ofs, const T &v) +[[nodiscard]] inline Result write_bytes(std::ostream& ofs, const T& v) { - if (!nn::write_pod(ofs, v)) - NN_FAIL("Write error"); + if (!nn::write_pod(ofs, v)) NN_FAIL("Write error"); return {}; } template requires std::is_trivially_copyable_v -[[nodiscard]] inline Result read_bytes(std::ifstream &ifs) +[[nodiscard]] inline Result read_bytes(std::ifstream& ifs) { T v{}; if (!nn::read_pod(ifs, v)) { // 诊断增强：带上失败位置（跨版本格式排查时先看"读到哪断了"） const auto pos = ifs.tellg(); - NN_FAIL("Unexpected end of file（读 " + std::to_string(sizeof(T)) + - " 字节失败，流位置 " + - (pos < 0 ? std::string("") - : std::to_string(static_cast(pos))) + "）"); + NN_FAIL("Unexpected end of file（读 " + std::to_string(sizeof(T)) + " 字节失败，流位置 " + + (pos < 0 ? std::string("") : std::to_string(static_cast(pos))) + + "）"); } return v; } @@ -138,11 +136,12 @@ template // variadic 字段批量写入：依次写入每个字段，遇错即停 template -[[nodiscard]] inline Result write_fields(std::ostream &ofs, const Ts &...fields) +[[nodiscard]] inline Result write_fields(std::ostream& ofs, const Ts&... fields) { Result result = {}; bool stop = false; - auto write_one = [&](const auto &field) { + auto write_one = [&](const auto& field) + { if (stop) return; result = write_bytes(ofs, field); if (!result) stop = true; @@ -157,16 +156,14 @@ namespace detail_read { // 递归终止：所有字段已读取 template -[[nodiscard]] inline Result read_fields_rec( - std::ifstream & /*ifs*/, Tuple & /*values*/, std::index_sequence<>) -{ - return {}; -} +[[nodiscard]] inline Result read_fields_rec(std::ifstream& /*ifs*/, Tuple& /*values*/, + std::index_sequence<>) +{ return {}; } // 递归读取第 I 个字段，然后读取剩余字段 template -[[nodiscard]] inline Result read_fields_rec( - std::ifstream &ifs, Tuple &values, std::index_sequence) +[[nodiscard]] inline Result read_fields_rec(std::ifstream& ifs, Tuple& values, + std::index_sequence) { using T = std::tuple_element_t; auto r = read_bytes(ifs); @@ -177,11 +174,11 @@ template } // namespace detail_read template -[[nodiscard]] inline Result> read_fields(std::ifstream &ifs) +[[nodiscard]] inline Result> read_fields(std::ifstream& ifs) { std::tuple values{}; - auto result = detail_read::read_fields_rec>( - ifs, values, std::index_sequence_for{}); + auto result = detail_read::read_fields_rec>(ifs, values, + std::index_sequence_for{}); NN_TRY_CHECK(result); return values; } @@ -195,14 +192,15 @@ template class HashingStreamBuf : public std::streambuf { public: - HashingStreamBuf(std::streambuf &sink, Sha256State &hasher) noexcept - : sink_(sink), hasher_(hasher) {} + HashingStreamBuf(std::streambuf& sink, Sha256State& hasher) noexcept + : sink_(sink), hasher_(hasher) + { + } protected: - std::streamsize xsputn(const char *s, std::streamsize n) override + std::streamsize xsputn(const char* s, std::streamsize n) override { - if (n > 0) - sha256_update(hasher_, std::string_view(s, static_cast(n))); + if (n > 0) sha256_update(hasher_, std::string_view(s, static_cast(n))); return sink_.sputn(s, n); } @@ -215,14 +213,14 @@ class HashingStreamBuf : public std::streambuf } private: - std::streambuf &sink_; - Sha256State &hasher_; + std::streambuf& sink_; + Sha256State& hasher_; }; // ── 矩阵读写（v5：per-tensor precision tag）────────────────────────────── // v5：写入矩阵时附带 precision tag（1B），tag 值见 precision_tag() -[[nodiscard]] inline Result write_matrix_v5(std::ostream &ofs, const Matrix &m, Precision p) +[[nodiscard]] inline Result write_matrix_v5(std::ostream& ofs, const Matrix& m, Precision p) { // 写入 precision tag（§11.3） NN_TRY(r, write_bytes(ofs, precision_tag(p))); @@ -230,24 +228,23 @@ class HashingStreamBuf : public std::streambuf NN_TRY(r2, write_bytes(ofs, static_cast(m.rows()))); NN_TRY(r3, write_bytes(ofs, static_cast(m.cols()))); // 写入数据 - if (!nn::write_pod_span(ofs, m.span())) - NN_FAIL("Write error while writing matrix data"); + if (!nn::write_pod_span(ofs, m.span())) NN_FAIL("Write error while writing matrix data"); return {}; } // v5：写入 f16 矩阵 -[[nodiscard]] inline Result write_matrix_fp16(std::ostream &ofs, const MatrixT &m) +[[nodiscard]] inline Result write_matrix_fp16(std::ostream& ofs, + const MatrixT& m) { NN_TRY(r, write_bytes(ofs, precision_tag(Precision::F16))); NN_TRY(r2, write_bytes(ofs, static_cast(m.rows()))); NN_TRY(r3, write_bytes(ofs, static_cast(m.cols()))); - if (!nn::write_pod_span(ofs, m.span())) - NN_FAIL("Write error while writing f16 matrix data"); + if (!nn::write_pod_span(ofs, m.span())) NN_FAIL("Write error while writing f16 matrix data"); return {}; } // v5：读取矩阵（自动识别 precision tag） -[[nodiscard]] inline Result> read_matrix_v5(std::ifstream &ifs) +[[nodiscard]] inline Result> read_matrix_v5(std::ifstream& ifs) { // 读取 precision tag auto tag_r = read_bytes(ifs); @@ -275,16 +272,15 @@ class HashingStreamBuf : public std::streambuf const std::streampos pos = ifs.tellg(); ifs.seekg(0, std::ios::end); const std::streampos end = ifs.tellg(); - if (pos < 0 || end < pos) - NN_FAIL("Cannot determine file size while reading matrix"); + if (pos < 0 || end < pos) NN_FAIL("Cannot determine file size while reading matrix"); ifs.seekg(pos); - if (!ifs) - NN_FAIL("Stream seek failed while reading matrix"); + if (!ifs) NN_FAIL("Stream seek failed while reading matrix"); const auto remaining = static_cast(end - pos); const std::uint64_t elem_bytes = (p == Precision::F16) ? 2u : 4u; const std::uint64_t max_elems = remaining / elem_bytes; if (rows != 0 && cols != 0 && static_cast(rows) > max_elems / cols) - NN_FAIL("Matrix shape " + std::to_string(rows) + "x" + std::to_string(cols) + " exceeds remaining file bytes; file is corrupt or malicious"); + NN_FAIL("Matrix shape " + std::to_string(rows) + "x" + std::to_string(cols) + + " exceeds remaining file bytes; file is corrupt or malicious"); } if (p == Precision::F16) @@ -309,16 +305,15 @@ class HashingStreamBuf : public std::streambuf } // ── v4 兼容：无 precision tag 的矩阵读写 ───────────────────────────────── -[[nodiscard]] inline Result write_matrix(std::ostream &ofs, const Matrix &m) +[[nodiscard]] inline Result write_matrix(std::ostream& ofs, const Matrix& m) { NN_TRY(r, write_bytes(ofs, static_cast(m.rows()))); NN_TRY(r2, write_bytes(ofs, static_cast(m.cols()))); - if (!nn::write_pod_span(ofs, m.span())) - NN_FAIL("Write error while writing matrix data"); + if (!nn::write_pod_span(ofs, m.span())) NN_FAIL("Write error while writing matrix data"); return {}; } -[[nodiscard]] inline Result read_matrix(std::ifstream &ifs, Matrix &m) +[[nodiscard]] inline Result read_matrix(std::ifstream& ifs, Matrix& m) { auto rows_r = read_bytes(ifs); NN_TRY_CHECK(rows_r); @@ -329,7 +324,9 @@ class HashingStreamBuf : public std::streambuf if (rows != m.rows() || cols != m.cols()) { - NN_FAIL("Matrix shape mismatch: expected (" + std::to_string(m.rows()) + ", " + std::to_string(m.cols()) + "), got (" + std::to_string(rows) + ", " + std::to_string(cols) + ")"); + NN_FAIL("Matrix shape mismatch: expected (" + std::to_string(m.rows()) + ", " + + std::to_string(m.cols()) + "), got (" + std::to_string(rows) + ", " + + std::to_string(cols) + ")"); } if (!nn::read_pod_span(ifs, m.span())) @@ -339,20 +336,18 @@ class HashingStreamBuf : public std::streambuf // ── ModelSpec ↔ KeyValueRecord（自描述规格头） ───────────────────────── -[[nodiscard]] inline KeyValueRecord spec_to_kv(const ModelSpec &spec) +[[nodiscard]] inline KeyValueRecord spec_to_kv(const ModelSpec& spec) { KeyValueRecord kv; // type = 字符串类型名（判别符随 ModelSpecBase 多态走；旧文件的数值 type // 仅在读侧做向下回落，见 spec_from_kv）。 kv.set("type", std::string(spec.type_name())); // 公共字段（RL-0 多输出）：旧文件缺键 = 单输出模型。 - if (spec.output_schema) - kv.set("output_schema", encode_tensor_schema(*spec.output_schema)); + if (spec.output_schema) kv.set("output_schema", encode_tensor_schema(*spec.output_schema)); if (!spec.output_packing.empty()) kv.set("output_packing", encode_tensor_name_list(spec.output_packing)); // 类型特有字段：各子类只写自己的键（键名与既有扁平布局一致）。 - if (spec.arch) - spec.arch->write_kv(kv); + if (spec.arch) spec.arch->write_kv(kv); return kv; } @@ -365,7 +360,7 @@ class HashingStreamBuf : public std::streambuf // 1. spec_to_kv 中写入； // 2. 此处追加一行（引入版本 + 缺失默认值）。 // ══════════════════════════════════════════════════════════════════════ -inline void apply_spec_version_defaults(KeyValueRecord &kv, uint32_t version) +inline void apply_spec_version_defaults(KeyValueRecord& kv, uint32_t version) { // norm_type 于 v4 引入；更早版本缺失时默认 LayerNorm。 // 注：v1/v2/v3 偏移量格式在 read_and_validate_header 已拒绝，此处为语义兜底。 @@ -378,7 +373,7 @@ inline void apply_spec_version_defaults(KeyValueRecord &kv, uint32_t version) // if (version < 5 && !kv.has("max_norm")) kv.set("max_norm", 0); } -[[nodiscard]] inline Result spec_from_kv(const KeyValueRecord &kv) +[[nodiscard]] inline Result spec_from_kv(const KeyValueRecord& kv) { ModelSpec spec; // 缺失字段保持各子类默认值（即版本默认值） @@ -400,15 +395,13 @@ inline void apply_spec_version_defaults(KeyValueRecord &kv, uint32_t version) if (tnum == 6) NN_FAIL("该模型文件为 AttnZip/ZiPT 架构，已整体移除" "（代码保留在 legacy/zipt 分支；恢复前提见 docs/history/history.md）"); - const char *n = legacy_spec_type_name(tnum); - if (!n) - NN_FAIL("模型文件规格 type 数值非法（" + std::to_string(tnum) + "）"); + const char* n = legacy_spec_type_name(tnum); + if (!n) NN_FAIL("模型文件规格 type 数值非法（" + std::to_string(tnum) + "）"); if (std::string(n) == "zipt") NN_FAIL("该模型文件为 AttnZip/ZiPT 架构，已整体移除" "（代码保留在 legacy/zipt 分支；恢复前提见 docs/history/history.md）"); spec.arch = make_spec(n); - if (!spec.arch) - NN_FAIL("模型文件规格 type 数值无法还原（" + std::to_string(tnum) + "）"); + if (!spec.arch) NN_FAIL("模型文件规格 type 数值无法还原（" + std::to_string(tnum) + "）"); } // ── 公共字段（output_schema / output_packing，RL-0）：缺键 → nullopt/空 ── @@ -442,27 +435,22 @@ inline void apply_spec_version_defaults(KeyValueRecord &kv, uint32_t version) } // ── 类型特有字段（各子类读自己的键；缺失键保持默认）────────────────── - if (spec.arch) - { - NN_TRY(x, spec.arch->read_kv(kv)); - } + if (spec.arch) { NN_TRY(x, spec.arch->read_kv(kv)); } - if (!spec.arch) - NN_FAIL("模型文件规格缺少有效的 type 字段"); + if (!spec.arch) NN_FAIL("模型文件规格缺少有效的 type 字段"); return spec; } // ── 规格头读写：长度前缀 + KeyValueRecord（无偏移量假设） ──────────────── -[[nodiscard]] inline Result write_spec_header(std::ostream &ofs, const ModelSpec &spec) +[[nodiscard]] inline Result write_spec_header(std::ostream& ofs, const ModelSpec& spec) { auto bytes = spec_to_kv(spec).serialize(); NN_TRY(r, write_bytes(ofs, static_cast(bytes.size()))); if (!bytes.empty()) { ofs.write(bytes.data(), static_cast(bytes.size())); - if (!ofs) - NN_FAIL("Write error while writing spec header"); + if (!ofs) NN_FAIL("Write error while writing spec header"); } return {}; } @@ -472,7 +460,7 @@ inline void apply_spec_version_defaults(KeyValueRecord &kv, uint32_t version) // 下连 Result 都回不去）。模型 spec / tokenizer JSON 远小于此值。 inline constexpr std::size_t kMaxSerializedStringBytes = 64u * 1024u * 1024u; -[[nodiscard]] inline Result read_spec_header(std::ifstream &ifs, uint32_t version) +[[nodiscard]] inline Result read_spec_header(std::ifstream& ifs, uint32_t version) { auto len_r = read_bytes(ifs); NN_TRY_CHECK(len_r); @@ -483,8 +471,7 @@ inline constexpr std::size_t kMaxSerializedStringBytes = 64u * 1024u * 1024u; if (len > 0) { ifs.read(bytes.data(), static_cast(len)); - if (!ifs) - NN_FAIL("Unexpected end while reading spec header"); + if (!ifs) NN_FAIL("Unexpected end while reading spec header"); } auto kv_r = KeyValueRecord::parse(bytes); NN_TRY_CHECK(kv_r); @@ -495,24 +482,22 @@ inline constexpr std::size_t kMaxSerializedStringBytes = 64u * 1024u * 1024u; // ── 训练状态段读取（v8）：长度前缀 + KeyValueRecord → TrainState ───────── // `state_len == 0` = 无该段（Model 角色，或断点由未接线路径写出）⇒ 默认值， // 续训路径按"冷启动语义"处理。 -[[nodiscard]] inline Result> read_train_state_section( - std::ifstream &ifs, std::uint64_t state_len) +[[nodiscard]] inline Result> +read_train_state_section(std::ifstream& ifs, std::uint64_t state_len) { - if (state_len == 0) - return std::unique_ptr(std::make_unique()); + if (state_len == 0) return std::unique_ptr(std::make_unique()); if (state_len > kMaxSerializedStringBytes) NN_FAIL("训练状态段长度越界（超过 64 MiB），文件可能损坏或恶意"); std::string bytes(static_cast(state_len), '\0'); ifs.read(bytes.data(), static_cast(bytes.size())); - if (!ifs) - NN_FAIL("Unexpected end while reading train state section"); + if (!ifs) NN_FAIL("Unexpected end while reading train state section"); NN_TRY(kv, KeyValueRecord::parse(bytes)); return train_state_from_kv(*kv); } // ── 文件头读写 ────────────────────────────────────────────────────── -[[nodiscard]] inline Result write_header(std::ostream &ofs, ModelRole role) +[[nodiscard]] inline Result write_header(std::ostream& ofs, ModelRole role) { NN_TRY(r, write_bytes(ofs, MODEL_MAGIC)); NN_TRY(r2, write_bytes(ofs, MODEL_VERSION)); @@ -523,37 +508,35 @@ inline constexpr std::size_t kMaxSerializedStringBytes = 64u * 1024u * 1024u; // 返回读到的 version，同时校验 magic number、格式版本和精度 // 仅接受自描述格式 v4+；v1/v2/v3 偏移量定长格式在此处拒绝。 -[[nodiscard]] inline Result read_and_validate_header( - std::ifstream &ifs, ModelRole* role_out = nullptr) +[[nodiscard]] inline Result read_and_validate_header(std::ifstream& ifs, + ModelRole* role_out = nullptr) { auto magic_r = read_bytes(ifs); NN_TRY_CHECK(magic_r); if (*magic_r != MODEL_MAGIC) - NN_FAIL("Invalid model file: bad magic number (0x" - + std::to_string(*magic_r) + ")"); + NN_FAIL("Invalid model file: bad magic number (0x" + std::to_string(*magic_r) + ")"); auto version_r = read_bytes(ifs); NN_TRY_CHECK(version_r); const auto version = *version_r; if (version < 4) - NN_FAIL("模型文件为旧格式 (v" + std::to_string(version) - + ")，已不再支持；请用当前版本重新训练/保存。"); + NN_FAIL("模型文件为旧格式 (v" + std::to_string(version) + + ")，已不再支持；请用当前版本重新训练/保存。"); if (version > MODEL_VERSION) - NN_FAIL("模型文件版本过新 (v" + std::to_string(version) - + " > v" + std::to_string(MODEL_VERSION) + ")，请升级程序。"); + NN_FAIL("模型文件版本过新 (v" + std::to_string(version) + " > v" + + std::to_string(MODEL_VERSION) + ")，请升级程序。"); auto pt_result = read_bytes(ifs); - if (!pt_result) - NN_FAIL("Unexpected end: missing precision tag"); + if (!pt_result) NN_FAIL("Unexpected end: missing precision tag"); uint8_t precision_tag = *pt_result; if (precision_tag != PRECISION_TAG) - NN_FAIL("Precision mismatch: file uses " + std::string(precision_tag == 0 ? "f32" : "f64") + ", but build uses " + std::string(PRECISION_TAG == 0 ? "f32" : "f64")); + NN_FAIL("Precision mismatch: file uses " + std::string(precision_tag == 0 ? "f32" : "f64") + + ", but build uses " + std::string(PRECISION_TAG == 0 ? "f32" : "f64")); // ── v7：文件角色字节（v5/v6 无该字段 → 一律按 Model 处理，零破坏）────── ModelRole role = ModelRole::Model; if (version >= 7) { auto role_r = read_bytes(ifs); - if (!role_r) - NN_FAIL("Unexpected end: missing file role byte (v7)"); + if (!role_r) NN_FAIL("Unexpected end: missing file role byte (v7)"); if (*role_r > static_cast(ModelRole::Checkpoint)) NN_FAIL("模型文件 role 字段非法（" + std::to_string(*role_r) + "）"); role = static_cast(*role_r); @@ -566,19 +549,18 @@ inline constexpr std::size_t kMaxSerializedStringBytes = 64u * 1024u * 1024u; // v5 及以下 = JSON 字符串；v6 起 = .nnvocab 文件镜像字节。 // 本层只搬字节，解释权在调用方（load_tokenizer_from_model_blob）。 -[[nodiscard]] inline Result write_tokenizer(std::ostream &ofs, const std::string &data) +[[nodiscard]] inline Result write_tokenizer(std::ostream& ofs, const std::string& data) { NN_TRY(r, write_bytes(ofs, static_cast(data.size()))); if (!data.empty()) { ofs.write(data.data(), static_cast(data.size())); - if (!ofs) - NN_FAIL("Write error while writing tokenizer"); + if (!ofs) NN_FAIL("Write error while writing tokenizer"); } return {}; } -[[nodiscard]] inline Result read_tokenizer(std::ifstream &ifs) +[[nodiscard]] inline Result read_tokenizer(std::ifstream& ifs) { auto len_r = read_bytes(ifs); NN_TRY_CHECK(len_r); @@ -589,8 +571,7 @@ inline constexpr std::size_t kMaxSerializedStringBytes = 64u * 1024u * 1024u; std::string data(len, '\0'); ifs.read(data.data(), static_cast(len)); - if (!ifs) - NN_FAIL("Unexpected end while reading tokenizer data"); + if (!ifs) NN_FAIL("Unexpected end while reading tokenizer data"); return data; } @@ -607,8 +588,8 @@ inline constexpr std::size_t kMaxSerializedStringBytes = 64u * 1024u * 1024u; static_assert(sizeof(Scalar) == 4, "v7 打包载荷按 f32 字承载（Scalar 必须是 float）"); // 单行：2-bit 容器字 -> 盘上 payload（2-bit / base-3） -inline void encode_packed_row(bool base3, const std::uint32_t *words, std::size_t k, - std::uint8_t *out) noexcept +inline void encode_packed_row(bool base3, const std::uint32_t* words, std::size_t k, + std::uint8_t* out) noexcept { if (base3) { @@ -619,8 +600,8 @@ inline void encode_packed_row(bool base3, const std::uint32_t *words, std::size_ } // 单行：盘上 payload -> 2-bit 容器字（运行期一律 2-bit） -inline void decode_packed_row(bool base3, const std::uint8_t *in, std::size_t k, - std::uint32_t *words) noexcept +inline void decode_packed_row(bool base3, const std::uint8_t* in, std::size_t k, + std::uint32_t* words) noexcept { if (base3) { @@ -642,14 +623,13 @@ struct ParamBlob // 未信任长度防御：条目字节数不得超过文件剩余长度 // （与 read_matrix_v5 同一威胁模型：先校验再分配/跳过） -[[nodiscard]] inline Result ensure_blob_fits_( - std::ifstream &ifs, std::uint64_t aux_len, std::uint64_t payload_len) +[[nodiscard]] inline Result ensure_blob_fits_(std::ifstream& ifs, std::uint64_t aux_len, + std::uint64_t payload_len) { const std::streampos pos = ifs.tellg(); ifs.seekg(0, std::ios::end); const std::streampos end = ifs.tellg(); - if (pos < 0 || end < pos) - NN_FAIL("Cannot determine file size while reading param entry"); + if (pos < 0 || end < pos) NN_FAIL("Cannot determine file size while reading param entry"); ifs.seekg(pos); const auto remain = static_cast(end - pos); if (aux_len > remain || payload_len > remain - aux_len) @@ -657,7 +637,7 @@ struct ParamBlob return {}; } -[[nodiscard]] inline Result write_param_blob(std::ostream &ofs, const ParamBlob &b) +[[nodiscard]] inline Result write_param_blob(std::ostream& ofs, const ParamBlob& b) { NN_TRY(r0, write_bytes(ofs, static_cast(b.kind))); NN_TRY(r1, write_bytes(ofs, precision_tag(b.precision))); @@ -667,20 +647,20 @@ struct ParamBlob NN_TRY(r5, write_bytes(ofs, static_cast(b.payload.size()))); if (!b.aux.empty()) { - ofs.write(reinterpret_cast(b.aux.data()), + ofs.write(reinterpret_cast(b.aux.data()), static_cast(b.aux.size())); if (!ofs) NN_FAIL("Write error while writing param aux"); } if (!b.payload.empty()) { - ofs.write(reinterpret_cast(b.payload.data()), + ofs.write(reinterpret_cast(b.payload.data()), static_cast(b.payload.size())); if (!ofs) NN_FAIL("Write error while writing param payload"); } return {}; } -[[nodiscard]] inline Result read_param_blob(std::ifstream &ifs) +[[nodiscard]] inline Result read_param_blob(std::ifstream& ifs) { auto k_r = read_bytes(ifs); NN_TRY_CHECK(k_r); @@ -706,13 +686,12 @@ struct ParamBlob b.payload.resize(static_cast(*plen_r)); if (!b.aux.empty()) { - ifs.read(reinterpret_cast(b.aux.data()), - static_cast(b.aux.size())); + ifs.read(reinterpret_cast(b.aux.data()), static_cast(b.aux.size())); if (!ifs) NN_FAIL("Unexpected end while reading param aux"); } if (!b.payload.empty()) { - ifs.read(reinterpret_cast(b.payload.data()), + ifs.read(reinterpret_cast(b.payload.data()), static_cast(b.payload.size())); if (!ifs) NN_FAIL("Unexpected end while reading param payload"); } @@ -722,7 +701,7 @@ struct ParamBlob // ── 跳过参数条目（只读条目头，载荷 seek 过去）──────────────────────────── // 读训练状态段（在参数之前）与读参数（状态段在前）时各需跳过对方；把载荷 // 读进内存再丢掉在大模型上是纯浪费（一条 embedding 就几百 MB）。 -[[nodiscard]] inline Result skip_param_blob(std::ifstream &ifs) +[[nodiscard]] inline Result skip_param_blob(std::ifstream& ifs) { auto k_r = read_bytes(ifs); NN_TRY_CHECK(k_r); @@ -738,196 +717,187 @@ struct ParamBlob NN_TRY_CHECK(ensure_blob_fits_(ifs, *alen_r, *plen_r)); const auto jump = static_cast(*alen_r + *plen_r); ifs.seekg(jump, std::ios::cur); - if (!ifs) - NN_FAIL("参数条目跳过失败（文件可能损坏或截断）"); + if (!ifs) NN_FAIL("参数条目跳过失败（文件可能损坏或截断）"); return {}; } } // namespace detail // ── ParamEntry ⇄ ParamBlob（宿主字节编解码；base-3 只在这里出现）───────── -[[nodiscard]] inline Result param_entry_to_blob( - ComputeEngine &engine, const ParamEntry &e, ModelRole role) +[[nodiscard]] inline Result +param_entry_to_blob(ComputeEngine& engine, const ParamEntry& e, ModelRole role) { detail::ParamBlob b; b.rows = e.rows; b.cols = e.cols; switch (e.kind) { - case ParamKind::Dense: - case ParamKind::TernaryLatent: + case ParamKind::Dense: + case ParamKind::TernaryLatent: + { + b.kind = e.kind; + b.precision = e.precision; + auto m32 = engine.to_matrix(e.payload, Precision::F32); + NN_TRY_CHECK(m32); + if (m32->rows() != e.rows || m32->cols() != e.cols) + NN_FAIL("序列化：参数张量形状与条目不一致"); + if (e.precision == Precision::F16) { - b.kind = e.kind; - b.precision = e.precision; - auto m32 = engine.to_matrix(e.payload, Precision::F32); - NN_TRY_CHECK(m32); - if (m32->rows() != e.rows || m32->cols() != e.cols) - NN_FAIL("序列化：参数张量形状与条目不一致"); - if (e.precision == Precision::F16) - { - MatrixT m16(e.rows, e.cols); - const auto s = m32->span(); - auto d = m16.span(); - for (std::size_t i = 0; i < s.size(); ++i) d[i] = s[i]; - const auto bytes = std::span( - reinterpret_cast(m16.span().data()), - m16.span().size() * sizeof(std::uint16_t)); - b.payload.assign(bytes.begin(), bytes.end()); - } - else - { - const auto bytes = std::span( - reinterpret_cast(m32->span().data()), - m32->span().size() * sizeof(Scalar)); - b.payload.assign(bytes.begin(), bytes.end()); - } - return b; + MatrixT m16(e.rows, e.cols); + const auto s = m32->span(); + auto d = m16.span(); + for (std::size_t i = 0; i < s.size(); ++i) + d[i] = s[i]; + const auto bytes = std::span( + reinterpret_cast(m16.span().data()), + m16.span().size() * sizeof(std::uint16_t)); + b.payload.assign(bytes.begin(), bytes.end()); } - case ParamKind::TernaryPacked2: + else { - auto cont = engine.to_matrix(e.payload, Precision::F32); - NN_TRY_CHECK(cont); - auto tau = engine.to_matrix(e.aux, Precision::F32); - NN_TRY_CHECK(tau); - const std::size_t k = e.cols; - const std::size_t rows = e.rows; - const std::size_t wpr = nn::trit_pack::words_per_row(k); - if (cont->rows() != rows || cont->cols() != wpr) - NN_FAIL("序列化：打包容器形状与条目不一致"); - if (tau->rows() != rows || tau->cols() != 1) - NN_FAIL("序列化：tau 形状与条目不一致"); - const auto tsp = tau->span(); - b.aux.resize(tsp.size() * sizeof(Scalar)); - std::memcpy(b.aux.data(), tsp.data(), b.aux.size()); - // 推理角色默认落盘为 base-3（1.6 bpw）；断点角色不会走到这里 - const bool base3 = (role == ModelRole::Model); - b.kind = base3 ? ParamKind::TernaryPacked3 : ParamKind::TernaryPacked2; - b.precision = Precision::F32; - const std::size_t per_row = base3 - ? nn::trit_pack::base3_bytes_per_row(k) - : nn::trit_pack::bytes_per_row(k); - b.payload.resize(rows * per_row); - const auto cs = cont->span(); - std::vector words(wpr); - for (std::size_t r = 0; r < rows; ++r) - { - for (std::size_t i = 0; i < wpr; ++i) - words[i] = std::bit_cast(cs[r * wpr + i]); - detail::encode_packed_row(base3, words.data(), k, - b.payload.data() + r * per_row); - } - return b; + const auto bytes = std::span( + reinterpret_cast(m32->span().data()), + m32->span().size() * sizeof(Scalar)); + b.payload.assign(bytes.begin(), bytes.end()); } - default: - NN_FAIL(std::string("序列化：不支持的参数形态（kind=") + - param_kind_name(e.kind) + "）"); + return b; + } + case ParamKind::TernaryPacked2: + { + auto cont = engine.to_matrix(e.payload, Precision::F32); + NN_TRY_CHECK(cont); + auto tau = engine.to_matrix(e.aux, Precision::F32); + NN_TRY_CHECK(tau); + const std::size_t k = e.cols; + const std::size_t rows = e.rows; + const std::size_t wpr = nn::trit_pack::words_per_row(k); + if (cont->rows() != rows || cont->cols() != wpr) + NN_FAIL("序列化：打包容器形状与条目不一致"); + if (tau->rows() != rows || tau->cols() != 1) NN_FAIL("序列化：tau 形状与条目不一致"); + const auto tsp = tau->span(); + b.aux.resize(tsp.size() * sizeof(Scalar)); + std::memcpy(b.aux.data(), tsp.data(), b.aux.size()); + // 推理角色默认落盘为 base-3（1.6 bpw）；断点角色不会走到这里 + const bool base3 = (role == ModelRole::Model); + b.kind = base3 ? ParamKind::TernaryPacked3 : ParamKind::TernaryPacked2; + b.precision = Precision::F32; + const std::size_t per_row = + base3 ? nn::trit_pack::base3_bytes_per_row(k) : nn::trit_pack::bytes_per_row(k); + b.payload.resize(rows * per_row); + const auto cs = cont->span(); + std::vector words(wpr); + for (std::size_t r = 0; r < rows; ++r) + { + for (std::size_t i = 0; i < wpr; ++i) + words[i] = std::bit_cast(cs[r * wpr + i]); + detail::encode_packed_row(base3, words.data(), k, b.payload.data() + r * per_row); + } + return b; + } + default: + NN_FAIL(std::string("序列化：不支持的参数形态（kind=") + param_kind_name(e.kind) + "）"); } } -[[nodiscard]] inline Result param_blob_to_entry( - ComputeEngine &engine, const detail::ParamBlob &b, - WeightLoadIntent intent = WeightLoadIntent::Inference) +[[nodiscard]] inline Result +param_blob_to_entry(ComputeEngine& engine, const detail::ParamBlob& b, + WeightLoadIntent intent = WeightLoadIntent::Inference) { ParamEntry e; e.rows = static_cast(b.rows); e.cols = static_cast(b.cols); switch (b.kind) { - case ParamKind::Dense: - case ParamKind::TernaryLatent: + case ParamKind::Dense: + case ParamKind::TernaryLatent: + { + Matrix m(e.rows, e.cols); + if (b.precision == Precision::F16) { - Matrix m(e.rows, e.cols); - if (b.precision == Precision::F16) - { - MatrixT m16(e.rows, e.cols); - if (b.payload.size() != m16.span().size() * sizeof(std::uint16_t)) - NN_FAIL("序列化：f16 载荷长度与形状不一致"); - std::memcpy(m16.span().data(), b.payload.data(), b.payload.size()); - const auto s = m16.span(); - auto d = m.span(); - for (std::size_t i = 0; i < s.size(); ++i) - d[i] = static_cast(s[i]); - } - else - { - if (b.payload.size() != m.size() * sizeof(Scalar)) - NN_FAIL("序列化：载荷长度与形状不一致"); - std::memcpy(m.span().data(), b.payload.data(), b.payload.size()); - } - // 载荷恒物化为 f32 张量（f16 升位精确无损）→ import_into_ 负责回写 - auto t = engine.from_matrix(m, Precision::F32); - NN_TRY_CHECK(t); - e.kind = b.kind; - e.precision = b.precision; - e.payload = std::move(*t); - return e; + MatrixT m16(e.rows, e.cols); + if (b.payload.size() != m16.span().size() * sizeof(std::uint16_t)) + NN_FAIL("序列化：f16 载荷长度与形状不一致"); + std::memcpy(m16.span().data(), b.payload.data(), b.payload.size()); + const auto s = m16.span(); + auto d = m.span(); + for (std::size_t i = 0; i < s.size(); ++i) + d[i] = static_cast(s[i]); } - case ParamKind::TernaryPacked2: - case ParamKind::TernaryPacked3: + else { - const std::size_t k = static_cast(b.cols); - const std::size_t rows = static_cast(b.rows); - const std::size_t wpr = nn::trit_pack::words_per_row(k); - const bool base3 = (b.kind == ParamKind::TernaryPacked3); - const std::size_t per_row = base3 - ? nn::trit_pack::base3_bytes_per_row(k) - : nn::trit_pack::bytes_per_row(k); - if (b.payload.size() != rows * per_row) - NN_FAIL("序列化：打包载荷长度与形状不一致"); - Matrix tm(rows, 1); - if (b.aux.size() != tm.size() * sizeof(Scalar)) - NN_FAIL("序列化：tau（aux）长度与形状不一致"); - std::memcpy(tm.span().data(), b.aux.data(), b.aux.size()); - // ── P1-7：Trainable ⇒ 反量化成 f32 latent（warm start / QAT 续训）── - // 载荷 → latent 的转换在**宿主侧**（L3 是宿主字节/Matrix 的唯一落点； - // L2 层不得碰 Matrix）。τ 随条目一起交给层（aux）：层按文件 τ 重建 - // 打包派生物，故首次 forward 的 wq/τ 与文件里的模型逐位一致。 - // 公式与两条不变量的依据见 ternary_pack.hpp 的 `latent_row_scale`。 - if (intent == WeightLoadIntent::Trainable) - { - Matrix lat(rows, k); - std::vector words(wpr); - const auto ts = tm.span(); - auto ls = lat.span(); - for (std::size_t r = 0; r < rows; ++r) - { - detail::decode_packed_row(base3, b.payload.data() + r * per_row, - k, words.data()); - nn::trit_pack::dequant_row(words.data(), k, ts[r], - ls.data() + r * k); - } - auto lt = engine.from_matrix(lat, Precision::F32); - NN_TRY_CHECK(lt); - auto tt = engine.from_matrix(tm, Precision::F32); - NN_TRY_CHECK(tt); - e.kind = ParamKind::TernaryLatent; - e.precision = Precision::F32; - e.payload = std::move(*lt); - e.aux = std::move(*tt); - return e; - } - Matrix cont(rows, wpr); + if (b.payload.size() != m.size() * sizeof(Scalar)) + NN_FAIL("序列化：载荷长度与形状不一致"); + std::memcpy(m.span().data(), b.payload.data(), b.payload.size()); + } + // 载荷恒物化为 f32 张量（f16 升位精确无损）→ import_into_ 负责回写 + auto t = engine.from_matrix(m, Precision::F32); + NN_TRY_CHECK(t); + e.kind = b.kind; + e.precision = b.precision; + e.payload = std::move(*t); + return e; + } + case ParamKind::TernaryPacked2: + case ParamKind::TernaryPacked3: + { + const std::size_t k = static_cast(b.cols); + const std::size_t rows = static_cast(b.rows); + const std::size_t wpr = nn::trit_pack::words_per_row(k); + const bool base3 = (b.kind == ParamKind::TernaryPacked3); + const std::size_t per_row = + base3 ? nn::trit_pack::base3_bytes_per_row(k) : nn::trit_pack::bytes_per_row(k); + if (b.payload.size() != rows * per_row) NN_FAIL("序列化：打包载荷长度与形状不一致"); + Matrix tm(rows, 1); + if (b.aux.size() != tm.size() * sizeof(Scalar)) + NN_FAIL("序列化：tau（aux）长度与形状不一致"); + std::memcpy(tm.span().data(), b.aux.data(), b.aux.size()); + // ── P1-7：Trainable ⇒ 反量化成 f32 latent（warm start / QAT 续训）── + // 载荷 → latent 的转换在**宿主侧**（L3 是宿主字节/Matrix 的唯一落点； + // L2 层不得碰 Matrix）。τ 随条目一起交给层（aux）：层按文件 τ 重建 + // 打包派生物，故首次 forward 的 wq/τ 与文件里的模型逐位一致。 + // 公式与两条不变量的依据见 ternary_pack.hpp 的 `latent_row_scale`。 + if (intent == WeightLoadIntent::Trainable) + { + Matrix lat(rows, k); std::vector words(wpr); - auto cs = cont.span(); + const auto ts = tm.span(); + auto ls = lat.span(); for (std::size_t r = 0; r < rows; ++r) { - detail::decode_packed_row(base3, b.payload.data() + r * per_row, k, - words.data()); - for (std::size_t i = 0; i < wpr; ++i) - cs[r * wpr + i] = std::bit_cast(words[i]); + detail::decode_packed_row(base3, b.payload.data() + r * per_row, k, words.data()); + nn::trit_pack::dequant_row(words.data(), k, ts[r], ls.data() + r * k); } - auto ct = engine.from_matrix(cont, Precision::F32); - NN_TRY_CHECK(ct); - e.kind = ParamKind::TernaryPacked2; // 运行期形态：内核只认 2-bit 容器 - e.payload = std::move(*ct); + auto lt = engine.from_matrix(lat, Precision::F32); + NN_TRY_CHECK(lt); auto tt = engine.from_matrix(tm, Precision::F32); NN_TRY_CHECK(tt); + e.kind = ParamKind::TernaryLatent; + e.precision = Precision::F32; + e.payload = std::move(*lt); e.aux = std::move(*tt); return e; } - default: - NN_FAIL(std::string("序列化：文件里是不支持的参数形态（kind=") + - param_kind_name(b.kind) + "）"); + Matrix cont(rows, wpr); + std::vector words(wpr); + auto cs = cont.span(); + for (std::size_t r = 0; r < rows; ++r) + { + detail::decode_packed_row(base3, b.payload.data() + r * per_row, k, words.data()); + for (std::size_t i = 0; i < wpr; ++i) + cs[r * wpr + i] = std::bit_cast(words[i]); + } + auto ct = engine.from_matrix(cont, Precision::F32); + NN_TRY_CHECK(ct); + e.kind = ParamKind::TernaryPacked2; // 运行期形态：内核只认 2-bit 容器 + e.payload = std::move(*ct); + auto tt = engine.from_matrix(tm, Precision::F32); + NN_TRY_CHECK(tt); + e.aux = std::move(*tt); + return e; + } + default: + NN_FAIL(std::string("序列化：文件里是不支持的参数形态（kind=") + param_kind_name(b.kind) + + "）"); } } @@ -947,31 +917,29 @@ struct ParamBlob // 截断/损坏的文件在解析参数之前就明确报错（旧缺陷 C5：静默载入错误权重）。 // 成本 = 一次顺序读（分块哈希，不把文件读进内存）；调用方是 // load_model / load_train_state，`peek_*` 只读头部故不校验。 -[[nodiscard]] inline Result verify_model_integrity(const std::string &filename) +[[nodiscard]] inline Result verify_model_integrity(const std::string& filename) { std::ifstream ifs(filename, std::ios::binary); - if (!ifs) - NN_FAIL("Cannot open file for reading: " + filename); + if (!ifs) NN_FAIL("Cannot open file for reading: " + filename); ifs.seekg(0, std::ios::end); const std::streampos end = ifs.tellg(); - if (end < 0) - NN_FAIL("Cannot determine file size: " + filename); + if (end < 0) NN_FAIL("Cannot determine file size: " + filename); const auto total = static_cast(end); if (total < MODEL_TAIL_BYTES) NN_FAIL("模型文件过小（" + std::to_string(total) + " 字节 < 尾部标记 " + - std::to_string(MODEL_TAIL_BYTES) + " 字节）：文件被截断或不是模型文件: " + filename); + std::to_string(MODEL_TAIL_BYTES) + + " 字节）：文件被截断或不是模型文件: " + filename); ifs.seekg(static_cast(total - MODEL_TAIL_BYTES), std::ios::beg); - if (!ifs) - NN_FAIL("Seek failed while verifying: " + filename); + if (!ifs) NN_FAIL("Seek failed while verifying: " + filename); std::array tail{}; - ifs.read(reinterpret_cast(tail.data()), - static_cast(MODEL_TAIL_BYTES)); + ifs.read(reinterpret_cast(tail.data()), static_cast(MODEL_TAIL_BYTES)); if (ifs.gcount() != static_cast(MODEL_TAIL_BYTES)) NN_FAIL("尾部完整性标记读取失败（文件被截断）: " + filename); - if (std::string_view(reinterpret_cast(tail.data()) + MODEL_DIGEST_BYTES, + if (std::string_view(reinterpret_cast(tail.data()) + MODEL_DIGEST_BYTES, MODEL_TAIL_MAGIC.size()) != MODEL_TAIL_MAGIC) NN_FAIL("模型文件缺少尾部完整性标记（\"NNEND\"）：文件被截断、损坏，" - "或是旧版本（v7 及以下）文件: " + filename); + "或是旧版本（v7 及以下）文件: " + + filename); std::array stored{}; std::memcpy(stored.data(), tail.data(), MODEL_DIGEST_BYTES); @@ -992,16 +960,18 @@ struct ParamBlob // （`Optimizer::save_state` 的产出顺序即载入时的期望顺序）。 struct CheckpointState { - const TrainState *state = nullptr; // 状态本体（空 = 写 0 长度段） + const TrainState* state = nullptr; // 状态本体（空 = 写 0 长度段） std::span optimizer_blobs{}; // 优化器状态张量（Dense 形态落盘） }; // ── 保存模型（v8：文件角色 + 参数级 kind + 训练状态段 + 尾部校验和）────── // role 决定"写哪些参数"：Model = 推理模型（三值层写打包形态，默认 base-3 载荷）； // Checkpoint = 训练断点（三值层写 f32 latent）+ 可选训练状态段。 -[[nodiscard]] inline Result save_model(const std::string &filename, - Model &model, const ModelSpec &spec, const std::string &tokenizer_data = {}, - ModelRole role = ModelRole::Model, const CheckpointState &ckpt = {}) +[[nodiscard]] inline Result save_model(const std::string& filename, Model& model, + const ModelSpec& spec, + const std::string& tokenizer_data = {}, + ModelRole role = ModelRole::Model, + const CheckpointState& ckpt = {}) { if (ckpt.state != nullptr && role != ModelRole::Checkpoint) NN_FAIL("训练状态段只能写入 Checkpoint 角色（.nnckpt）：推理模型（.nnmodel）" @@ -1010,8 +980,7 @@ struct CheckpointState NN_FAIL("优化器状态张量必须与训练状态一同写入（CheckpointState::state 为空）"); std::ofstream ofs(filename, std::ios::binary); - if (!ofs) - NN_FAIL("Cannot open file for writing: " + filename); + if (!ofs) NN_FAIL("Cannot open file for writing: " + filename); // v8 整文件校验和：所有内容字节经哈希流写入；摘要与尾部标记最后**不经**该流 detail::Sha256State hasher; @@ -1037,17 +1006,15 @@ struct CheckpointState if (state_len > 0) { os.write(state_bytes.data(), static_cast(state_bytes.size())); - if (!os) - NN_FAIL("Write error while writing train state to: " + filename); + if (!os) NN_FAIL("Write error while writing train state to: " + filename); } } auto& engine = model.engine(); // 优化器状态张量（Dense 条目；形状/精度随张量自描述，载入方自行校验） - for (const Tensor &t : ckpt.optimizer_blobs) + for (const Tensor& t : ckpt.optimizer_blobs) { - if (!t.valid()) - NN_FAIL("优化器状态张量非法（未绑定存储）"); + if (!t.valid()) NN_FAIL("优化器状态张量非法（未绑定存储）"); ParamEntry e; e.kind = ParamKind::Dense; e.precision = t.precision(); @@ -1061,7 +1028,7 @@ struct CheckpointState // 逐层逐参数（索引 = 本层参数扁平顺序；`Model::parameters()` 即各层的级联） for (std::size_t li = 0; li < model.num_layers(); ++li) { - Layer &L = model.layer_at(li); + Layer& L = model.layer_at(li); const std::size_t n = L.param_count(); for (std::size_t i = 0; i < n; ++i) { @@ -1087,39 +1054,34 @@ struct CheckpointState // tokenizer 数据（长度前缀，0 = 未嵌入；v5=JSON / v6=.nnvocab，按版本分派） NN_TRY(r4, detail::write_tokenizer(os, tokenizer_data)); - if (!os) - NN_FAIL("Write error while saving model to: " + filename); + if (!os) NN_FAIL("Write error while saving model to: " + filename); // ── 尾部（v8）：整文件校验和 + 完整性标记 ───────────────────────────── // 注意：摘要与标记本身**不算进**校验和 ⇒ 直接写 ofs（绕过哈希流）。 const auto digest = detail::sha256_finalize(hasher); - ofs.write(reinterpret_cast(digest.data()), + ofs.write(reinterpret_cast(digest.data()), static_cast(digest.size())); - ofs.write(MODEL_TAIL_MAGIC.data(), - static_cast(MODEL_TAIL_MAGIC.size())); - if (!ofs) - NN_FAIL("Write error while writing file tail to: " + filename); + ofs.write(MODEL_TAIL_MAGIC.data(), static_cast(MODEL_TAIL_MAGIC.size())); + if (!ofs) NN_FAIL("Write error while writing file tail to: " + filename); return {}; } // ── 读取模型规格（只读头部，不读参数/tokenizer） ────────────────────────── -[[nodiscard]] inline Result peek_model_spec(const std::string &filename) +[[nodiscard]] inline Result peek_model_spec(const std::string& filename) { std::ifstream ifs(filename, std::ios::binary); - if (!ifs) - NN_FAIL("Cannot open file for reading: " + filename); + if (!ifs) NN_FAIL("Cannot open file for reading: " + filename); auto version_r = detail::read_and_validate_header(ifs); NN_TRY_CHECK(version_r); return detail::read_spec_header(ifs, *version_r); } // ── 只读文件角色（v7；v5/v6 无该字段 → Model）──────────────────────────── -[[nodiscard]] inline Result peek_model_role(const std::string &filename, - uint32_t *version_out = nullptr) +[[nodiscard]] inline Result peek_model_role(const std::string& filename, + uint32_t* version_out = nullptr) { std::ifstream ifs(filename, std::ios::binary); - if (!ifs) - NN_FAIL("Cannot open file for reading: " + filename); + if (!ifs) NN_FAIL("Cannot open file for reading: " + filename); ModelRole role = ModelRole::Model; auto v = detail::read_and_validate_header(ifs, &role); NN_TRY_CHECK(v); @@ -1131,24 +1093,22 @@ struct CheckpointState // 供续训路径在建模型**之前**取精度 profile 与计数器、供 CLI/GUI 显示断点状态。 // 没有该段（role=Model，或 v7 及以下 / 未接线的断点）→ 明确报错，调用方据此 // 回落"冷启动语义"（docs/internal/18 §4.2 P1-2）。 -[[nodiscard]] inline Result> peek_train_state( - const std::string &filename) +[[nodiscard]] inline Result> +peek_train_state(const std::string& filename) { std::ifstream ifs(filename, std::ios::binary); - if (!ifs) - NN_FAIL("Cannot open file for reading: " + filename); + if (!ifs) NN_FAIL("Cannot open file for reading: " + filename); auto version_r = detail::read_and_validate_header(ifs); NN_TRY_CHECK(version_r); if (*version_r < 8) - NN_FAIL("训练状态段自 v8 起才有：" + filename + " 是 v" + - std::to_string(*version_r) + "（按冷启动语义处理）"); + NN_FAIL("训练状态段自 v8 起才有：" + filename + " 是 v" + std::to_string(*version_r) + + "（按冷启动语义处理）"); auto spec_r = detail::read_spec_header(ifs, *version_r); NN_TRY_CHECK(spec_r); auto len_r = detail::read_bytes(ifs); NN_TRY_CHECK(len_r); if (*len_r == 0) - NN_FAIL("文件没有训练状态段（推理模型 .nnmodel，或该断点由旧版本写出）: " + - filename); + NN_FAIL("文件没有训练状态段（推理模型 .nnmodel，或该断点由旧版本写出）: " + filename); NN_TRY(state, detail::read_train_state_section(ifs, *len_r)); return *std::move(state); } @@ -1163,25 +1123,23 @@ struct LoadedTrainState std::vector optimizer_blob_precisions; }; -[[nodiscard]] inline Result load_train_state( - ComputeEngine &engine, const std::string &filename) +[[nodiscard]] inline Result load_train_state(ComputeEngine& engine, + const std::string& filename) { std::ifstream ifs(filename, std::ios::binary); - if (!ifs) - NN_FAIL("Cannot open file for reading: " + filename); + if (!ifs) NN_FAIL("Cannot open file for reading: " + filename); auto version_r = detail::read_and_validate_header(ifs); NN_TRY_CHECK(version_r); if (*version_r < 8) - NN_FAIL("训练状态段自 v8 起才有：" + filename + " 是 v" + - std::to_string(*version_r) + "（按冷启动语义处理）"); + NN_FAIL("训练状态段自 v8 起才有：" + filename + " 是 v" + std::to_string(*version_r) + + "（按冷启动语义处理）"); NN_TRY_CHECK(verify_model_integrity(filename)); auto spec_r = detail::read_spec_header(ifs, *version_r); NN_TRY_CHECK(spec_r); auto len_r = detail::read_bytes(ifs); NN_TRY_CHECK(len_r); if (*len_r == 0) - NN_FAIL("文件没有训练状态段（推理模型 .nnmodel，或该断点由旧版本写出）: " + - filename); + NN_FAIL("文件没有训练状态段（推理模型 .nnmodel，或该断点由旧版本写出）: " + filename); LoadedTrainState out; NN_TRY(st, detail::read_train_state_section(ifs, *len_r)); @@ -1210,15 +1168,12 @@ struct LoadedTrainState // intent（P1-7）：三值层的**打包载荷**是"进部署只读态"还是"反量化成可训练 // latent"——推理侧默认 `Inference`（0.25 B/权重）；训练侧的 warm start / // QAT 续训传 `Trainable`（`.nnckpt` 的 TernaryLatent 载荷两种意图同结果）。 -[[nodiscard]] inline Result load_model(const std::string &filename, Model &model, - uint32_t *version_out = nullptr, - ModelRole *role_out = nullptr, - WeightLoadIntent intent = - WeightLoadIntent::Inference) +[[nodiscard]] inline Result +load_model(const std::string& filename, Model& model, uint32_t* version_out = nullptr, + ModelRole* role_out = nullptr, WeightLoadIntent intent = WeightLoadIntent::Inference) { std::ifstream ifs(filename, std::ios::binary); - if (!ifs) - NN_FAIL("Cannot open file for reading: " + filename); + if (!ifs) NN_FAIL("Cannot open file for reading: " + filename); ModelRole role = ModelRole::Model; auto version_r = detail::read_and_validate_header(ifs, &role); NN_TRY_CHECK(version_r); @@ -1227,8 +1182,7 @@ struct LoadedTrainState if (role_out) *role_out = role; // v8：整文件校验和（截断/损坏在读参数之前明确报错） - if (version >= 8) - NN_TRY_CHECK(verify_model_integrity(filename)); + if (version >= 8) NN_TRY_CHECK(verify_model_integrity(filename)); // 读取并校验规格头 auto spec_r = detail::read_spec_header(ifs, version); @@ -1237,7 +1191,10 @@ struct LoadedTrainState { if (!spec_matches(*stored, *spec_r)) { - NN_FAIL("Model architecture mismatch while loading '" + filename + "': " "model expects " + spec_summary(*stored) + " but file contains " + spec_summary(*spec_r)); + NN_FAIL("Model architecture mismatch while loading '" + filename + + "': " + "model expects " + + spec_summary(*stored) + " but file contains " + spec_summary(*spec_r)); } } @@ -1260,7 +1217,7 @@ struct LoadedTrainState // ── v7：参数级 kind（由拥有该参数的那个层决定形态）───────────── for (std::size_t li = 0; li < model.num_layers(); ++li) { - Layer &L = model.layer_at(li); + Layer& L = model.layer_at(li); const std::size_t n = L.param_count(); for (std::size_t i = 0; i < n; ++i) { @@ -1346,20 +1303,17 @@ struct LoadedTrainState tokenizer_data = std::move(*tok_r); } - if (!ifs) - NN_FAIL("Read error while loading model from: " + filename); + if (!ifs) NN_FAIL("Read error while loading model from: " + filename); return tokenizer_data; } // ── 统一入口（v5/v6/v7 同一 magic，自动按版本分流）─────────────────────── -[[nodiscard]] inline Result peek_any_model_spec(const std::string &filename) -{ - return peek_model_spec(filename); -} +[[nodiscard]] inline Result peek_any_model_spec(const std::string& filename) +{ return peek_model_spec(filename); } -[[nodiscard]] inline Result load_any_model(const std::string &filename, Model &model, - std::string *tokenizer_out = nullptr, - ModelRole *role_out = nullptr) +[[nodiscard]] inline Result load_any_model(const std::string& filename, Model& model, + std::string* tokenizer_out = nullptr, + ModelRole* role_out = nullptr) { auto tok_r = load_model(filename, model, nullptr, role_out); NN_TRY_CHECK(tok_r); diff --git a/include/neuralnet.cpp/model_spec.hpp b/include/neuralnet.cpp/model_spec.hpp index 38f3e735..e7616ba9 100644 --- a/include/neuralnet.cpp/model_spec.hpp +++ b/include/neuralnet.cpp/model_spec.hpp @@ -37,27 +37,27 @@ namespace nn // ── 位置编码类型 ───────────────────────────────────────────────────────── enum class PosEncodingType : uint32_t { - Learned = 0, // 可学习位置嵌入（GPT 默认） + Learned = 0, // 可学习位置嵌入（GPT 默认） Sinusoidal = 1, // 正弦波固定位置编码 - ALiBi = 2, // 线性偏置注意力（无位置嵌入） - RoPE = 3, // 旋转位置编码（在注意力 Q/K 上施加，无位置嵌入） + ALiBi = 2, // 线性偏置注意力（无位置嵌入） + RoPE = 3, // 旋转位置编码（在注意力 Q/K 上施加，无位置嵌入） }; // ── FFN 激活类型 ───────────────────────────────────────────────────────── enum class ActivationType : uint32_t { - GeLU = 0, // QuickGeLU（GPT-2 风格，默认） + GeLU = 0, // QuickGeLU（GPT-2 风格，默认） SwiGLU = 1, // SwiGLU（LLaMA/Mistral 风格，每参数效率更高） // 门控平方 ReLU（BitNet b1.58 2B4T 的 `hidden_act = "relu2"`）： // out = relu²(gate(x)) ⊙ up(x) —— 与 SwiGLU 同接线，只换逐元素函数。 - ReLU2 = 2, + ReLU2 = 2, }; // ── 归一化层类型 ───────────────────────────────────────────────────────── enum class NormType : uint32_t { LayerNorm = 0, // LayerNorm（GPT-2 风格；MLP/ViT 默认） - RMSNorm = 1, // RMSNorm（LLaMA/Mistral 风格，更快更稳） + RMSNorm = 1, // RMSNorm（LLaMA/Mistral 风格，更快更稳） BatchNorm = 2, // BatchNorm（沿 batch 维归一化，CNN 默认） }; @@ -77,10 +77,10 @@ enum class NormType : uint32_t // Final = ViT：编码器末端 final norm（全局平均池化之前，原版 ViT 的 ln_f） enum class NormPlace : uint32_t { - None = 0, - Conv = 1, - Head = 2, - Both = 3, + None = 0, + Conv = 1, + Head = 2, + Both = 3, Final = 4, }; @@ -89,11 +89,16 @@ enum class NormPlace : uint32_t { switch (p) { - case NormPlace::Conv: return "conv"; - case NormPlace::Head: return "head"; - case NormPlace::Both: return "both"; - case NormPlace::Final: return "final"; - default: return "none"; + case NormPlace::Conv: + return "conv"; + case NormPlace::Head: + return "head"; + case NormPlace::Both: + return "both"; + case NormPlace::Final: + return "final"; + default: + return "none"; } } @@ -106,32 +111,30 @@ enum class NormPlace : uint32_t // 参数条数/顺序不变 → 无需升 MODEL_VERSION，与 norm_place 的缺键处理同款。 enum class WeightQuant : uint32_t { - None = 0, + None = 0, T1_58 = 1, }; // WeightQuant → CLI/日志/序列化用短名（唯一命名口径；与 Precision::T1_58 同词法） [[nodiscard]] inline const char* weight_quant_name(WeightQuant q) noexcept -{ - return q == WeightQuant::T1_58 ? "t1_58" : "none"; -} +{ return q == WeightQuant::T1_58 ? "t1_58" : "none"; } // ── 旧文件数值 type（仅作向下回落识别；规格判别符 = ModelSpecBase::type_name）── // 规格对象不再持有本枚举：新文件的 `type` 键写**类型名字符串**（如 "GPT"）， // 读到数值（旧文件）时经 `legacy_spec_type_name` 还原为类型名再构造对应子类。 enum class ModelType : uint32_t { - Unknown = 0, - MLP = 1, + Unknown = 0, + MLP = 1, Transformer = 2, - GPT = 3, - ALiBi_GPT = 4, // 使用 ALiBi 的 GPT 模型（向后兼容） - CNN = 5, // 卷积神经网络（LeNet 风格，MNIST） + GPT = 3, + ALiBi_GPT = 4, // 使用 ALiBi 的 GPT 模型（向后兼容） + CNN = 5, // 卷积神经网络（LeNet 风格，MNIST） // 6 = AttnZip / ZiPT —— **已整体移除**（代码保留在 `legacy/zipt` // 分支，恢复前提见 docs/history/history.md「ZiPT 移除」条）。此枚举值仅作占位： // 加载旧模型文件（type=6）时据此给出明确错误，而不是含糊的 "Unknown"。 Reserved_ZiPT = 6, - RAPT = 7, // ReLU 激活线性注意力（ReLU-Linear Attention，causal LM） + RAPT = 7, // ReLU 激活线性注意力（ReLU-Linear Attention，causal LM） }; // 前置声明（`summary` 需要读外壳的公共字段） @@ -203,10 +206,10 @@ class TransformerSpec : public ModelSpecBase [[nodiscard]] bool arch_compatible(const ModelSpecBase& o) const override { const auto* p = dynamic_cast(&o); - return p && d_model == p->d_model && num_heads == p->num_heads && - d_ff == p->d_ff && num_layers == p->num_layers && - patch_size == p->patch_size && norm_type == p->norm_type && - norm_place == p->norm_place && weight_quant == p->weight_quant; + return p && d_model == p->d_model && num_heads == p->num_heads && d_ff == p->d_ff && + num_layers == p->num_layers && patch_size == p->patch_size && + norm_type == p->norm_type && norm_place == p->norm_place && + weight_quant == p->weight_quant; } [[nodiscard]] std::string summary(const ModelSpec& shell) const override; void write_kv(KeyValueRecord& kv) const override; @@ -243,14 +246,12 @@ class GptSpec : public ModelSpecBase const auto* p = dynamic_cast(&o); // ALiBi_GPT 与 GPT 同族，统一由 GPTModel 承载，用 pos_encoding 区分 // Learned/Sinusoidal/ALiBi/RoPE —— 同为 GptSpec 即同族。 - return p && vocab_size == p->vocab_size && d_model == p->d_model && - seq_len == p->seq_len && num_heads == p->num_heads && - d_ff == p->d_ff && num_layers == p->num_layers && + return p && vocab_size == p->vocab_size && d_model == p->d_model && seq_len == p->seq_len && + num_heads == p->num_heads && d_ff == p->d_ff && num_layers == p->num_layers && pos_encoding == p->pos_encoding && activation == p->activation && norm_type == p->norm_type && subln == p->subln && n_head_kv_or_heads() == p->n_head_kv_or_heads() && - tie_embeddings == p->tie_embeddings && - weight_quant == p->weight_quant; + tie_embeddings == p->tie_embeddings && weight_quant == p->weight_quant; } [[nodiscard]] std::string summary(const ModelSpec& shell) const override; void write_kv(KeyValueRecord& kv) const override; @@ -283,12 +284,12 @@ class CnnSpec : public ModelSpecBase [[nodiscard]] bool arch_compatible(const ModelSpecBase& o) const override { const auto* p = dynamic_cast(&o); - return p && cnn_in_channels == p->cnn_in_channels && - cnn_in_size == p->cnn_in_size && cnn_pool == p->cnn_pool && - cnn_channels == p->cnn_channels && cnn_kernels == p->cnn_kernels && - cnn_strides == p->cnn_strides && cnn_paddings == p->cnn_paddings && - layer_dims == p->layer_dims && norm_type == p->norm_type && - norm_place == p->norm_place && weight_quant == p->weight_quant; + return p && cnn_in_channels == p->cnn_in_channels && cnn_in_size == p->cnn_in_size && + cnn_pool == p->cnn_pool && cnn_channels == p->cnn_channels && + cnn_kernels == p->cnn_kernels && cnn_strides == p->cnn_strides && + cnn_paddings == p->cnn_paddings && layer_dims == p->layer_dims && + norm_type == p->norm_type && norm_place == p->norm_place && + weight_quant == p->weight_quant; } [[nodiscard]] std::string summary(const ModelSpec& shell) const override; void write_kv(KeyValueRecord& kv) const override; @@ -317,9 +318,8 @@ class RaptSpec : public ModelSpecBase [[nodiscard]] bool arch_compatible(const ModelSpecBase& o) const override { const auto* p = dynamic_cast(&o); - return p && vocab_size == p->vocab_size && d_model == p->d_model && - seq_len == p->seq_len && num_heads == p->num_heads && - d_ff == p->d_ff && num_layers == p->num_layers && + return p && vocab_size == p->vocab_size && d_model == p->d_model && seq_len == p->seq_len && + num_heads == p->num_heads && d_ff == p->d_ff && num_layers == p->num_layers && pos_encoding == p->pos_encoding && activation == p->activation && norm_type == p->norm_type && weight_quant == p->weight_quant; } @@ -344,7 +344,9 @@ struct ModelSpec explicit ModelSpec(std::unique_ptr a) : arch(std::move(a)) {} ModelSpec(const ModelSpec& o) : output_schema(o.output_schema), output_packing(o.output_packing), - arch(o.arch ? o.arch->clone() : nullptr) {} + arch(o.arch ? o.arch->clone() : nullptr) + { + } ModelSpec& operator=(const ModelSpec& o) { if (this != &o) @@ -359,8 +361,7 @@ struct ModelSpec ModelSpec& operator=(ModelSpec&&) noexcept = default; // 就地构造类型化规格，返回引用便于逐字段填充 - template - T& emplace() + template T& emplace() { auto p = std::make_unique(); T& ref = *p; @@ -372,16 +373,13 @@ struct ModelSpec { return arch ? arch->type_name() : "Unknown"; } [[nodiscard]] bool has_arch() const noexcept { return arch != nullptr; } - template - [[nodiscard]] T* as() noexcept { return dynamic_cast(arch.get()); } - template - [[nodiscard]] const T* as() const noexcept + template [[nodiscard]] T* as() noexcept { return dynamic_cast(arch.get()); } + template [[nodiscard]] const T* as() const noexcept { return dynamic_cast(arch.get()); } // ── 判别（调用点兼容层；实现 = 类型化 dynamic_cast）──────────────── [[nodiscard]] bool is_mlp() const noexcept { return as() != nullptr; } - [[nodiscard]] bool is_transformer() const noexcept - { return as() != nullptr; } + [[nodiscard]] bool is_transformer() const noexcept { return as() != nullptr; } [[nodiscard]] bool is_gpt() const noexcept { return as() != nullptr; } [[nodiscard]] bool is_alibi_gpt() const noexcept { @@ -395,12 +393,9 @@ struct ModelSpec // ── 架构一致性校验（公共字段 + arch_compatible；语义见子类注释）──────────── [[nodiscard]] inline bool spec_matches(const ModelSpec& a, const ModelSpec& b) { - if (a.output_schema.has_value() != b.output_schema.has_value()) - return false; - if (a.output_schema.has_value() && *a.output_schema != *b.output_schema) - return false; - if (a.output_packing != b.output_packing) - return false; + if (a.output_schema.has_value() != b.output_schema.has_value()) return false; + if (a.output_schema.has_value() && *a.output_schema != *b.output_schema) return false; + if (a.output_packing != b.output_packing) return false; if (!a.arch || !b.arch) return false; // 任一侧无类型化规格（Unknown）⇒ 不匹配（与旧 default 一致） return a.arch->arch_compatible(*b.arch); @@ -435,15 +430,24 @@ struct ModelSpec { switch (v) { - case 0: return nullptr; // Unknown：不构造 arch - case 1: return "MLP"; - case 2: return "Transformer"; - case 3: return "GPT"; - case 4: return "GPT"; // ALiBi_GPT：同族，pos_encoding 从字段还原 - case 5: return "CNN"; - case 6: return "zipt"; // 哨兵：ZiPT 已移除（序列化层显式报错） - case 7: return "RAPT"; - default: return nullptr; + case 0: + return nullptr; // Unknown：不构造 arch + case 1: + return "MLP"; + case 2: + return "Transformer"; + case 3: + return "GPT"; + case 4: + return "GPT"; // ALiBi_GPT：同族，pos_encoding 从字段还原 + case 5: + return "CNN"; + case 6: + return "zipt"; // 哨兵：ZiPT 已移除（序列化层显式报错） + case 7: + return "RAPT"; + default: + return nullptr; } } @@ -461,8 +465,7 @@ class ModelRegistry return reg; } - void register_type(std::string name, Factory f) - { types_[std::move(name)] = std::move(f); } + void register_type(std::string name, Factory f) { types_[std::move(name)] = std::move(f); } [[nodiscard]] std::unique_ptr make(const std::string& name) const { @@ -486,8 +489,7 @@ class ModelRegistry { // 旧扁平布局的默认分支口径：layers 恒 0（MLP 用 layer_dims 表达深度） std::string out = std::string("MLP(layers=0"); - if (shell.output_schema) - out += ",outputs=" + std::to_string(shell.output_schema->size()); + if (shell.output_schema) out += ",outputs=" + std::to_string(shell.output_schema->size()); if (weight_quant != WeightQuant::None) out += ",wq=" + std::string(weight_quant_name(weight_quant)); return out + ")"; @@ -496,8 +498,7 @@ class ModelRegistry [[nodiscard]] inline std::string TransformerSpec::summary(const ModelSpec& shell) const { std::string out = std::string("Transformer(layers=") + std::to_string(num_layers); - if (shell.output_schema) - out += ",outputs=" + std::to_string(shell.output_schema->size()); + if (shell.output_schema) out += ",outputs=" + std::to_string(shell.output_schema->size()); if (weight_quant != WeightQuant::None) out += ",wq=" + std::string(weight_quant_name(weight_quant)); return out + ")"; @@ -506,21 +507,17 @@ class ModelRegistry [[nodiscard]] inline std::string GptSpec::summary(const ModelSpec&) const { return std::string("GPT(vocab=") + std::to_string(vocab_size) + - ",d_model=" + std::to_string(d_model) + - ",seq_len=" + std::to_string(seq_len) + - ",heads=" + std::to_string(num_heads) + - ",d_ff=" + std::to_string(d_ff) + - ",layers=" + std::to_string(num_layers) + - (subln ? ",subln=1" : "") + - (n_head_kv_or_heads() != num_heads - ? ",kv_heads=" + std::to_string(n_head_kv_or_heads()) : "") + + ",d_model=" + std::to_string(d_model) + ",seq_len=" + std::to_string(seq_len) + + ",heads=" + std::to_string(num_heads) + ",d_ff=" + std::to_string(d_ff) + + ",layers=" + std::to_string(num_layers) + (subln ? ",subln=1" : "") + + (n_head_kv_or_heads() != num_heads ? ",kv_heads=" + std::to_string(n_head_kv_or_heads()) + : "") + (tie_embeddings ? ",tied=1" : "") + ")"; } [[nodiscard]] inline std::string CnnSpec::summary(const ModelSpec&) const { - std::string c = std::to_string(cnn_in_channels) + "x" + - std::to_string(cnn_in_size) + "->"; + std::string c = std::to_string(cnn_in_channels) + "x" + std::to_string(cnn_in_size) + "->"; for (std::size_t i = 0; i < cnn_channels.size(); ++i) { if (i) c += "-"; @@ -532,10 +529,8 @@ class ModelRegistry [[nodiscard]] inline std::string RaptSpec::summary(const ModelSpec&) const { return std::string("RAPT(vocab=") + std::to_string(vocab_size) + - ",d_model=" + std::to_string(d_model) + - ",seq_len=" + std::to_string(seq_len) + - ",heads=" + std::to_string(num_heads) + - ",d_ff=" + std::to_string(d_ff) + + ",d_model=" + std::to_string(d_model) + ",seq_len=" + std::to_string(seq_len) + + ",heads=" + std::to_string(num_heads) + ",d_ff=" + std::to_string(d_ff) + ",layers=" + std::to_string(num_layers) + ",pos_enc=" + std::to_string(static_cast(pos_encoding)) + ")"; } @@ -545,7 +540,8 @@ inline void MlpSpec::write_kv(KeyValueRecord& kv) const { std::vector dims; dims.reserve(layer_dims.size()); - for (auto d : layer_dims) dims.push_back(static_cast(d)); + for (auto d : layer_dims) + dims.push_back(static_cast(d)); kv.set("layer_dims", dims); kv.set("norm_type", static_cast(norm_type)); kv.set("weight_quant", static_cast(weight_quant)); @@ -585,10 +581,12 @@ inline void CnnSpec::write_kv(KeyValueRecord& kv) const kv.set("cnn_in_channels", static_cast(cnn_in_channels)); kv.set("cnn_in_size", static_cast(cnn_in_size)); kv.set("cnn_pool", static_cast(cnn_pool)); - auto arr = [](const std::vector& v) { + auto arr = [](const std::vector& v) + { std::vector out; out.reserve(v.size()); - for (auto x : v) out.push_back(static_cast(x)); + for (auto x : v) + out.push_back(static_cast(x)); return out; }; kv.set("cnn_channels", arr(cnn_channels)); @@ -628,15 +626,15 @@ inline void get_size(const KeyValueRecord& kv, const char* key, std::size_t& out std::uint64_t v = 0; if (kv.get(key, v)) out = static_cast(v); } -inline void get_dims(const KeyValueRecord& kv, const char* key, - std::vector& out) +inline void get_dims(const KeyValueRecord& kv, const char* key, std::vector& out) { std::vector v; if (kv.get(key, v)) { out.clear(); out.reserve(v.size()); - for (auto x : v) out.push_back(static_cast(x)); + for (auto x : v) + out.push_back(static_cast(x)); } } } // namespace spec_kv_detail diff --git a/include/neuralnet.cpp/named_tensor_set.hpp b/include/neuralnet.cpp/named_tensor_set.hpp index 8d7cb959..1ec11f4b 100644 --- a/include/neuralnet.cpp/named_tensor_set.hpp +++ b/include/neuralnet.cpp/named_tensor_set.hpp @@ -48,11 +48,10 @@ class NamedTensorSet return values_[handle.index].valid() ? &values_[handle.index] : nullptr; } - [[nodiscard]] static Result validate_value_( - const TensorFieldSpec& field, const Tensor& value) + [[nodiscard]] static Result validate_value_(const TensorFieldSpec& field, + const Tensor& value) { - if (!value.valid()) - NN_FAIL("NamedTensorSet: 字段 '" + field.name + "' 是无效张量"); + if (!value.valid()) NN_FAIL("NamedTensorSet: 字段 '" + field.name + "' 是无效张量"); if (!field.shape.accepts(value.rows(), value.cols())) NN_FAIL("NamedTensorSet: 字段 '" + field.name + "' 形状不匹配（期望 " + field.shape.to_string() + "，实际 " + value.shape_str() + "）"); @@ -62,12 +61,10 @@ class NamedTensorSet precision_name(value.precision()) + "）"); if (field.placement != TensorPlacement::Any) { - const TensorPlacement actual = value.device() == Device::CPU - ? TensorPlacement::CPU - : TensorPlacement::GPU; + const TensorPlacement actual = + value.device() == Device::CPU ? TensorPlacement::CPU : TensorPlacement::GPU; if (field.placement != actual) - NN_FAIL("NamedTensorSet: 字段 '" + field.name + - "' 设备不匹配（期望 " + + NN_FAIL("NamedTensorSet: 字段 '" + field.name + "' 设备不匹配（期望 " + tensor_placement_name(field.placement) + "，实际 " + tensor_placement_name(actual) + "）"); } @@ -82,11 +79,10 @@ class NamedTensorSet NamedTensorSet& operator=(NamedTensorSet&&) noexcept = default; ~NamedTensorSet() = default; - [[nodiscard]] static Result create( - TensorSchema schema, std::vector tensors) + [[nodiscard]] static Result create(TensorSchema schema, + std::vector tensors) { - if (!schema.valid()) - NN_FAIL("NamedTensorSet::create: schema 未初始化"); + if (!schema.valid()) NN_FAIL("NamedTensorSet::create: schema 未初始化"); std::vector ordered(schema.size()); std::vector seen(schema.size(), false); @@ -97,14 +93,12 @@ class NamedTensorSet auto h_r = schema.optional_field(nt.name); NN_TRY_CHECK(h_r); if (!h_r->valid()) - NN_FAIL("NamedTensorSet::create: schema '" + schema.name() + - "' 不认识字段 '" + nt.name + "'"); + NN_FAIL("NamedTensorSet::create: schema '" + schema.name() + "' 不认识字段 '" + + nt.name + "'"); if (seen[h_r->index]) - NN_FAIL("NamedTensorSet::create: 字段 '" + nt.name + - "' 重复出现"); + NN_FAIL("NamedTensorSet::create: 字段 '" + nt.name + "' 重复出现"); seen[h_r->index] = true; - NN_TRY_CHECK(validate_value_( - schema.fields()[h_r->index], nt.value)); + NN_TRY_CHECK(validate_value_(schema.fields()[h_r->index], nt.value)); ordered[h_r->index] = nt.value; ++matched; } @@ -112,11 +106,10 @@ class NamedTensorSet for (std::size_t i = 0; i < schema.size(); ++i) { if (!seen[i] && schema.fields()[i].required) - NN_FAIL("NamedTensorSet::create: schema '" + schema.name() + - "' 缺少必需字段 '" + schema.fields()[i].name + "'"); + NN_FAIL("NamedTensorSet::create: schema '" + schema.name() + "' 缺少必需字段 '" + + schema.fields()[i].name + "'"); } - if (tensors.size() != matched) - NN_FAIL("NamedTensorSet::create: 输入字段计数不一致"); + if (tensors.size() != matched) NN_FAIL("NamedTensorSet::create: 输入字段计数不一致"); return NamedTensorSet(std::move(schema), std::move(ordered)); } @@ -126,30 +119,18 @@ class NamedTensorSet [[nodiscard]] bool empty() const noexcept { return values_.empty(); } [[nodiscard]] Result field(std::string_view name) const - { - return schema_.field(name); - } + { return schema_.field(name); } - [[nodiscard]] Result optional_field( - std::string_view name) const - { - return schema_.optional_field(name); - } + [[nodiscard]] Result optional_field(std::string_view name) const + { return schema_.optional_field(name); } - [[nodiscard]] Result require(const FieldHandle& handle) - { - return value_(handle); - } + [[nodiscard]] Result require(const FieldHandle& handle) { return value_(handle); } - [[nodiscard]] Result require( - const FieldHandle& handle) const - { - return value_(handle); - } + [[nodiscard]] Result require(const FieldHandle& handle) const + { return value_(handle); } // setup 期便捷入口；热路径请先缓存 FieldHandle。 - [[nodiscard]] Result require( - std::string_view name) const + [[nodiscard]] Result require(std::string_view name) const { NN_TRY(h, schema_.field(name)); return require(*h); @@ -157,8 +138,7 @@ class NamedTensorSet [[nodiscard]] const Tensor* find(const FieldHandle& handle) const noexcept { - if (!handle.valid() || - handle.schema_fingerprint != schema_.fingerprint() || + if (!handle.valid() || handle.schema_fingerprint != schema_.fingerprint() || handle.index >= values_.size()) return nullptr; return values_[handle.index].valid() ? &values_[handle.index] : nullptr; @@ -167,16 +147,12 @@ class NamedTensorSet [[nodiscard]] const Tensor* find(std::string_view name) const noexcept { for (std::size_t i = 0; i < schema_.size(); ++i) - if (schema_.fields()[i].name == name) - return values_[i].valid() ? &values_[i] : nullptr; + if (schema_.fields()[i].name == name) return values_[i].valid() ? &values_[i] : nullptr; return nullptr; } [[nodiscard]] Tensor* find(const FieldHandle& handle) noexcept - { - return const_cast( - static_cast(this)->find(handle)); - } + { return const_cast(static_cast(this)->find(handle)); } }; // 梯度与输出端口共享同一 schema；类型独立，避免把推理值误当梯度传入。 @@ -185,10 +161,7 @@ class GradientBundle private: NamedTensorSet values_; - explicit GradientBundle(NamedTensorSet values) - : values_(std::move(values)) - { - } + explicit GradientBundle(NamedTensorSet values) : values_(std::move(values)) {} public: GradientBundle() = default; @@ -198,36 +171,23 @@ class GradientBundle GradientBundle& operator=(GradientBundle&&) noexcept = default; ~GradientBundle() = default; - [[nodiscard]] static Result create( - TensorSchema schema, std::vector gradients) + [[nodiscard]] static Result create(TensorSchema schema, + std::vector gradients) { - NN_TRY(values, - NamedTensorSet::create(std::move(schema), std::move(gradients))); + NN_TRY(values, NamedTensorSet::create(std::move(schema), std::move(gradients))); return GradientBundle(std::move(*values)); } - [[nodiscard]] const TensorSchema& schema() const noexcept - { - return values_.schema(); - } + [[nodiscard]] const TensorSchema& schema() const noexcept { return values_.schema(); } [[nodiscard]] std::size_t size() const noexcept { return values_.size(); } [[nodiscard]] Result field(std::string_view name) const - { - return values_.field(name); - } - [[nodiscard]] Result require( - const FieldHandle& handle) const - { - return values_.require(handle); - } + { return values_.field(name); } + [[nodiscard]] Result require(const FieldHandle& handle) const + { return values_.require(handle); } [[nodiscard]] Result require(const FieldHandle& handle) - { - return values_.require(handle); - } + { return values_.require(handle); } [[nodiscard]] const Tensor* find(const FieldHandle& handle) const noexcept - { - return values_.find(handle); - } + { return values_.find(handle); } }; } // namespace nn diff --git a/include/neuralnet.cpp/nn.hpp b/include/neuralnet.cpp/nn.hpp index 3d4ae115..ab52f6d9 100644 --- a/include/neuralnet.cpp/nn.hpp +++ b/include/neuralnet.cpp/nn.hpp @@ -63,4 +63,3 @@ // 注：训练驱动 `text_trainer.hpp`（L5：数据集 + 样本 + 循环）**不**进聚合头， // 需要时 include 薄转发头 `train.hpp` 或直接 include `text_trainer.hpp`。 // 训练配方 `text_train_config.hpp` 随 train.hpp / text_trainer.hpp 一起提供。 - diff --git a/include/neuralnet.cpp/rl_experience.hpp b/include/neuralnet.cpp/rl_experience.hpp index e8d9e3a6..3d1e1c0c 100644 --- a/include/neuralnet.cpp/rl_experience.hpp +++ b/include/neuralnet.cpp/rl_experience.hpp @@ -34,10 +34,7 @@ class TensorRecord private: NamedTensorSet values_; - explicit TensorRecord(NamedTensorSet values) - : values_(std::move(values)) - { - } + explicit TensorRecord(NamedTensorSet values) : values_(std::move(values)) {} public: TensorRecord() = default; @@ -47,43 +44,27 @@ class TensorRecord TensorRecord& operator=(TensorRecord&&) noexcept = default; ~TensorRecord() = default; - [[nodiscard]] static Result create( - TensorSchema schema, std::vector values) + [[nodiscard]] static Result create(TensorSchema schema, + std::vector values) { - NN_TRY(set, NamedTensorSet::create(std::move(schema), - std::move(values))); + NN_TRY(set, NamedTensorSet::create(std::move(schema), std::move(values))); return TensorRecord(std::move(*set)); } - [[nodiscard]] const TensorSchema& schema() const noexcept - { - return values_.schema(); - } + [[nodiscard]] const TensorSchema& schema() const noexcept { return values_.schema(); } [[nodiscard]] std::size_t size() const noexcept { return values_.size(); } [[nodiscard]] Result field(std::string_view name) const - { - return values_.field(name); - } + { return values_.field(name); } [[nodiscard]] Result optional_field(std::string_view name) const - { - return values_.optional_field(name); - } + { return values_.optional_field(name); } [[nodiscard]] Result require(const FieldHandle& handle) const - { - return values_.require(handle); - } + { return values_.require(handle); } [[nodiscard]] Result require(std::string_view name) const - { - return values_.require(name); - } + { return values_.require(name); } [[nodiscard]] const Tensor* find(const FieldHandle& handle) const noexcept - { - return values_.find(handle); - } + { return values_.find(handle); } [[nodiscard]] const Tensor* find(std::string_view name) const noexcept - { - return values_.find(name); - } + { return values_.find(name); } }; class ExperienceSchema @@ -101,17 +82,14 @@ class ExperienceSchema std::optional done_; std::optional group_id_; - static void resolve_(const TensorSchema& schema, - std::string_view name, + static void resolve_(const TensorSchema& schema, std::string_view name, std::optional& out) { auto h = NN_UNWRAP(schema.optional_field(name)); if (h.valid()) out = h; } - ExperienceSchema(TensorSchema schema, - std::vector) - : schema_(std::move(schema)) + ExperienceSchema(TensorSchema schema, std::vector) : schema_(std::move(schema)) { resolve_(schema_, "action", action_); resolve_(schema_, "old_log_prob", old_log_prob_); @@ -135,21 +113,15 @@ class ExperienceSchema // 标准字段先全部登记为可选，再按 required 列表收紧；custom 字段追加在 // 标准字段之后。required 中的未知名称 -> 明确错误。 - [[nodiscard]] static Result create( - std::vector required = {}, - std::vector custom = {}) + [[nodiscard]] static Result create(std::vector required = {}, + std::vector custom = {}) { std::vector fields{ - tensor_field("action", {}, false), - tensor_field("old_log_prob", {}, false), - tensor_field("advantage", {}, false), - tensor_field("return", {}, false), - tensor_field("old_value", {}, false), - tensor_field("action_mask", {}, false), - tensor_field("value_mask", {}, false), - tensor_field("reward", {}, false), - tensor_field("done", {}, false), - tensor_field("group_id", {}, false), + tensor_field("action", {}, false), tensor_field("old_log_prob", {}, false), + tensor_field("advantage", {}, false), tensor_field("return", {}, false), + tensor_field("old_value", {}, false), tensor_field("action_mask", {}, false), + tensor_field("value_mask", {}, false), tensor_field("reward", {}, false), + tensor_field("done", {}, false), tensor_field("group_id", {}, false), }; for (const std::string& name : required) { @@ -164,8 +136,8 @@ class ExperienceSchema } } if (!found) - NN_FAIL("ExperienceSchema::create: required 含未知标准字段 '" + - name + "'（自定义字段请在定义中自行设 required）"); + NN_FAIL("ExperienceSchema::create: required 含未知标准字段 '" + name + + "'（自定义字段请在定义中自行设 required）"); } fields.insert(fields.end(), custom.begin(), custom.end()); NN_TRY(schema, TensorSchema::create("experience", std::move(fields))); @@ -173,46 +145,16 @@ class ExperienceSchema } [[nodiscard]] const TensorSchema& schema() const noexcept { return schema_; } - [[nodiscard]] std::optional action() const noexcept - { - return action_; - } - [[nodiscard]] std::optional old_log_prob() const noexcept - { - return old_log_prob_; - } - [[nodiscard]] std::optional advantage() const noexcept - { - return advantage_; - } - [[nodiscard]] std::optional returns() const noexcept - { - return returns_; - } - [[nodiscard]] std::optional old_value() const noexcept - { - return old_value_; - } - [[nodiscard]] std::optional action_mask() const noexcept - { - return action_mask_; - } - [[nodiscard]] std::optional value_mask() const noexcept - { - return value_mask_; - } - [[nodiscard]] std::optional reward() const noexcept - { - return reward_; - } - [[nodiscard]] std::optional done() const noexcept - { - return done_; - } - [[nodiscard]] std::optional group_id() const noexcept - { - return group_id_; - } + [[nodiscard]] std::optional action() const noexcept { return action_; } + [[nodiscard]] std::optional old_log_prob() const noexcept { return old_log_prob_; } + [[nodiscard]] std::optional advantage() const noexcept { return advantage_; } + [[nodiscard]] std::optional returns() const noexcept { return returns_; } + [[nodiscard]] std::optional old_value() const noexcept { return old_value_; } + [[nodiscard]] std::optional action_mask() const noexcept { return action_mask_; } + [[nodiscard]] std::optional value_mask() const noexcept { return value_mask_; } + [[nodiscard]] std::optional reward() const noexcept { return reward_; } + [[nodiscard]] std::optional done() const noexcept { return done_; } + [[nodiscard]] std::optional group_id() const noexcept { return group_id_; } }; class PolicyBatchView @@ -228,10 +170,7 @@ class PolicyBatchView std::optional value_mask_; std::optional group_id_; - PolicyBatchView(const TensorRecord& record) - : record_(&record) - { - } + PolicyBatchView(const TensorRecord& record) : record_(&record) {} public: PolicyBatchView() = default; @@ -241,20 +180,17 @@ class PolicyBatchView PolicyBatchView& operator=(PolicyBatchView&&) noexcept = default; ~PolicyBatchView() = default; - [[nodiscard]] static Result create( - const TensorRecord& record) + [[nodiscard]] static Result create(const TensorRecord& record) { auto action = record.field("action"); - if (!action) - NN_FAIL("PolicyBatchView::create: 缺少 action"); + if (!action) NN_FAIL("PolicyBatchView::create: 缺少 action"); auto action_mask = record.field("action_mask"); - if (!action_mask) - NN_FAIL("PolicyBatchView::create: 缺少 action_mask"); + if (!action_mask) NN_FAIL("PolicyBatchView::create: 缺少 action_mask"); PolicyBatchView view(record); view.action_ = *action; view.action_mask_ = *action_mask; - auto set_optional = [&](std::string_view name, - std::optional& out) { + auto set_optional = [&](std::string_view name, std::optional& out) + { auto h = record.optional_field(name); if (h && h->valid()) out = *h; }; @@ -268,38 +204,20 @@ class PolicyBatchView } [[nodiscard]] const TensorRecord& record() const noexcept { return *record_; } - [[nodiscard]] const Tensor& action() const - { - return *record_->find(action_); - } - [[nodiscard]] const Tensor& action_mask() const - { - return *record_->find(action_mask_); - } + [[nodiscard]] const Tensor& action() const { return *record_->find(action_); } + [[nodiscard]] const Tensor& action_mask() const { return *record_->find(action_mask_); } [[nodiscard]] const Tensor* old_log_prob() const noexcept - { - return old_log_prob_ ? record_->find(*old_log_prob_) : nullptr; - } + { return old_log_prob_ ? record_->find(*old_log_prob_) : nullptr; } [[nodiscard]] const Tensor* advantage() const noexcept - { - return advantage_ ? record_->find(*advantage_) : nullptr; - } + { return advantage_ ? record_->find(*advantage_) : nullptr; } [[nodiscard]] const Tensor* returns() const noexcept - { - return returns_ ? record_->find(*returns_) : nullptr; - } + { return returns_ ? record_->find(*returns_) : nullptr; } [[nodiscard]] const Tensor* old_value() const noexcept - { - return old_value_ ? record_->find(*old_value_) : nullptr; - } + { return old_value_ ? record_->find(*old_value_) : nullptr; } [[nodiscard]] const Tensor* value_mask() const noexcept - { - return value_mask_ ? record_->find(*value_mask_) : nullptr; - } + { return value_mask_ ? record_->find(*value_mask_) : nullptr; } [[nodiscard]] const Tensor* group_id() const noexcept - { - return group_id_ ? record_->find(*group_id_) : nullptr; - } + { return group_id_ ? record_->find(*group_id_) : nullptr; } }; struct AdvantageConfig @@ -311,36 +229,29 @@ struct AdvantageConfig namespace detail { -[[nodiscard]] inline Result rl_zero( - ComputeEngine& engine, std::size_t rows, std::size_t cols, Precision p) +[[nodiscard]] inline Result rl_zero(ComputeEngine& engine, std::size_t rows, + std::size_t cols, Precision p) { NN_TRY(t, engine.try_create_tensor(rows, cols, p)); NN_TRY_CHECK(engine.zero(*t)); return t; } -[[nodiscard]] inline Result rl_constant( - ComputeEngine& engine, std::size_t rows, std::size_t cols, - Scalar value, Precision p) +[[nodiscard]] inline Result rl_constant(ComputeEngine& engine, std::size_t rows, + std::size_t cols, Scalar value, Precision p) { NN_TRY(zero, rl_zero(engine, rows, cols, p)); - if (value == Scalar{0}) - return zero; - return dsl::compute(engine, - dsl::leaf(*zero) + dsl::rparam(value), - rows, cols, p); + if (value == Scalar{0}) return zero; + return dsl::compute(engine, dsl::leaf(*zero) + dsl::rparam(value), rows, cols, p); } -inline Result rl_check_same_shape( - const Tensor& reference, const Tensor* other, const char* name) +inline Result rl_check_same_shape(const Tensor& reference, const Tensor* other, + const char* name) { if (other == nullptr) return {}; - if (other->rows() != reference.rows() || - other->cols() != reference.cols()) - NN_FAIL(std::string("RL transform: ") + name + - " 形状与 reward/value 不一致（期望 " + - reference.shape_str() + "，实际 " + - other->shape_str() + "）"); + if (other->rows() != reference.rows() || other->cols() != reference.cols()) + NN_FAIL(std::string("RL transform: ") + name + " 形状与 reward/value 不一致（期望 " + + reference.shape_str() + "，实际 " + other->shape_str() + "）"); return {}; } } // namespace detail @@ -351,10 +262,7 @@ class GaeAdvantage AdvantageConfig cfg_; public: - explicit GaeAdvantage(AdvantageConfig cfg = {}) - : cfg_(cfg) - { - } + explicit GaeAdvantage(AdvantageConfig cfg = {}) : cfg_(cfg) {} struct Output { @@ -362,22 +270,20 @@ class GaeAdvantage Tensor returns; }; - [[nodiscard]] Result apply( - ComputeEngine& engine, const Tensor& rewards, const Tensor& values, - const Tensor* done = nullptr, const Tensor* mask = nullptr, - const Tensor* bootstrap_value = nullptr) const + [[nodiscard]] Result apply(ComputeEngine& engine, const Tensor& rewards, + const Tensor& values, const Tensor* done = nullptr, + const Tensor* mask = nullptr, + const Tensor* bootstrap_value = nullptr) const { const std::size_t steps = rewards.rows(); const std::size_t batch = rewards.cols(); - if (steps == 0 || batch == 0) - NN_FAIL("GaeAdvantage: reward 不能为空"); + if (steps == 0 || batch == 0) NN_FAIL("GaeAdvantage: reward 不能为空"); if (values.rows() != steps || values.cols() != batch) NN_FAIL("GaeAdvantage: value 与 reward 形状不一致"); NN_TRY_CHECK(detail::rl_check_same_shape(rewards, done, "done")); NN_TRY_CHECK(detail::rl_check_same_shape(rewards, mask, "mask")); if (bootstrap_value != nullptr && - (bootstrap_value->rows() != 1 || - bootstrap_value->cols() != batch)) + (bootstrap_value->rows() != 1 || bootstrap_value->cols() != batch)) NN_FAIL("GaeAdvantage: bootstrap_value 形状必须为 (1,batch)"); const Precision p = cfg_.precision; @@ -410,8 +316,8 @@ class GaeAdvantage { NN_TRY(d, engine.slice_rows(*done, ti, 1)); NN_TRY(c, dsl::compute(engine, - dsl::leaf(*reward) * dsl::rparam(0) + - dsl::rparam(1) - dsl::leaf(*d), + dsl::leaf(*reward) * dsl::rparam(0) + dsl::rparam(1) - + dsl::leaf(*d), 1, batch, p)); continuation = std::move(*c); } @@ -427,38 +333,32 @@ class GaeAdvantage dsl::leaf(continuation) - dsl::leaf(*value), 1, batch, p)); - NN_TRY(advance, dsl::compute( - engine, - dsl::leaf(*delta) + - dsl::leaf(*next_adv) * - dsl::rparam(cfg_.gamma * cfg_.lambda) * - dsl::leaf(continuation), + NN_TRY(advance, + dsl::compute(engine, + dsl::leaf(*delta) + dsl::leaf(*next_adv) * + dsl::rparam(cfg_.gamma * cfg_.lambda) * + dsl::leaf(continuation), 1, batch, p)); if (mask_row != nullptr) { - NN_TRY(masked_a, dsl::compute(engine, - dsl::leaf(*advance) * - dsl::leaf(*mask_row), - 1, batch, p)); + NN_TRY(masked_a, dsl::compute(engine, dsl::leaf(*advance) * dsl::leaf(*mask_row), 1, + batch, p)); advance = std::move(*masked_a); } Tensor ret_row; if (mask_row != nullptr) { - NN_TRY(ret, dsl::compute( - engine, - (dsl::leaf(*advance) + dsl::leaf(*value)) * - dsl::leaf(*mask_row), - 1, batch, p)); + NN_TRY(ret, + dsl::compute( + engine, (dsl::leaf(*advance) + dsl::leaf(*value)) * dsl::leaf(*mask_row), + 1, batch, p)); ret_row = std::move(*ret); } else { - NN_TRY(ret, dsl::compute( - engine, - dsl::leaf(*advance) + dsl::leaf(*value), - 1, batch, p)); + NN_TRY(ret, + dsl::compute(engine, dsl::leaf(*advance) + dsl::leaf(*value), 1, batch, p)); ret_row = std::move(*ret); } @@ -476,8 +376,7 @@ class GaeAdvantage next_value = *value; } } - return Output{.advantage = std::move(*out_a), - .returns = std::move(*out_r)}; + return Output{.advantage = std::move(*out_a), .returns = std::move(*out_r)}; } }; @@ -488,20 +387,18 @@ class MonteCarloReturn Precision precision_ = Precision::F32; public: - explicit MonteCarloReturn(Scalar gamma = 0.99f, - Precision precision = Precision::F32) + explicit MonteCarloReturn(Scalar gamma = 0.99f, Precision precision = Precision::F32) : gamma_(gamma), precision_(precision) { } - [[nodiscard]] Result apply( - ComputeEngine& engine, const Tensor& rewards, - const Tensor* done = nullptr, const Tensor* mask = nullptr) const + [[nodiscard]] Result apply(ComputeEngine& engine, const Tensor& rewards, + const Tensor* done = nullptr, + const Tensor* mask = nullptr) const { const std::size_t steps = rewards.rows(); const std::size_t batch = rewards.cols(); - if (steps == 0 || batch == 0) - NN_FAIL("MonteCarloReturn: reward 不能为空"); + if (steps == 0 || batch == 0) NN_FAIL("MonteCarloReturn: reward 不能为空"); NN_TRY_CHECK(detail::rl_check_same_shape(rewards, done, "done")); NN_TRY_CHECK(detail::rl_check_same_shape(rewards, mask, "mask")); @@ -515,29 +412,25 @@ class MonteCarloReturn { NN_TRY(d, engine.slice_rows(*done, ti, 1)); NN_TRY(c, dsl::compute(engine, - dsl::leaf(*reward) * dsl::rparam(0) + - dsl::rparam(1) - dsl::leaf(*d), + dsl::leaf(*reward) * dsl::rparam(0) + dsl::rparam(1) - + dsl::leaf(*d), 1, batch, precision_)); continuation = std::move(*c); } else { - NN_TRY(c, detail::rl_constant(engine, 1, batch, 1, - precision_)); + NN_TRY(c, detail::rl_constant(engine, 1, batch, 1, precision_)); continuation = std::move(*c); } - NN_TRY(row, dsl::compute( - engine, - dsl::leaf(*reward) + dsl::leaf(*next) * - dsl::rparam(gamma_) * - dsl::leaf(continuation), - 1, batch, precision_)); + NN_TRY(row, dsl::compute(engine, + dsl::leaf(*reward) + dsl::leaf(*next) * dsl::rparam(gamma_) * + dsl::leaf(continuation), + 1, batch, precision_)); if (mask != nullptr) { NN_TRY(m, engine.slice_rows(*mask, ti, 1)); - NN_TRY(masked, dsl::compute(engine, - dsl::leaf(*row) * dsl::leaf(*m), - 1, batch, precision_)); + NN_TRY(masked, + dsl::compute(engine, dsl::leaf(*row) * dsl::leaf(*m), 1, batch, precision_)); row = std::move(*masked); } NN_TRY_CHECK(engine.insert_rows(*out, ti, *row)); @@ -555,47 +448,33 @@ class MaskedNormalizer Precision precision_ = Precision::F32; public: - explicit MaskedNormalizer(Scalar epsilon = 1e-8f, - Precision precision = Precision::F32) + explicit MaskedNormalizer(Scalar epsilon = 1e-8f, Precision precision = Precision::F32) : epsilon_(epsilon), precision_(precision) { } - [[nodiscard]] Result apply( - ComputeEngine& engine, const Tensor& values, - const Tensor& mask) const + [[nodiscard]] Result apply(ComputeEngine& engine, const Tensor& values, + const Tensor& mask) const { - if (values.rows() == 0 || values.cols() == 0) - NN_FAIL("MaskedNormalizer: values 不能为空"); + if (values.rows() == 0 || values.cols() == 0) NN_FAIL("MaskedNormalizer: values 不能为空"); NN_TRY_CHECK(detail::rl_check_same_shape(values, &mask, "mask")); - NN_TRY(weighted, dsl::compute(engine, - dsl::leaf(values) * dsl::leaf(mask), - values.rows(), values.cols(), - precision_)); - NN_TRY(square, dsl::compute(engine, - dsl::leaf(values) * dsl::leaf(values) * - dsl::leaf(mask), - values.rows(), values.cols(), - precision_)); - NN_TRY(sum_w, dsl::compute_reduce( - engine, dsl::col_reduce_sum(dsl::leaf(*weighted)), - values.rows(), values.cols(), precision_)); - NN_TRY(sum_sq, dsl::compute_reduce( - engine, dsl::col_reduce_sum(dsl::leaf(*square)), - values.rows(), values.cols(), precision_)); - NN_TRY(sum_m, dsl::compute_reduce( - engine, dsl::col_reduce_sum(dsl::leaf(mask)), - values.rows(), values.cols(), precision_)); - NN_TRY(total_w, dsl::compute_reduce( - engine, dsl::row_reduce_sum(dsl::leaf(*sum_w)), - 1, values.cols(), precision_)); - NN_TRY(total_sq, dsl::compute_reduce( - engine, dsl::row_reduce_sum(dsl::leaf(*sum_sq)), - 1, values.cols(), precision_)); - NN_TRY(total_m, dsl::compute_reduce( - engine, dsl::row_reduce_sum(dsl::leaf(*sum_m)), - 1, values.cols(), precision_)); + NN_TRY(weighted, dsl::compute(engine, dsl::leaf(values) * dsl::leaf(mask), values.rows(), + values.cols(), precision_)); + NN_TRY(square, dsl::compute(engine, dsl::leaf(values) * dsl::leaf(values) * dsl::leaf(mask), + values.rows(), values.cols(), precision_)); + NN_TRY(sum_w, dsl::compute_reduce(engine, dsl::col_reduce_sum(dsl::leaf(*weighted)), + values.rows(), values.cols(), precision_)); + NN_TRY(sum_sq, dsl::compute_reduce(engine, dsl::col_reduce_sum(dsl::leaf(*square)), + values.rows(), values.cols(), precision_)); + NN_TRY(sum_m, dsl::compute_reduce(engine, dsl::col_reduce_sum(dsl::leaf(mask)), + values.rows(), values.cols(), precision_)); + NN_TRY(total_w, dsl::compute_reduce(engine, dsl::row_reduce_sum(dsl::leaf(*sum_w)), 1, + values.cols(), precision_)); + NN_TRY(total_sq, dsl::compute_reduce(engine, dsl::row_reduce_sum(dsl::leaf(*sum_sq)), 1, + values.cols(), precision_)); + NN_TRY(total_m, dsl::compute_reduce(engine, dsl::row_reduce_sum(dsl::leaf(*sum_m)), 1, + values.cols(), precision_)); auto w = detail::download_vector(engine, *total_w); NN_TRY_CHECK(w); auto sq = detail::download_vector(engine, *total_sq); @@ -605,18 +484,16 @@ class MaskedNormalizer const Scalar count = (*m)[0]; if (count <= Scalar{0}) - return detail::rl_zero(engine, values.rows(), values.cols(), - precision_); + return detail::rl_zero(engine, values.rows(), values.cols(), precision_); const Scalar mean = (*w)[0] / count; Scalar variance = (*sq)[0] / count - mean * mean; if (variance < Scalar{0}) variance = Scalar{0}; if (variance <= epsilon_ * epsilon_) - return detail::rl_zero(engine, values.rows(), values.cols(), - precision_); + return detail::rl_zero(engine, values.rows(), values.cols(), precision_); const Scalar inv_std = Scalar{1} / std::sqrt(variance); return dsl::compute(engine, - (dsl::leaf(values) - dsl::rparam(mean)) * - dsl::rparam(inv_std) * dsl::leaf(mask), + (dsl::leaf(values) - dsl::rparam(mean)) * dsl::rparam(inv_std) * + dsl::leaf(mask), values.rows(), values.cols(), precision_); } }; @@ -634,40 +511,32 @@ class GroupAdvantage Precision precision_ = Precision::F32; public: - explicit GroupAdvantage(Scalar epsilon = 1e-8f, - Precision precision = Precision::F32) + explicit GroupAdvantage(Scalar epsilon = 1e-8f, Precision precision = Precision::F32) : epsilon_(epsilon), precision_(precision) { } // returns: (N,1)；spans 必须有序、互不重叠并连续覆盖全部 N 行。 - [[nodiscard]] Result apply( - ComputeEngine& engine, const Tensor& returns, - std::span spans) const + [[nodiscard]] Result apply(ComputeEngine& engine, const Tensor& returns, + std::span spans) const { if (returns.cols() != 1 || returns.rows() == 0) NN_FAIL("GroupAdvantage: returns 形状必须为 (N,1)"); std::size_t cursor = 0; for (const GroupSpan& span : spans) { - if (span.row_count == 0) - NN_FAIL("GroupAdvantage: group row_count 必须 > 0"); - if (span.first_row != cursor) - NN_FAIL("GroupAdvantage: group spans 必须连续且无重叠"); + if (span.row_count == 0) NN_FAIL("GroupAdvantage: group row_count 必须 > 0"); + if (span.first_row != cursor) NN_FAIL("GroupAdvantage: group spans 必须连续且无重叠"); cursor += span.row_count; } - if (cursor != returns.rows()) - NN_FAIL("GroupAdvantage: group spans 未覆盖全部 returns 行"); + if (cursor != returns.rows()) NN_FAIL("GroupAdvantage: group spans 未覆盖全部 returns 行"); NN_TRY(out, engine.try_create_tensor(returns.rows(), 1, precision_)); MaskedNormalizer normalizer(epsilon_, precision_); for (const GroupSpan& span : spans) { - NN_TRY(group, - engine.slice_rows(returns, span.first_row, span.row_count)); - NN_TRY(mask, - detail::rl_constant(engine, span.row_count, 1, 1, - precision_)); + NN_TRY(group, engine.slice_rows(returns, span.first_row, span.row_count)); + NN_TRY(mask, detail::rl_constant(engine, span.row_count, 1, 1, precision_)); NN_TRY(normalized, normalizer.apply(engine, *group, *mask)); NN_TRY_CHECK(engine.insert_rows(*out, span.first_row, *normalized)); } diff --git a/include/neuralnet.cpp/rl_pipeline.hpp b/include/neuralnet.cpp/rl_pipeline.hpp index fa5491ed..1cb53919 100644 --- a/include/neuralnet.cpp/rl_pipeline.hpp +++ b/include/neuralnet.cpp/rl_pipeline.hpp @@ -33,22 +33,14 @@ class DeterministicRng std::mt19937_64 engine_; public: - explicit DeterministicRng(std::uint64_t seed = 0) - : seed_(seed), engine_(seed) - { - } + explicit DeterministicRng(std::uint64_t seed = 0) : seed_(seed), engine_(seed) {} [[nodiscard]] std::uint64_t seed() const noexcept { return seed_; } - [[nodiscard]] std::uint64_t next_u64() - { - return engine_(); - } + [[nodiscard]] std::uint64_t next_u64() { return engine_(); } [[nodiscard]] std::size_t uniform_index(std::size_t upper_exclusive) { - if (upper_exclusive == 0) - NN_UNWRAP(false, "DeterministicRng: uniform_index 上界必须 > 0"); - std::uniform_int_distribution dist(0, - upper_exclusive - 1); + if (upper_exclusive == 0) NN_UNWRAP(false, "DeterministicRng: uniform_index 上界必须 > 0"); + std::uniform_int_distribution dist(0, upper_exclusive - 1); return dist(engine_); } @@ -65,11 +57,9 @@ class DeterministicRng std::string tag; std::uint64_t saved_seed = 0; in >> tag >> saved_seed; - if (!in || tag != "RNGS1") - NN_FAIL("DeterministicRng::restore_state: 状态头非法"); + if (!in || tag != "RNGS1") NN_FAIL("DeterministicRng::restore_state: 状态头非法"); in >> engine_; - if (!in) - NN_FAIL("DeterministicRng::restore_state: 状态不完整"); + if (!in) NN_FAIL("DeterministicRng::restore_state: 状态不完整"); seed_ = saved_seed; return {}; } @@ -89,17 +79,11 @@ class ExperienceBuffer std::size_t capacity_ = 0; public: - explicit ExperienceBuffer(std::size_t capacity = 0) - : capacity_(capacity) - { - } + explicit ExperienceBuffer(std::size_t capacity = 0) : capacity_(capacity) {} [[nodiscard]] Result add(TensorRecord record) { - if (capacity_ > 0 && records_.size() >= capacity_) - { - records_.erase(records_.begin()); - } + if (capacity_ > 0 && records_.size() >= capacity_) { records_.erase(records_.begin()); } records_.push_back(std::move(record)); return {}; } @@ -111,15 +95,13 @@ class ExperienceBuffer [[nodiscard]] Result at(std::size_t index) { - if (index >= records_.size()) - NN_FAIL("ExperienceBuffer::at: 下标越界"); + if (index >= records_.size()) NN_FAIL("ExperienceBuffer::at: 下标越界"); return &records_[index]; } [[nodiscard]] Result at(std::size_t index) const { - if (index >= records_.size()) - NN_FAIL("ExperienceBuffer::at: 下标越界"); + if (index >= records_.size()) NN_FAIL("ExperienceBuffer::at: 下标越界"); return &records_[index]; } }; @@ -128,9 +110,8 @@ class BatchSampler { public: virtual ~BatchSampler() = default; - virtual Result> sample( - std::size_t available, std::size_t batch_size, - DeterministicRng& rng) = 0; + virtual Result> sample(std::size_t available, std::size_t batch_size, + DeterministicRng& rng) = 0; }; class UniformBatchSampler final : public BatchSampler @@ -139,24 +120,19 @@ class UniformBatchSampler final : public BatchSampler bool replacement_ = false; public: - explicit UniformBatchSampler(bool replacement = false) - : replacement_(replacement) - { - } + explicit UniformBatchSampler(bool replacement = false) : replacement_(replacement) {} - [[nodiscard]] Result> sample( - std::size_t available, std::size_t batch_size, - DeterministicRng& rng) override + [[nodiscard]] Result> + sample(std::size_t available, std::size_t batch_size, DeterministicRng& rng) override { - if (batch_size == 0) - NN_FAIL("UniformBatchSampler: batch_size 必须 > 0"); - if (available == 0) - NN_FAIL("UniformBatchSampler: buffer 为空"); + if (batch_size == 0) NN_FAIL("UniformBatchSampler: batch_size 必须 > 0"); + if (available == 0) NN_FAIL("UniformBatchSampler: buffer 为空"); if (!replacement_ && batch_size > available) NN_FAIL("UniformBatchSampler: 无放回采样时 batch_size > buffer size"); std::vector pool(available); - for (std::size_t i = 0; i < available; ++i) pool[i] = i; + for (std::size_t i = 0; i < available; ++i) + pool[i] = i; std::vector out; out.reserve(batch_size); if (replacement_) @@ -191,10 +167,7 @@ class MetricSink class NullMetricSink final : public MetricSink { public: - [[nodiscard]] Result emit(const MetricRecord&) override - { - return {}; - } + [[nodiscard]] Result emit(const MetricRecord&) override { return {}; } }; class CollectingMetricSink final : public MetricSink @@ -209,10 +182,7 @@ class CollectingMetricSink final : public MetricSink return {}; } - [[nodiscard]] const std::vector& records() const noexcept - { - return records_; - } + [[nodiscard]] const std::vector& records() const noexcept { return records_; } void clear() noexcept { records_.clear(); } }; @@ -220,8 +190,7 @@ class CheckpointProvider { public: virtual ~CheckpointProvider() = default; - virtual Result save_section( - std::string_view name, std::string_view bytes) = 0; + virtual Result save_section(std::string_view name, std::string_view bytes) = 0; virtual Result load_section(std::string_view name) const = 0; }; @@ -231,11 +200,9 @@ class InMemoryCheckpointProvider final : public CheckpointProvider std::vector> sections_; public: - [[nodiscard]] Result save_section( - std::string_view name, std::string_view bytes) override + [[nodiscard]] Result save_section(std::string_view name, std::string_view bytes) override { - if (name.empty()) - NN_FAIL("InMemoryCheckpointProvider: section 名不能为空"); + if (name.empty()) NN_FAIL("InMemoryCheckpointProvider: section 名不能为空"); for (auto& section : sections_) { if (section.first == name) @@ -248,22 +215,18 @@ class InMemoryCheckpointProvider final : public CheckpointProvider return {}; } - [[nodiscard]] Result load_section( - std::string_view name) const override + [[nodiscard]] Result load_section(std::string_view name) const override { for (const auto& section : sections_) if (section.first == name) return section.second; - NN_FAIL("InMemoryCheckpointProvider: 缺少 section '" + - std::string(name) + "'"); + NN_FAIL("InMemoryCheckpointProvider: 缺少 section '" + std::string(name) + "'"); } }; struct TrainerHooks { - std::function(std::span)> - before_update; - std::function(std::span)> - after_update; + std::function(std::span)> before_update; + std::function(std::span)> after_update; }; struct RLPipelineConfig @@ -290,19 +253,12 @@ class RLTrainingPipeline std::uint64_t update_count_ = 0; std::uint64_t rollout_cursor_ = 0; - RLTrainingPipeline(RLPipelineConfig cfg, - std::unique_ptr source, - std::unique_ptr sampler, - std::shared_ptr metrics, - std::shared_ptr checkpoints, - std::uint64_t seed) - : cfg_(std::move(cfg)), - source_(std::move(source)), - sampler_(std::move(sampler)), - metrics_(std::move(metrics)), - checkpoints_(std::move(checkpoints)), - buffer_(cfg_.buffer_capacity), - rng_(seed) + RLTrainingPipeline(RLPipelineConfig cfg, std::unique_ptr source, + std::unique_ptr sampler, std::shared_ptr metrics, + std::shared_ptr checkpoints, std::uint64_t seed) + : cfg_(std::move(cfg)), source_(std::move(source)), sampler_(std::move(sampler)), + metrics_(std::move(metrics)), checkpoints_(std::move(checkpoints)), + buffer_(cfg_.buffer_capacity), rng_(seed) { } @@ -313,23 +269,17 @@ class RLTrainingPipeline state.set("update_count", update_count_); state.set("rollout_cursor", rollout_cursor_); state.set("rng_state", rng_.save_state()); - return checkpoints_->save_section( - cfg_.checkpoint_section, - state.serialize()); + return checkpoints_->save_section(cfg_.checkpoint_section, state.serialize()); } public: - [[nodiscard]] static Result create( - RLPipelineConfig cfg, std::unique_ptr source, - std::unique_ptr sampler, - std::shared_ptr metrics = {}, - std::shared_ptr checkpoints = {}, - std::uint64_t seed = 0) + [[nodiscard]] static Result + create(RLPipelineConfig cfg, std::unique_ptr source, + std::unique_ptr sampler, std::shared_ptr metrics = {}, + std::shared_ptr checkpoints = {}, std::uint64_t seed = 0) { - if (!source || !sampler) - NN_FAIL("RLTrainingPipeline::create: source/sampler 不能为空"); - if (cfg.batch_size == 0 || cfg.collect_per_update == 0 || - cfg.min_buffer_size == 0) + if (!source || !sampler) NN_FAIL("RLTrainingPipeline::create: source/sampler 不能为空"); + if (cfg.batch_size == 0 || cfg.collect_per_update == 0 || cfg.min_buffer_size == 0) NN_FAIL("RLTrainingPipeline::create: batch/collect/min_buffer 必须 > 0"); if (cfg.min_buffer_size < cfg.batch_size) NN_FAIL("RLTrainingPipeline::create: min_buffer_size 不能小于 batch_size"); @@ -338,13 +288,10 @@ class RLTrainingPipeline " buffer_capacity"); if (cfg.checkpoint_section.empty()) NN_FAIL("RLTrainingPipeline::create: checkpoint_section 不能为空"); - if (!metrics) - metrics = std::make_shared(); - if (!checkpoints) - checkpoints = std::make_shared(); - return RLTrainingPipeline(std::move(cfg), std::move(source), - std::move(sampler), std::move(metrics), - std::move(checkpoints), seed); + if (!metrics) metrics = std::make_shared(); + if (!checkpoints) checkpoints = std::make_shared(); + return RLTrainingPipeline(std::move(cfg), std::move(source), std::move(sampler), + std::move(metrics), std::move(checkpoints), seed); } RLTrainingPipeline(RLTrainingPipeline&&) noexcept = default; @@ -353,36 +300,17 @@ class RLTrainingPipeline RLTrainingPipeline& operator=(const RLTrainingPipeline&) = delete; ~RLTrainingPipeline() = default; - [[nodiscard]] const RLPipelineConfig& config() const noexcept - { - return cfg_; - } - [[nodiscard]] const ExperienceBuffer& buffer() const noexcept - { - return buffer_; - } - [[nodiscard]] std::uint64_t update_count() const noexcept - { - return update_count_; - } - [[nodiscard]] std::uint64_t rollout_cursor() const noexcept - { - return rollout_cursor_; - } - void set_rollout_cursor(std::uint64_t cursor) noexcept - { - rollout_cursor_ = cursor; - } + [[nodiscard]] const RLPipelineConfig& config() const noexcept { return cfg_; } + [[nodiscard]] const ExperienceBuffer& buffer() const noexcept { return buffer_; } + [[nodiscard]] std::uint64_t update_count() const noexcept { return update_count_; } + [[nodiscard]] std::uint64_t rollout_cursor() const noexcept { return rollout_cursor_; } + void set_rollout_cursor(std::uint64_t cursor) noexcept { rollout_cursor_ = cursor; } [[nodiscard]] DeterministicRng& rng() noexcept { return rng_; } [[nodiscard]] MetricSink& metrics() noexcept { return *metrics_; } - [[nodiscard]] CheckpointProvider& checkpoints() noexcept - { - return *checkpoints_; - } + [[nodiscard]] CheckpointProvider& checkpoints() noexcept { return *checkpoints_; } template - [[nodiscard]] Result run(std::size_t updates, UpdateFn&& update, - TrainerHooks hooks = {}) + [[nodiscard]] Result run(std::size_t updates, UpdateFn&& update, TrainerHooks hooks = {}) { for (std::size_t i = 0; i < updates; ++i) { @@ -394,8 +322,7 @@ class RLTrainingPipeline } if (buffer_.size() < cfg_.min_buffer_size) NN_FAIL("RLTrainingPipeline: buffer 未达到 min_buffer_size"); - NN_TRY(indices, - sampler_->sample(buffer_.size(), cfg_.batch_size, rng_)); + NN_TRY(indices, sampler_->sample(buffer_.size(), cfg_.batch_size, rng_)); if (indices->size() != cfg_.batch_size) NN_FAIL("RLTrainingPipeline: sampler 返回 batch size 不符"); @@ -406,8 +333,7 @@ class RLTrainingPipeline NN_TRY(record, buffer_.at(index)); selected.push_back(*record); } - if (hooks.before_update) - NN_TRY_CHECK(hooks.before_update(selected)); + if (hooks.before_update) NN_TRY_CHECK(hooks.before_update(selected)); NN_TRY_CHECK(update(selected, rng_, *metrics_)); // update 已成功修改外部模型；即使后置 hook 失败，也必须提交 @@ -415,20 +341,15 @@ class RLTrainingPipeline ++update_count_; Result after_r{}; if (hooks.after_update) after_r = hooks.after_update(selected); - if (cfg_.clear_buffer_after_update) - buffer_.clear(); + if (cfg_.clear_buffer_after_update) buffer_.clear(); NN_TRY_CHECK(after_r); - if (cfg_.checkpoint_every > 0 && - update_count_ % cfg_.checkpoint_every == 0) + if (cfg_.checkpoint_every > 0 && update_count_ % cfg_.checkpoint_every == 0) NN_TRY_CHECK(save_pipeline_state_()); } return {}; } - [[nodiscard]] Result save_checkpoint() - { - return save_pipeline_state_(); - } + [[nodiscard]] Result save_checkpoint() { return save_pipeline_state_(); } [[nodiscard]] Result restore_checkpoint() { @@ -439,8 +360,7 @@ class RLTrainingPipeline NN_FAIL("RLTrainingPipeline::restore_checkpoint: 版本不兼容"); std::uint64_t updates = 0, cursor = 0; std::string rng_state; - if (!kv->get("update_count", updates) || - !kv->get("rollout_cursor", cursor) || + if (!kv->get("update_count", updates) || !kv->get("rollout_cursor", cursor) || !kv->get("rng_state", rng_state)) NN_FAIL("RLTrainingPipeline::restore_checkpoint: 状态字段缺失"); NN_TRY_CHECK(rng_.restore_state(rng_state)); diff --git a/include/neuralnet.cpp/rl_policy.hpp b/include/neuralnet.cpp/rl_policy.hpp index e1cffbda..8a51aa70 100644 --- a/include/neuralnet.cpp/rl_policy.hpp +++ b/include/neuralnet.cpp/rl_policy.hpp @@ -44,20 +44,18 @@ class PolicyDistribution }; virtual ~PolicyDistribution() = default; - virtual Result log_prob( - ComputeEngine& engine, const Tensor& parameters, - const Tensor& actions) const = 0; - virtual Result entropy( - ComputeEngine& engine, const Tensor& parameters) const = 0; - virtual Result log_prob_with_gradient( - ComputeEngine& engine, const Tensor& parameters, - const Tensor& actions, const Tensor& upstream) const = 0; - virtual Result entropy_with_gradient( - ComputeEngine& engine, const Tensor& parameters, - const Tensor& upstream) const = 0; - virtual Result sample( - ComputeEngine& engine, const Tensor& parameters, - std::uint64_t seed) const = 0; + virtual Result log_prob(ComputeEngine& engine, const Tensor& parameters, + const Tensor& actions) const = 0; + virtual Result entropy(ComputeEngine& engine, const Tensor& parameters) const = 0; + virtual Result log_prob_with_gradient(ComputeEngine& engine, + const Tensor& parameters, + const Tensor& actions, + const Tensor& upstream) const = 0; + virtual Result entropy_with_gradient(ComputeEngine& engine, + const Tensor& parameters, + const Tensor& upstream) const = 0; + virtual Result sample(ComputeEngine& engine, const Tensor& parameters, + std::uint64_t seed) const = 0; }; class CategoricalDistribution final : public PolicyDistribution @@ -71,9 +69,8 @@ class CategoricalDistribution final : public PolicyDistribution Precision precision_ = Precision::F32; - [[nodiscard]] Result validate_actions( - ComputeEngine& engine, const Tensor& logits, - const Tensor& actions) const + [[nodiscard]] Result validate_actions(ComputeEngine& engine, const Tensor& logits, + const Tensor& actions) const { if (logits.rows() == 0 || logits.cols() == 0) NN_FAIL("CategoricalDistribution: logits 不能为空"); @@ -83,163 +80,118 @@ class CategoricalDistribution final : public PolicyDistribution for (std::size_t i = 0; i < v->size(); ++i) { const Scalar a = (*v)[i]; - if (!std::isfinite(a) || a < Scalar{0} || - a >= static_cast(logits.rows()) || + if (!std::isfinite(a) || a < Scalar{0} || a >= static_cast(logits.rows()) || std::fabs(a - std::round(a)) > Scalar{1e-4}) - NN_FAIL("CategoricalDistribution: action[" + - std::to_string(i) + "]=" + std::to_string(a) + - " 越界或不是整数（动作数=" + - std::to_string(logits.rows()) + "）"); + NN_FAIL("CategoricalDistribution: action[" + std::to_string(i) + + "]=" + std::to_string(a) + + " 越界或不是整数（动作数=" + std::to_string(logits.rows()) + "）"); } return {}; } - [[nodiscard]] Result probabilities( - ComputeEngine& engine, const Tensor& logits) const + [[nodiscard]] Result probabilities(ComputeEngine& engine, + const Tensor& logits) const { if (logits.rows() == 0 || logits.cols() == 0) NN_FAIL("CategoricalDistribution: logits 不能为空"); - NN_TRY(col_max, - dsl::compute_reduce(engine, - dsl::col_reduce_max(dsl::leaf(logits)), - logits.rows(), logits.cols(), - precision_)); - NN_TRY(col_sum, - dsl::compute_reduce( - engine, - dsl::col_reduce_sum( - dsl::exp(dsl::leaf(logits) - - dsl::col_broadcast(*col_max))), - logits.rows(), logits.cols(), precision_)); - NN_TRY(shifted, - dsl::compute(engine, - dsl::leaf(logits) - - dsl::col_broadcast(*col_max), - logits.rows(), logits.cols(), precision_)); + NN_TRY(col_max, dsl::compute_reduce(engine, dsl::col_reduce_max(dsl::leaf(logits)), + logits.rows(), logits.cols(), precision_)); + NN_TRY(col_sum, dsl::compute_reduce(engine, + dsl::col_reduce_sum(dsl::exp( + dsl::leaf(logits) - dsl::col_broadcast(*col_max))), + logits.rows(), logits.cols(), precision_)); + NN_TRY(shifted, dsl::compute(engine, dsl::leaf(logits) - dsl::col_broadcast(*col_max), + logits.rows(), logits.cols(), precision_)); NN_TRY(softmax, - dsl::compute(engine, - dsl::exp(dsl::leaf(*shifted)) / - dsl::col_broadcast(*col_sum), + dsl::compute(engine, dsl::exp(dsl::leaf(*shifted)) / dsl::col_broadcast(*col_sum), logits.rows(), logits.cols(), precision_)); - NN_TRY(log_col_sum, - dsl::compute(engine, dsl::log(dsl::leaf(*col_sum)), - col_sum->rows(), col_sum->cols(), precision_)); + NN_TRY(log_col_sum, dsl::compute(engine, dsl::log(dsl::leaf(*col_sum)), col_sum->rows(), + col_sum->cols(), precision_)); NN_TRY(log_softmax, - dsl::compute(engine, - dsl::leaf(*shifted) - - dsl::col_broadcast(*log_col_sum), + dsl::compute(engine, dsl::leaf(*shifted) - dsl::col_broadcast(*log_col_sum), logits.rows(), logits.cols(), precision_)); return Probabilities{.softmax = std::move(*softmax), .log_softmax = std::move(*log_softmax)}; } public: - explicit CategoricalDistribution(Precision precision = Precision::F32) - : precision_(precision) + explicit CategoricalDistribution(Precision precision = Precision::F32) : precision_(precision) { } - [[nodiscard]] Result log_prob( - ComputeEngine& engine, const Tensor& logits, - const Tensor& actions) const override + [[nodiscard]] Result log_prob(ComputeEngine& engine, const Tensor& logits, + const Tensor& actions) const override { NN_TRY_CHECK(validate_actions(engine, logits, actions)); NN_TRY(probs, probabilities(engine, logits)); return dsl::compute_reduce( engine, - dsl::col_reduce_sum( - dsl::select(dsl::row() == - dsl::col_broadcast(actions), - dsl::leaf(probs->log_softmax), - Scalar{0})), + dsl::col_reduce_sum(dsl::select(dsl::row() == dsl::col_broadcast(actions), + dsl::leaf(probs->log_softmax), Scalar{0})), logits.rows(), logits.cols(), precision_); } - [[nodiscard]] Result entropy( - ComputeEngine& engine, const Tensor& logits) const override + [[nodiscard]] Result entropy(ComputeEngine& engine, const Tensor& logits) const override { NN_TRY(probs, probabilities(engine, logits)); - NN_TRY(term, - dsl::compute(engine, - dsl::leaf(probs->softmax) * - dsl::leaf(probs->log_softmax), - logits.rows(), logits.cols(), precision_)); - NN_TRY(sum, - dsl::compute_reduce( - engine, dsl::col_reduce_sum(dsl::leaf(*term)), - logits.rows(), logits.cols(), precision_)); - return dsl::compute(engine, dsl::leaf(*sum) * dsl::rparam(-1), - sum->rows(), sum->cols(), precision_); + NN_TRY(term, dsl::compute(engine, dsl::leaf(probs->softmax) * dsl::leaf(probs->log_softmax), + logits.rows(), logits.cols(), precision_)); + NN_TRY(sum, dsl::compute_reduce(engine, dsl::col_reduce_sum(dsl::leaf(*term)), + logits.rows(), logits.cols(), precision_)); + return dsl::compute(engine, dsl::leaf(*sum) * dsl::rparam(-1), sum->rows(), sum->cols(), + precision_); } - [[nodiscard]] Result log_prob_with_gradient( - ComputeEngine& engine, const Tensor& logits, const Tensor& actions, - const Tensor& upstream) const override + [[nodiscard]] Result log_prob_with_gradient(ComputeEngine& engine, + const Tensor& logits, + const Tensor& actions, + const Tensor& upstream) const override { NN_TRY_CHECK(validate_actions(engine, logits, actions)); if (upstream.rows() != 1 || upstream.cols() != logits.cols()) NN_FAIL("CategoricalDistribution: log_prob upstream 必须为 (1,T)"); NN_TRY(probs, probabilities(engine, logits)); - NN_TRY(value, - dsl::compute_reduce( - engine, - dsl::col_reduce_sum( - dsl::select( - dsl::row() == dsl::col_broadcast(actions), - dsl::leaf(probs->log_softmax), Scalar{0})), - logits.rows(), logits.cols(), precision_)); - NN_TRY(grad, - dsl::compute( - engine, - (dsl::select( - dsl::row() == dsl::col_broadcast(actions), - Scalar{1}, Scalar{0}) - - dsl::leaf(probs->softmax)) * - dsl::col_broadcast(upstream), - logits.rows(), logits.cols(), precision_)); - return Evaluation{.value = std::move(*value), - .gradient = std::move(*grad)}; + NN_TRY(value, dsl::compute_reduce(engine, + dsl::col_reduce_sum(dsl::select( + dsl::row() == dsl::col_broadcast(actions), + dsl::leaf(probs->log_softmax), Scalar{0})), + logits.rows(), logits.cols(), precision_)); + NN_TRY(grad, dsl::compute(engine, + (dsl::select(dsl::row() == dsl::col_broadcast(actions), Scalar{1}, + Scalar{0}) - + dsl::leaf(probs->softmax)) * + dsl::col_broadcast(upstream), + logits.rows(), logits.cols(), precision_)); + return Evaluation{.value = std::move(*value), .gradient = std::move(*grad)}; } - [[nodiscard]] Result entropy_with_gradient( - ComputeEngine& engine, const Tensor& logits, - const Tensor& upstream) const override + [[nodiscard]] Result entropy_with_gradient(ComputeEngine& engine, + const Tensor& logits, + const Tensor& upstream) const override { if (upstream.rows() != 1 || upstream.cols() != logits.cols()) NN_FAIL("CategoricalDistribution: entropy upstream 必须为 (1,T)"); NN_TRY(probs, probabilities(engine, logits)); - NN_TRY(term, - dsl::compute(engine, - dsl::leaf(probs->softmax) * - dsl::leaf(probs->log_softmax), - logits.rows(), logits.cols(), precision_)); - NN_TRY(sum, - dsl::compute_reduce( - engine, dsl::col_reduce_sum(dsl::leaf(*term)), - logits.rows(), logits.cols(), precision_)); - NN_TRY(value, - dsl::compute(engine, dsl::leaf(*sum) * dsl::rparam(-1), - sum->rows(), sum->cols(), precision_)); - NN_TRY(grad, - dsl::compute( - engine, - (dsl::leaf(probs->softmax) * - (dsl::leaf(probs->log_softmax) + - dsl::col_broadcast(*value))) * - dsl::col_broadcast(upstream) * dsl::rparam(-1), - logits.rows(), logits.cols(), precision_)); - return Evaluation{.value = std::move(*value), - .gradient = std::move(*grad)}; + NN_TRY(term, dsl::compute(engine, dsl::leaf(probs->softmax) * dsl::leaf(probs->log_softmax), + logits.rows(), logits.cols(), precision_)); + NN_TRY(sum, dsl::compute_reduce(engine, dsl::col_reduce_sum(dsl::leaf(*term)), + logits.rows(), logits.cols(), precision_)); + NN_TRY(value, dsl::compute(engine, dsl::leaf(*sum) * dsl::rparam(-1), sum->rows(), + sum->cols(), precision_)); + NN_TRY(grad, dsl::compute(engine, + (dsl::leaf(probs->softmax) * + (dsl::leaf(probs->log_softmax) + dsl::col_broadcast(*value))) * + dsl::col_broadcast(upstream) * dsl::rparam(-1), + logits.rows(), logits.cols(), precision_)); + return Evaluation{.value = std::move(*value), .gradient = std::move(*grad)}; } - [[nodiscard]] Result sample( - ComputeEngine& engine, const Tensor& logits, - std::uint64_t seed) const override + [[nodiscard]] Result sample(ComputeEngine& engine, const Tensor& logits, + std::uint64_t seed) const override { if (logits.rows() == 0 || logits.cols() == 0) NN_FAIL("CategoricalDistribution::sample: logits 不能为空"); - NN_TRY(values, - detail::download_vector(engine, logits)); + NN_TRY(values, detail::download_vector(engine, logits)); std::mt19937_64 rng(seed); std::uniform_real_distribution uniform(0.0, 1.0); std::vector actions(logits.cols(), Scalar{0}); @@ -252,8 +204,7 @@ class CategoricalDistribution final : public PolicyDistribution std::vector probs(logits.rows()); for (std::size_t r = 0; r < logits.rows(); ++r) { - probs[r] = std::exp(static_cast( - (*values)[r * logits.cols() + c] - mx)); + probs[r] = std::exp(static_cast((*values)[r * logits.cols() + c] - mx)); total += probs[r]; } const double u = uniform(rng) * total; @@ -270,13 +221,11 @@ class CategoricalDistribution final : public PolicyDistribution } actions[c] = static_cast(chosen); } - NN_TRY(action, - detail::upload_span(engine, 1, logits.cols(), precision_, - std::span(actions))); + NN_TRY(action, detail::upload_span(engine, 1, logits.cols(), precision_, + std::span(actions))); NN_TRY_CHECK(validate_actions(engine, logits, *action)); NN_TRY(lp, log_prob(engine, logits, *action)); - return SampleResult{.action = std::move(*action), - .log_prob = std::move(*lp)}; + return SampleResult{.action = std::move(*action), .log_prob = std::move(*lp)}; } }; @@ -292,23 +241,20 @@ class ContextualObjective : public Loss Tensor cached_grad_; protected: - [[nodiscard]] virtual Result evaluate( - ComputeEngine& engine, const Tensor& prediction, - const PolicyBatchView& batch) = 0; + [[nodiscard]] virtual Result + evaluate(ComputeEngine& engine, const Tensor& prediction, const PolicyBatchView& batch) = 0; [[nodiscard]] Precision precision() const noexcept { return p_.stable; } public: - [[nodiscard]] Result forward( - ComputeEngine&, const Tensor&, const Tensor&) override + [[nodiscard]] Result forward(ComputeEngine&, const Tensor&, const Tensor&) override { NN_FAIL("ContextualObjective::forward: policy/value objective 需要强类型" " PolicyBatchView；请调用 forward(engine, prediction, batch)"); } - [[nodiscard]] Result forward( - ComputeEngine& engine, const Tensor& prediction, - const PolicyBatchView& batch) + [[nodiscard]] Result forward(ComputeEngine& engine, const Tensor& prediction, + const PolicyBatchView& batch) { auto result = evaluate(engine, prediction, batch); if (!result) @@ -321,12 +267,9 @@ class ContextualObjective : public Loss } // 供组合目标复用组件，而不依赖其隐式缓存。 - [[nodiscard]] Result evaluate_context( - ComputeEngine& engine, const Tensor& prediction, - const PolicyBatchView& batch) - { - return evaluate(engine, prediction, batch); - } + [[nodiscard]] Result + evaluate_context(ComputeEngine& engine, const Tensor& prediction, const PolicyBatchView& batch) + { return evaluate(engine, prediction, batch); } [[nodiscard]] Result backward() override { @@ -348,32 +291,26 @@ namespace detail return {}; } -[[nodiscard]] inline Result rl_masked_units( - ComputeEngine& engine, const Tensor& mask, Precision p) +[[nodiscard]] inline Result rl_masked_units(ComputeEngine& engine, const Tensor& mask, + Precision p) { - NN_TRY(sum, dsl::compute_reduce( - engine, dsl::col_reduce_sum(dsl::leaf(mask)), - mask.rows(), mask.cols(), p)); - NN_TRY(total, dsl::compute_reduce( - engine, dsl::row_reduce_sum(dsl::leaf(*sum)), - 1, mask.cols(), p)); + NN_TRY(sum, dsl::compute_reduce(engine, dsl::col_reduce_sum(dsl::leaf(mask)), mask.rows(), + mask.cols(), p)); + NN_TRY(total, + dsl::compute_reduce(engine, dsl::row_reduce_sum(dsl::leaf(*sum)), 1, mask.cols(), p)); NN_TRY(v, download_vector(engine, *total)); return (*v)[0]; } -[[nodiscard]] inline Result rl_masked_sum( - ComputeEngine& engine, const Tensor& values, const Tensor& mask, - Precision p) +[[nodiscard]] inline Result rl_masked_sum(ComputeEngine& engine, const Tensor& values, + const Tensor& mask, Precision p) { - NN_TRY(weighted, dsl::compute(engine, - dsl::leaf(values) * dsl::leaf(mask), - values.rows(), values.cols(), p)); - NN_TRY(col, dsl::compute_reduce( - engine, dsl::col_reduce_sum(dsl::leaf(*weighted)), - values.rows(), values.cols(), p)); - NN_TRY(total, dsl::compute_reduce( - engine, dsl::row_reduce_sum(dsl::leaf(*col)), - 1, values.cols(), p)); + NN_TRY(weighted, dsl::compute(engine, dsl::leaf(values) * dsl::leaf(mask), values.rows(), + values.cols(), p)); + NN_TRY(col, dsl::compute_reduce(engine, dsl::col_reduce_sum(dsl::leaf(*weighted)), + values.rows(), values.cols(), p)); + NN_TRY(total, + dsl::compute_reduce(engine, dsl::row_reduce_sum(dsl::leaf(*col)), 1, values.cols(), p)); NN_TRY(v, download_vector(engine, *total)); return (*v)[0]; } @@ -386,15 +323,12 @@ class ReinforceObjective final : public ContextualObjective std::shared_ptr distribution_; public: - explicit ReinforceObjective( - bool use_old_value_baseline = false, - Precision precision = Precision::F32, - std::shared_ptr distribution = {}) + explicit ReinforceObjective(bool use_old_value_baseline = false, + Precision precision = Precision::F32, + std::shared_ptr distribution = {}) : use_old_value_baseline_(use_old_value_baseline), - distribution_(distribution - ? std::move(distribution) - : std::make_shared( - precision)) + distribution_(distribution ? std::move(distribution) + : std::make_shared(precision)) { PrecisionProfile p; p.stable = precision; @@ -402,9 +336,8 @@ class ReinforceObjective final : public ContextualObjective } protected: - [[nodiscard]] Result evaluate( - ComputeEngine& engine, const Tensor& logits, - const PolicyBatchView& batch) override + [[nodiscard]] Result evaluate(ComputeEngine& engine, const Tensor& logits, + const PolicyBatchView& batch) override { NN_TRY_CHECK(detail::rl_check_policy_row(logits, &batch.action(), "action")); NN_TRY_CHECK(detail::rl_check_policy_row(logits, &batch.action_mask(), "action_mask")); @@ -419,51 +352,37 @@ class ReinforceObjective final : public ContextualObjective NN_FAIL("ReinforceObjective: baseline 需要 return 与 old_value"); NN_TRY_CHECK(detail::rl_check_policy_row(logits, returns, "return")); NN_TRY_CHECK(detail::rl_check_policy_row(logits, old_value, "old_value")); - NN_TRY(adv, dsl::compute(engine, - dsl::leaf(*returns) - - dsl::leaf(*old_value), - returns->rows(), returns->cols(), - precision())); + NN_TRY(adv, dsl::compute(engine, dsl::leaf(*returns) - dsl::leaf(*old_value), + returns->rows(), returns->cols(), precision())); baseline_adv = std::move(*adv); advantage = &baseline_adv; } if (advantage == nullptr) NN_FAIL("ReinforceObjective: 缺少 advantage（或启用 old_value baseline）"); const Tensor& mask = batch.action_mask(); - if (advantage->rows() != mask.rows() || - advantage->cols() != mask.cols()) + if (advantage->rows() != mask.rows() || advantage->cols() != mask.cols()) NN_FAIL("ReinforceObjective: advantage/action_mask 形状不一致"); - NN_TRY(units_r, - detail::rl_masked_units(engine, mask, precision())); + NN_TRY(units_r, detail::rl_masked_units(engine, mask, precision())); const Scalar units = *units_r; if (units <= Scalar{0}) { - NN_TRY(zero_grad, - detail::rl_zero(engine, logits.rows(), logits.cols(), - precision())); - return ObjectiveOutput{.loss = 0, - .gradient = std::move(*zero_grad)}; + NN_TRY(zero_grad, detail::rl_zero(engine, logits.rows(), logits.cols(), precision())); + return ObjectiveOutput{.loss = 0, .gradient = std::move(*zero_grad)}; } - NN_TRY(logp, - distribution_->log_prob(engine, logits, batch.action())); - NN_TRY(weighted_adv, - dsl::compute(engine, - dsl::leaf(*advantage) * dsl::leaf(mask) * - dsl::rparam(-Scalar{1} / units), - 1, logits.cols(), precision())); - NN_TRY(eval, - distribution_->log_prob_with_gradient( - engine, logits, batch.action(), *weighted_adv)); + NN_TRY(logp, distribution_->log_prob(engine, logits, batch.action())); + NN_TRY(weighted_adv, dsl::compute(engine, + dsl::leaf(*advantage) * dsl::leaf(mask) * + dsl::rparam(-Scalar{1} / units), + 1, logits.cols(), precision())); + NN_TRY(eval, distribution_->log_prob_with_gradient(engine, logits, batch.action(), + *weighted_adv)); NN_TRY(loss_terms, dsl::compute(engine, - dsl::leaf(eval->value) * dsl::leaf(*advantage) * - dsl::leaf(mask), - 1, logits.cols(), precision())); - NN_TRY(loss_sum, - detail::rl_masked_sum(engine, *loss_terms, mask, - precision())); + dsl::leaf(eval->value) * dsl::leaf(*advantage) * dsl::leaf(mask), 1, + logits.cols(), precision())); + NN_TRY(loss_sum, detail::rl_masked_sum(engine, *loss_terms, mask, precision())); return ObjectiveOutput{ .loss = -*loss_sum / units, .gradient = std::move(eval->gradient), @@ -477,8 +396,7 @@ class ValueRegressionObjective final : public ContextualObjective Scalar coefficient_ = 1; public: - explicit ValueRegressionObjective(Scalar coefficient = 1, - Precision precision = Precision::F32) + explicit ValueRegressionObjective(Scalar coefficient = 1, Precision precision = Precision::F32) : coefficient_(coefficient) { PrecisionProfile p; @@ -487,44 +405,31 @@ class ValueRegressionObjective final : public ContextualObjective } protected: - [[nodiscard]] Result evaluate( - ComputeEngine& engine, const Tensor& value, - const PolicyBatchView& batch) override + [[nodiscard]] Result evaluate(ComputeEngine& engine, const Tensor& value, + const PolicyBatchView& batch) override { const Tensor* returns = batch.returns(); const Tensor* mask = batch.value_mask(); - if (returns == nullptr) - NN_FAIL("ValueRegressionObjective: 缺少 return"); - if (mask == nullptr) - NN_FAIL("ValueRegressionObjective: 缺少显式 value_mask"); - if (returns->rows() != value.rows() || - returns->cols() != value.cols() || - mask->rows() != value.rows() || - mask->cols() != value.cols()) + if (returns == nullptr) NN_FAIL("ValueRegressionObjective: 缺少 return"); + if (mask == nullptr) NN_FAIL("ValueRegressionObjective: 缺少显式 value_mask"); + if (returns->rows() != value.rows() || returns->cols() != value.cols() || + mask->rows() != value.rows() || mask->cols() != value.cols()) NN_FAIL("ValueRegressionObjective: value/return/mask 形状不一致"); NN_TRY(units_r, detail::rl_masked_units(engine, *mask, precision())); const Scalar units = *units_r; if (units <= Scalar{0}) { - NN_TRY(zero_grad, - detail::rl_zero(engine, value.rows(), value.cols(), - precision())); - return ObjectiveOutput{.loss = 0, - .gradient = std::move(*zero_grad)}; + NN_TRY(zero_grad, detail::rl_zero(engine, value.rows(), value.cols(), precision())); + return ObjectiveOutput{.loss = 0, .gradient = std::move(*zero_grad)}; } - NN_TRY(diff, dsl::compute(engine, - dsl::leaf(value) - dsl::leaf(*returns), - value.rows(), value.cols(), precision())); - NN_TRY(square, dsl::compute(engine, - dsl::leaf(*diff) * dsl::leaf(*diff), - value.rows(), value.cols(), precision())); - NN_TRY(loss_sum, - detail::rl_masked_sum(engine, *square, *mask, precision())); + NN_TRY(diff, dsl::compute(engine, dsl::leaf(value) - dsl::leaf(*returns), value.rows(), + value.cols(), precision())); + NN_TRY(square, dsl::compute(engine, dsl::leaf(*diff) * dsl::leaf(*diff), value.rows(), + value.cols(), precision())); + NN_TRY(loss_sum, detail::rl_masked_sum(engine, *square, *mask, precision())); NN_TRY(grad, dsl::compute(engine, - dsl::leaf(*diff) * - dsl::leaf(*mask) * - dsl::rparam(Scalar{2} * - coefficient_ / units), + dsl::leaf(*diff) * dsl::leaf(*mask) * + dsl::rparam(Scalar{2} * coefficient_ / units), value.rows(), value.cols(), precision())); return ObjectiveOutput{ .loss = coefficient_ * *loss_sum / units, @@ -540,15 +445,11 @@ class EntropyBonusObjective final : public ContextualObjective std::shared_ptr distribution_; public: - explicit EntropyBonusObjective( - Scalar coefficient = 0.01f, - Precision precision = Precision::F32, - std::shared_ptr distribution = {}) + explicit EntropyBonusObjective(Scalar coefficient = 0.01f, Precision precision = Precision::F32, + std::shared_ptr distribution = {}) : coefficient_(coefficient), - distribution_(distribution - ? std::move(distribution) - : std::make_shared( - precision)) + distribution_(distribution ? std::move(distribution) + : std::make_shared(precision)) { PrecisionProfile p; p.stable = precision; @@ -556,35 +457,23 @@ class EntropyBonusObjective final : public ContextualObjective } protected: - [[nodiscard]] Result evaluate( - ComputeEngine& engine, const Tensor& logits, - const PolicyBatchView& batch) override + [[nodiscard]] Result evaluate(ComputeEngine& engine, const Tensor& logits, + const PolicyBatchView& batch) override { const Tensor& mask = batch.action_mask(); NN_TRY_CHECK(detail::rl_check_policy_row(logits, &batch.action(), "action")); NN_TRY_CHECK(detail::rl_check_policy_row(logits, &mask, "action_mask")); - NN_TRY(units_r, - detail::rl_masked_units(engine, mask, precision())); + NN_TRY(units_r, detail::rl_masked_units(engine, mask, precision())); const Scalar units = *units_r; if (units <= Scalar{0}) { - NN_TRY(zero_grad, - detail::rl_zero(engine, logits.rows(), logits.cols(), - precision())); - return ObjectiveOutput{.loss = 0, - .gradient = std::move(*zero_grad)}; + NN_TRY(zero_grad, detail::rl_zero(engine, logits.rows(), logits.cols(), precision())); + return ObjectiveOutput{.loss = 0, .gradient = std::move(*zero_grad)}; } - NN_TRY(upstream, - dsl::compute(engine, - dsl::leaf(mask) * - dsl::rparam(-coefficient_ / units), - 1, logits.cols(), precision())); - NN_TRY(eval, - distribution_->entropy_with_gradient( - engine, logits, *upstream)); - NN_TRY(total_h, - detail::rl_masked_sum(engine, eval->value, mask, - precision())); + NN_TRY(upstream, dsl::compute(engine, dsl::leaf(mask) * dsl::rparam(-coefficient_ / units), + 1, logits.cols(), precision())); + NN_TRY(eval, distribution_->entropy_with_gradient(engine, logits, *upstream)); + NN_TRY(total_h, detail::rl_masked_sum(engine, eval->value, mask, precision())); return ObjectiveOutput{ .loss = -coefficient_ * *total_h / units, .gradient = std::move(eval->gradient), @@ -599,15 +488,11 @@ class SampledKlObjective final : public ContextualObjective std::shared_ptr distribution_; public: - explicit SampledKlObjective( - Scalar coefficient = 1, - Precision precision = Precision::F32, - std::shared_ptr distribution = {}) + explicit SampledKlObjective(Scalar coefficient = 1, Precision precision = Precision::F32, + std::shared_ptr distribution = {}) : coefficient_(coefficient), - distribution_(distribution - ? std::move(distribution) - : std::make_shared( - precision)) + distribution_(distribution ? std::move(distribution) + : std::make_shared(precision)) { PrecisionProfile p; p.stable = precision; @@ -615,42 +500,29 @@ class SampledKlObjective final : public ContextualObjective } protected: - [[nodiscard]] Result evaluate( - ComputeEngine& engine, const Tensor& logits, - const PolicyBatchView& batch) override + [[nodiscard]] Result evaluate(ComputeEngine& engine, const Tensor& logits, + const PolicyBatchView& batch) override { const Tensor* old_logp = batch.old_log_prob(); - if (old_logp == nullptr) - NN_FAIL("SampledKlObjective: 缺少 old_log_prob"); + if (old_logp == nullptr) NN_FAIL("SampledKlObjective: 缺少 old_log_prob"); const Tensor& mask = batch.action_mask(); NN_TRY_CHECK(detail::rl_check_policy_row(logits, &batch.action(), "action")); NN_TRY_CHECK(detail::rl_check_policy_row(logits, old_logp, "old_log_prob")); NN_TRY_CHECK(detail::rl_check_policy_row(logits, &mask, "action_mask")); - NN_TRY(units_r, - detail::rl_masked_units(engine, mask, precision())); + NN_TRY(units_r, detail::rl_masked_units(engine, mask, precision())); const Scalar units = *units_r; if (units <= Scalar{0}) { - NN_TRY(zero_grad, - detail::rl_zero(engine, logits.rows(), logits.cols(), - precision())); - return ObjectiveOutput{.loss = 0, - .gradient = std::move(*zero_grad)}; + NN_TRY(zero_grad, detail::rl_zero(engine, logits.rows(), logits.cols(), precision())); + return ObjectiveOutput{.loss = 0, .gradient = std::move(*zero_grad)}; } - NN_TRY(upstream, - dsl::compute(engine, - dsl::leaf(mask) * - dsl::rparam(coefficient_ / units), - 1, logits.cols(), precision())); + NN_TRY(upstream, dsl::compute(engine, dsl::leaf(mask) * dsl::rparam(coefficient_ / units), + 1, logits.cols(), precision())); NN_TRY(eval, - distribution_->log_prob_with_gradient( - engine, logits, batch.action(), *upstream)); - NN_TRY(diff, - dsl::compute(engine, - dsl::leaf(eval->value) - dsl::leaf(*old_logp), - 1, logits.cols(), precision())); - NN_TRY(total, - detail::rl_masked_sum(engine, *diff, mask, precision())); + distribution_->log_prob_with_gradient(engine, logits, batch.action(), *upstream)); + NN_TRY(diff, dsl::compute(engine, dsl::leaf(eval->value) - dsl::leaf(*old_logp), 1, + logits.cols(), precision())); + NN_TRY(total, detail::rl_masked_sum(engine, *diff, mask, precision())); return ObjectiveOutput{ .loss = coefficient_ * *total / units, .gradient = std::move(eval->gradient), @@ -674,101 +546,73 @@ class ClippedPolicyObjective final : public ContextualObjective PpoPolicyMetrics metrics_; public: - explicit ClippedPolicyObjective( - Scalar clip_epsilon = 0.2f, - Precision precision = Precision::F32, - std::shared_ptr distribution = {}) + explicit ClippedPolicyObjective(Scalar clip_epsilon = 0.2f, + Precision precision = Precision::F32, + std::shared_ptr distribution = {}) : clip_epsilon_(clip_epsilon), - distribution_(distribution - ? std::move(distribution) - : std::make_shared( - precision)) + distribution_(distribution ? std::move(distribution) + : std::make_shared(precision)) { PrecisionProfile p; p.stable = precision; set_precision_profile(p); } - [[nodiscard]] const PpoPolicyMetrics& metrics() const noexcept - { - return metrics_; - } + [[nodiscard]] const PpoPolicyMetrics& metrics() const noexcept { return metrics_; } protected: - [[nodiscard]] Result evaluate( - ComputeEngine& engine, const Tensor& logits, - const PolicyBatchView& batch) override + [[nodiscard]] Result evaluate(ComputeEngine& engine, const Tensor& logits, + const PolicyBatchView& batch) override { const Tensor* old_logp = batch.old_log_prob(); const Tensor* advantage = batch.advantage(); - if (old_logp == nullptr) - NN_FAIL("ClippedPolicyObjective: 缺少 old_log_prob"); - if (advantage == nullptr) - NN_FAIL("ClippedPolicyObjective: 缺少 advantage"); + if (old_logp == nullptr) NN_FAIL("ClippedPolicyObjective: 缺少 old_log_prob"); + if (advantage == nullptr) NN_FAIL("ClippedPolicyObjective: 缺少 advantage"); const Tensor& mask = batch.action_mask(); NN_TRY_CHECK(detail::rl_check_policy_row(logits, &batch.action(), "action")); NN_TRY_CHECK(detail::rl_check_policy_row(logits, old_logp, "old_log_prob")); NN_TRY_CHECK(detail::rl_check_policy_row(logits, advantage, "advantage")); NN_TRY_CHECK(detail::rl_check_policy_row(logits, &mask, "action_mask")); - NN_TRY(units_r, - detail::rl_masked_units(engine, mask, precision())); + NN_TRY(units_r, detail::rl_masked_units(engine, mask, precision())); const Scalar units = *units_r; if (units <= Scalar{0}) { metrics_ = {}; - NN_TRY(zero_grad, - detail::rl_zero(engine, logits.rows(), logits.cols(), - precision())); - return ObjectiveOutput{.loss = 0, - .gradient = std::move(*zero_grad)}; + NN_TRY(zero_grad, detail::rl_zero(engine, logits.rows(), logits.cols(), precision())); + return ObjectiveOutput{.loss = 0, .gradient = std::move(*zero_grad)}; } - NN_TRY(new_logp, - distribution_->log_prob(engine, logits, batch.action())); + NN_TRY(new_logp, distribution_->log_prob(engine, logits, batch.action())); NN_TRY_MSG(diff, "PPO policy log-ratio", - dsl::compute(engine, - dsl::leaf(*new_logp) - - dsl::leaf(*old_logp), - 1, logits.cols(), precision())); + dsl::compute(engine, dsl::leaf(*new_logp) - dsl::leaf(*old_logp), 1, + logits.cols(), precision())); NN_TRY_MSG(ratio, "PPO policy ratio", - dsl::compute(engine, dsl::exp(dsl::leaf(*diff)), - 1, logits.cols(), precision())); + dsl::compute(engine, dsl::exp(dsl::leaf(*diff)), 1, logits.cols(), precision())); const Scalar lower = Scalar{1} - clip_epsilon_; const Scalar upper = Scalar{1} + clip_epsilon_; - NN_TRY_MSG(clipped_ratio, "PPO policy clipped ratio", - dsl::compute( - engine, - dsl::select( - dsl::leaf(*ratio) < dsl::rparam(lower), - dsl::rparam(lower), - dsl::select( - dsl::leaf(*ratio) > dsl::rparam(upper), - dsl::rparam(upper), dsl::leaf(*ratio))), - 1, logits.cols(), precision())); + NN_TRY_MSG( + clipped_ratio, "PPO policy clipped ratio", + dsl::compute(engine, + dsl::select(dsl::leaf(*ratio) < dsl::rparam(lower), dsl::rparam(lower), + dsl::select(dsl::leaf(*ratio) > dsl::rparam(upper), + dsl::rparam(upper), dsl::leaf(*ratio))), + 1, logits.cols(), precision())); NN_TRY_MSG(objective, "PPO policy ratio*advantage", - dsl::compute(engine, - dsl::leaf(*ratio) * - dsl::leaf(*advantage), - 1, logits.cols(), precision())); + dsl::compute(engine, dsl::leaf(*ratio) * dsl::leaf(*advantage), 1, logits.cols(), + precision())); NN_TRY_MSG(clipped_objective, "PPO clipped ratio*advantage", - dsl::compute(engine, - dsl::leaf(*clipped_ratio) * - dsl::leaf(*advantage), - 1, logits.cols(), precision())); + dsl::compute(engine, dsl::leaf(*clipped_ratio) * dsl::leaf(*advantage), 1, + logits.cols(), precision())); NN_TRY_MSG(surrogate, "PPO policy min", - dsl::compute( - engine, - dsl::min(dsl::leaf(*objective), - dsl::leaf(*clipped_objective)), - 1, logits.cols(), precision())); - NN_TRY_MSG(loss_terms, "PPO policy masked surrogate", dsl::compute(engine, - dsl::leaf(*surrogate) * dsl::leaf(mask), - 1, logits.cols(), precision())); + dsl::min(dsl::leaf(*objective), dsl::leaf(*clipped_objective)), 1, + logits.cols(), precision())); + NN_TRY_MSG(loss_terms, "PPO policy masked surrogate", + dsl::compute(engine, dsl::leaf(*surrogate) * dsl::leaf(mask), 1, logits.cols(), + precision())); NN_TRY_MSG(loss_sum, "PPO policy reduce", - detail::rl_masked_sum(engine, *loss_terms, mask, - precision())); + detail::rl_masked_sum(engine, *loss_terms, mask, precision())); const auto ratio_v = dsl::leaf(*ratio); const auto adv_v = dsl::leaf(*advantage); @@ -776,48 +620,35 @@ class ClippedPolicyObjective final : public ContextualObjective // A>0 时只裁上界，A<0 时只裁下界。 const auto positive = dsl::select(ratio_v > dsl::rparam(upper), Scalar{0}, ratio_v * adv_v); const auto negative = dsl::select(ratio_v < dsl::rparam(lower), Scalar{0}, ratio_v * adv_v); - const auto factor = dsl::select( - adv_v > Scalar{0}, positive, negative); + const auto factor = dsl::select(adv_v > Scalar{0}, positive, negative); NN_TRY_MSG(factor_tensor, "PPO policy factor", - dsl::compute(engine, factor, 1, logits.cols(), - precision())); + dsl::compute(engine, factor, 1, logits.cols(), precision())); auto upstream_expr = - dsl::leaf(*factor_tensor) * dsl::leaf(mask) * - dsl::rparam(-Scalar{1} / units); + dsl::leaf(*factor_tensor) * dsl::leaf(mask) * dsl::rparam(-Scalar{1} / units); NN_TRY_MSG(upstream, "PPO policy upstream", - dsl::compute(engine, upstream_expr, - 1, logits.cols(), precision())); - NN_TRY_MSG(eval, "PPO policy log-prob gradient", - distribution_->log_prob_with_gradient( - engine, logits, batch.action(), *upstream)); + dsl::compute(engine, upstream_expr, 1, logits.cols(), precision())); + NN_TRY_MSG( + eval, "PPO policy log-prob gradient", + distribution_->log_prob_with_gradient(engine, logits, batch.action(), *upstream)); NN_TRY_MSG(ratio_weighted, "PPO ratio weighted", - dsl::compute(engine, - ratio_v * dsl::leaf(mask), - 1, logits.cols(), precision())); + dsl::compute(engine, ratio_v * dsl::leaf(mask), 1, logits.cols(), precision())); NN_TRY_MSG(ratio_sum, "PPO ratio reduce", - detail::rl_masked_sum(engine, *ratio_weighted, mask, - precision())); + detail::rl_masked_sum(engine, *ratio_weighted, mask, precision())); NN_TRY_MSG(kl_terms, "PPO KL terms", + dsl::compute(engine, dsl::leaf(*diff) * dsl::leaf(*diff) * dsl::leaf(mask), 1, + logits.cols(), precision())); + NN_TRY_MSG(kl_sum, "PPO KL reduce", + detail::rl_masked_sum(engine, *kl_terms, mask, precision())); + NN_TRY_MSG(outside, "PPO clip indicator", dsl::compute(engine, - dsl::leaf(*diff) * dsl::leaf(*diff) * + dsl::select(ratio_v < dsl::rparam(lower), Scalar{1}, + dsl::select(ratio_v > dsl::rparam(upper), Scalar{1}, + Scalar{0})) * dsl::leaf(mask), 1, logits.cols(), precision())); - NN_TRY_MSG(kl_sum, "PPO KL reduce", - detail::rl_masked_sum(engine, *kl_terms, mask, - precision())); - NN_TRY_MSG(outside, "PPO clip indicator", - dsl::compute( - engine, - dsl::select( - ratio_v < dsl::rparam(lower), Scalar{1}, - dsl::select(ratio_v > dsl::rparam(upper), - Scalar{1}, Scalar{0})) * - dsl::leaf(mask), - 1, logits.cols(), precision())); NN_TRY_MSG(outside_sum, "PPO clip reduce", - detail::rl_masked_sum(engine, *outside, mask, - precision())); + detail::rl_masked_sum(engine, *outside, mask, precision())); metrics_.policy_loss = -*loss_sum / units; metrics_.ratio_mean = *ratio_sum / units; @@ -844,9 +675,8 @@ class ClippedValueObjective final : public ContextualObjective PpoValueMetrics metrics_; public: - explicit ClippedValueObjective( - Scalar clip_epsilon = 0.2f, Scalar coefficient = 0.5f, - Precision precision = Precision::F32) + explicit ClippedValueObjective(Scalar clip_epsilon = 0.2f, Scalar coefficient = 0.5f, + Precision precision = Precision::F32) : clip_epsilon_(clip_epsilon), coefficient_(coefficient) { PrecisionProfile p; @@ -854,128 +684,85 @@ class ClippedValueObjective final : public ContextualObjective set_precision_profile(p); } - [[nodiscard]] const PpoValueMetrics& metrics() const noexcept - { - return metrics_; - } + [[nodiscard]] const PpoValueMetrics& metrics() const noexcept { return metrics_; } protected: - [[nodiscard]] Result evaluate( - ComputeEngine& engine, const Tensor& value, - const PolicyBatchView& batch) override + [[nodiscard]] Result evaluate(ComputeEngine& engine, const Tensor& value, + const PolicyBatchView& batch) override { const Tensor* returns = batch.returns(); const Tensor* old_value = batch.old_value(); const Tensor* mask = batch.value_mask(); - if (returns == nullptr) - NN_FAIL("ClippedValueObjective: 缺少 return"); - if (old_value == nullptr) - NN_FAIL("ClippedValueObjective: 缺少 old_value"); - if (mask == nullptr) - NN_FAIL("ClippedValueObjective: 缺少 value_mask"); - if (returns->rows() != value.rows() || - returns->cols() != value.cols() || - old_value->rows() != value.rows() || - old_value->cols() != value.cols() || - mask->rows() != value.rows() || - mask->cols() != value.cols()) + if (returns == nullptr) NN_FAIL("ClippedValueObjective: 缺少 return"); + if (old_value == nullptr) NN_FAIL("ClippedValueObjective: 缺少 old_value"); + if (mask == nullptr) NN_FAIL("ClippedValueObjective: 缺少 value_mask"); + if (returns->rows() != value.rows() || returns->cols() != value.cols() || + old_value->rows() != value.rows() || old_value->cols() != value.cols() || + mask->rows() != value.rows() || mask->cols() != value.cols()) NN_FAIL("ClippedValueObjective: value/return/old_value/mask 形状不一致"); - NN_TRY(units_r, - detail::rl_masked_units(engine, *mask, precision())); + NN_TRY(units_r, detail::rl_masked_units(engine, *mask, precision())); const Scalar units = *units_r; if (units <= Scalar{0}) { metrics_ = {}; - NN_TRY(zero_grad, - detail::rl_zero(engine, value.rows(), value.cols(), - precision())); - return ObjectiveOutput{.loss = 0, - .gradient = std::move(*zero_grad)}; + NN_TRY(zero_grad, detail::rl_zero(engine, value.rows(), value.cols(), precision())); + return ObjectiveOutput{.loss = 0, .gradient = std::move(*zero_grad)}; } - NN_TRY(clip_low, - dsl::compute(engine, - dsl::leaf(*old_value) - - dsl::rparam(clip_epsilon_), - value.rows(), value.cols(), precision())); - NN_TRY(clip_high, - dsl::compute(engine, - dsl::leaf(*old_value) + - dsl::rparam(clip_epsilon_), - value.rows(), value.cols(), precision())); - NN_TRY_MSG(clipped, "PPO value clip", - dsl::compute( - engine, - dsl::select( - dsl::leaf(value) < dsl::leaf(*clip_low), - dsl::leaf(*clip_low), - dsl::select( - dsl::leaf(value) > dsl::leaf(*clip_high), - dsl::leaf(*clip_high), dsl::leaf(value))), - value.rows(), value.cols(), precision())); + NN_TRY(clip_low, dsl::compute(engine, dsl::leaf(*old_value) - dsl::rparam(clip_epsilon_), + value.rows(), value.cols(), precision())); + NN_TRY(clip_high, dsl::compute(engine, dsl::leaf(*old_value) + dsl::rparam(clip_epsilon_), + value.rows(), value.cols(), precision())); + NN_TRY_MSG( + clipped, "PPO value clip", + dsl::compute(engine, + dsl::select(dsl::leaf(value) < dsl::leaf(*clip_low), dsl::leaf(*clip_low), + dsl::select(dsl::leaf(value) > dsl::leaf(*clip_high), + dsl::leaf(*clip_high), dsl::leaf(value))), + value.rows(), value.cols(), precision())); NN_TRY_MSG(diff, "PPO value diff", - dsl::compute(engine, - dsl::leaf(value) - dsl::leaf(*returns), - value.rows(), value.cols(), precision())); + dsl::compute(engine, dsl::leaf(value) - dsl::leaf(*returns), value.rows(), + value.cols(), precision())); NN_TRY_MSG(clipped_diff, "PPO value clipped diff", - dsl::compute(engine, - dsl::leaf(*clipped) - dsl::leaf(*returns), - value.rows(), value.cols(), precision())); + dsl::compute(engine, dsl::leaf(*clipped) - dsl::leaf(*returns), value.rows(), + value.cols(), precision())); NN_TRY_MSG(diff_sq, "PPO value diff square", - dsl::compute(engine, - dsl::leaf(*diff) * dsl::leaf(*diff), - value.rows(), value.cols(), precision())); + dsl::compute(engine, dsl::leaf(*diff) * dsl::leaf(*diff), value.rows(), + value.cols(), precision())); NN_TRY_MSG(clipped_sq, "PPO value clipped square", - dsl::compute(engine, - dsl::leaf(*clipped_diff) * - dsl::leaf(*clipped_diff), + dsl::compute(engine, dsl::leaf(*clipped_diff) * dsl::leaf(*clipped_diff), value.rows(), value.cols(), precision())); - auto loss_expr = dsl::max( - dsl::leaf(*diff_sq), dsl::leaf(*clipped_sq)) * - dsl::leaf(*mask); + auto loss_expr = dsl::max(dsl::leaf(*diff_sq), dsl::leaf(*clipped_sq)) * dsl::leaf(*mask); NN_TRY_MSG(loss_terms, "PPO value max loss", - dsl::compute(engine, loss_expr, value.rows(), - value.cols(), precision())); + dsl::compute(engine, loss_expr, value.rows(), value.cols(), precision())); NN_TRY_MSG(loss_sum, "PPO value reduce", - detail::rl_masked_sum(engine, *loss_terms, *mask, - precision())); + detail::rl_masked_sum(engine, *loss_terms, *mask, precision())); const auto outside_clip = dsl::select( dsl::leaf(value) < dsl::leaf(*clip_low), Scalar{1}, - dsl::select(dsl::leaf(value) > dsl::leaf(*clip_high), - Scalar{1}, Scalar{0})); + dsl::select(dsl::leaf(value) > dsl::leaf(*clip_high), Scalar{1}, Scalar{0})); NN_TRY_MSG(outside_tensor, "PPO value outside tensor", - dsl::compute(engine, outside_clip, value.rows(), - value.cols(), precision())); + dsl::compute(engine, outside_clip, value.rows(), value.cols(), precision())); NN_TRY_MSG(clipped_grad_tensor, "PPO value clipped grad", dsl::compute(engine, - dsl::select( - dsl::leaf(*outside_tensor) > Scalar{0}, - Scalar{0}, - dsl::leaf(*diff) * dsl::rparam(2)), + dsl::select(dsl::leaf(*outside_tensor) > Scalar{0}, Scalar{0}, + dsl::leaf(*diff) * dsl::rparam(2)), value.rows(), value.cols(), precision())); NN_TRY_MSG(unclipped_grad_tensor, "PPO value unclipped grad", - dsl::compute(engine, - dsl::leaf(*diff) * dsl::rparam(2), - value.rows(), value.cols(), precision())); - auto choose_clipped = - dsl::leaf(*clipped_sq) > dsl::leaf(*diff_sq); - auto grad_expr = dsl::select( - choose_clipped, dsl::leaf(*clipped_grad_tensor), - dsl::leaf(*unclipped_grad_tensor)) * - dsl::leaf(*mask) * - dsl::rparam(coefficient_ / units); - NN_TRY_MSG(grad, "PPO value gradient", - dsl::compute(engine, grad_expr, value.rows(), + dsl::compute(engine, dsl::leaf(*diff) * dsl::rparam(2), value.rows(), value.cols(), precision())); + auto choose_clipped = dsl::leaf(*clipped_sq) > dsl::leaf(*diff_sq); + auto grad_expr = dsl::select(choose_clipped, dsl::leaf(*clipped_grad_tensor), + dsl::leaf(*unclipped_grad_tensor)) * + dsl::leaf(*mask) * dsl::rparam(coefficient_ / units); + NN_TRY_MSG(grad, "PPO value gradient", + dsl::compute(engine, grad_expr, value.rows(), value.cols(), precision())); NN_TRY_MSG(clip_terms, "PPO value clip indicator", - dsl::compute(engine, - outside_clip * dsl::leaf(*mask), - value.rows(), value.cols(), precision())); + dsl::compute(engine, outside_clip * dsl::leaf(*mask), value.rows(), value.cols(), + precision())); NN_TRY_MSG(clip_sum, "PPO value clip reduce", - detail::rl_masked_sum(engine, *clip_terms, *mask, - precision())); + detail::rl_masked_sum(engine, *clip_terms, *mask, precision())); metrics_.value_loss = coefficient_ * *loss_sum / units; metrics_.clip_fraction = *clip_sum / units; return ObjectiveOutput{ @@ -1015,50 +802,36 @@ class PpoObjective PpoConfig cfg_; public: - explicit PpoObjective(PpoConfig cfg = {}) - : cfg_(std::move(cfg)) - { - } + explicit PpoObjective(PpoConfig cfg = {}) : cfg_(std::move(cfg)) {} - [[nodiscard]] Result evaluate( - ComputeEngine& engine, const TensorSchema& output_schema, - const Tensor& policy_logits, const Tensor& value, - const PolicyBatchView& batch) const + [[nodiscard]] Result evaluate(ComputeEngine& engine, + const TensorSchema& output_schema, + const Tensor& policy_logits, const Tensor& value, + const PolicyBatchView& batch) const { NN_TRY_CHECK(output_schema.field(cfg_.policy_field)); NN_TRY_CHECK(output_schema.field(cfg_.value_field)); - ClippedPolicyObjective policy(cfg_.clip_epsilon, cfg_.precision, - cfg_.distribution); - EntropyBonusObjective entropy(cfg_.entropy_coefficient, - cfg_.precision, cfg_.distribution); - ClippedValueObjective value_loss( - cfg_.value_clip_epsilon, cfg_.value_coefficient, - cfg_.precision); - NN_TRY(policy_out, - policy.evaluate_context(engine, policy_logits, batch)); - NN_TRY(entropy_out, - entropy.evaluate_context(engine, policy_logits, batch)); - NN_TRY(value_out, - value_loss.evaluate_context(engine, value, batch)); - - NN_TRY(policy_grad, - engine.clone(policy_out->gradient)); - NN_TRY_CHECK(engine.accumulate(*policy_grad, - entropy_out->gradient)); + ClippedPolicyObjective policy(cfg_.clip_epsilon, cfg_.precision, cfg_.distribution); + EntropyBonusObjective entropy(cfg_.entropy_coefficient, cfg_.precision, cfg_.distribution); + ClippedValueObjective value_loss(cfg_.value_clip_epsilon, cfg_.value_coefficient, + cfg_.precision); + NN_TRY(policy_out, policy.evaluate_context(engine, policy_logits, batch)); + NN_TRY(entropy_out, entropy.evaluate_context(engine, policy_logits, batch)); + NN_TRY(value_out, value_loss.evaluate_context(engine, value, batch)); + + NN_TRY(policy_grad, engine.clone(policy_out->gradient)); + NN_TRY_CHECK(engine.accumulate(*policy_grad, entropy_out->gradient)); NN_TRY(gradients, GradientBundle::create( output_schema, - {NamedTensor{.name = cfg_.policy_field, - .value = *policy_grad}, - NamedTensor{.name = cfg_.value_field, - .value = value_out->gradient}})); + {NamedTensor{.name = cfg_.policy_field, .value = *policy_grad}, + NamedTensor{.name = cfg_.value_field, .value = value_out->gradient}})); const PpoPolicyMetrics& pm = policy.metrics(); const PpoValueMetrics& vm = value_loss.metrics(); - const Scalar entropy_value = - cfg_.entropy_coefficient == Scalar{0} - ? Scalar{0} - : -entropy_out->loss / cfg_.entropy_coefficient; + const Scalar entropy_value = cfg_.entropy_coefficient == Scalar{0} + ? Scalar{0} + : -entropy_out->loss / cfg_.entropy_coefficient; return PpoReport{ .loss = policy_out->loss + entropy_out->loss + value_out->loss, .policy_loss = policy_out->loss + entropy_out->loss, @@ -1095,8 +868,8 @@ class GrpoObjective private: GrpoConfig cfg_; - [[nodiscard]] static Result with_advantage( - const TensorRecord& source, const Tensor& advantage) + [[nodiscard]] static Result with_advantage(const TensorRecord& source, + const Tensor& advantage) { std::vector values; values.reserve(source.size()); @@ -1104,29 +877,23 @@ class GrpoObjective { if (field.name == "advantage") { - values.push_back(NamedTensor{ - .name = field.name, .value = advantage}); + values.push_back(NamedTensor{.name = field.name, .value = advantage}); continue; } const Tensor* value = source.find(field.name); if (value != nullptr) - values.push_back(NamedTensor{.name = field.name, - .value = *value}); + values.push_back(NamedTensor{.name = field.name, .value = *value}); } return TensorRecord::create(source.schema(), std::move(values)); } public: - explicit GrpoObjective(GrpoConfig cfg = {}) - : cfg_(std::move(cfg)) - { - } + explicit GrpoObjective(GrpoConfig cfg = {}) : cfg_(std::move(cfg)) {} // advantage 已存在时直接使用；否则按 spans 对 returns 做组内标准化。 - [[nodiscard]] Result evaluate( - ComputeEngine& engine, const TensorSchema& output_schema, - const Tensor& policy_logits, const PolicyBatchView& batch, - std::span spans) const + [[nodiscard]] Result + evaluate(ComputeEngine& engine, const TensorSchema& output_schema, const Tensor& policy_logits, + const PolicyBatchView& batch, std::span spans) const { NN_TRY_CHECK(output_schema.field(cfg_.policy_field)); TensorRecord derived; @@ -1134,15 +901,13 @@ class GrpoObjective if (batch.advantage() == nullptr) { const Tensor* returns = batch.returns(); - if (returns == nullptr) - NN_FAIL("GrpoObjective: 缺少 advantage 或 return"); + if (returns == nullptr) NN_FAIL("GrpoObjective: 缺少 advantage 或 return"); Tensor flat_returns; const Tensor* group_input = returns; if (returns->cols() != 1) { - NN_TRY(flat, engine.reshape( - *returns, returns->rows() * returns->cols(), 1)); + NN_TRY(flat, engine.reshape(*returns, returns->rows() * returns->cols(), 1)); flat_returns = std::move(*flat); group_input = &flat_returns; } @@ -1152,9 +917,7 @@ class GrpoObjective Tensor advantage; if (advantage_flat.cols() == 1 && batch.action_mask().rows() == 1) { - NN_TRY(a, - engine.reshape(advantage_flat, 1, - advantage_flat.rows())); + NN_TRY(a, engine.reshape(advantage_flat, 1, advantage_flat.rows())); advantage = std::move(*a); } else @@ -1168,26 +931,18 @@ class GrpoObjective } NN_TRY(active_view, PolicyBatchView::create(*active_record)); - ClippedPolicyObjective policy(cfg_.clip_epsilon, cfg_.precision, - cfg_.distribution); - EntropyBonusObjective entropy(cfg_.entropy_coefficient, - cfg_.precision, cfg_.distribution); - NN_TRY(policy_out, - policy.evaluate_context(engine, policy_logits, *active_view)); - NN_TRY(entropy_out, - entropy.evaluate_context(engine, policy_logits, *active_view)); + ClippedPolicyObjective policy(cfg_.clip_epsilon, cfg_.precision, cfg_.distribution); + EntropyBonusObjective entropy(cfg_.entropy_coefficient, cfg_.precision, cfg_.distribution); + NN_TRY(policy_out, policy.evaluate_context(engine, policy_logits, *active_view)); + NN_TRY(entropy_out, entropy.evaluate_context(engine, policy_logits, *active_view)); NN_TRY(policy_grad, engine.clone(policy_out->gradient)); - NN_TRY_CHECK(engine.accumulate(*policy_grad, - entropy_out->gradient)); + NN_TRY_CHECK(engine.accumulate(*policy_grad, entropy_out->gradient)); NN_TRY(gradients, GradientBundle::create( - output_schema, - {NamedTensor{.name = cfg_.policy_field, - .value = *policy_grad}})); - const Scalar entropy_value = - cfg_.entropy_coefficient == Scalar{0} - ? Scalar{0} - : -entropy_out->loss / cfg_.entropy_coefficient; + output_schema, {NamedTensor{.name = cfg_.policy_field, .value = *policy_grad}})); + const Scalar entropy_value = cfg_.entropy_coefficient == Scalar{0} + ? Scalar{0} + : -entropy_out->loss / cfg_.entropy_coefficient; return GrpoReport{ .loss = policy_out->loss + entropy_out->loss, .policy_loss = policy_out->loss, diff --git a/include/neuralnet.cpp/run_driver.hpp b/include/neuralnet.cpp/run_driver.hpp index 60a3736e..c2785035 100644 --- a/include/neuralnet.cpp/run_driver.hpp +++ b/include/neuralnet.cpp/run_driver.hpp @@ -130,17 +130,13 @@ struct RunDriverConfig class RunDriver { public: - [[nodiscard]] static Result create(TrainingSession& session, - BatchSource& source, + [[nodiscard]] static Result create(TrainingSession& session, BatchSource& source, std::unique_ptr state, RunDriverConfig cfg) { - if (!state) - NN_FAIL("RunDriver::create: state 不能为空（域需提供其 TrainState 子类）"); - if (cfg.epochs < 0) - NN_FAIL("RunDriver::create: epochs 不能为负"); - if (cfg.stop_after_epoch < 0) - NN_FAIL("RunDriver::create: stop_after_epoch 不能为负"); + if (!state) NN_FAIL("RunDriver::create: state 不能为空（域需提供其 TrainState 子类）"); + if (cfg.epochs < 0) NN_FAIL("RunDriver::create: epochs 不能为负"); + if (cfg.stop_after_epoch < 0) NN_FAIL("RunDriver::create: stop_after_epoch 不能为负"); if (cfg.stop_after_epoch > cfg.epochs && cfg.epochs > 0) NN_FAIL("RunDriver::create: stop_after_epoch 超出 epochs"); return RunDriver(session, source, std::move(state), std::move(cfg)); @@ -152,8 +148,7 @@ class RunDriver // 驱动：epoch×step 循环 + lr 接线 + 断点节奏 + 报告。 [[nodiscard]] Result fit() { - if (!hooks_.do_update) - NN_FAIL("RunDriver::fit: do_update 钩子必选（域的单步更新逻辑）"); + if (!hooks_.do_update) NN_FAIL("RunDriver::fit: do_update 钩子必选（域的单步更新逻辑）"); const auto t0 = std::chrono::steady_clock::now(); RunReport report; @@ -164,8 +159,7 @@ class RunDriver // 起点（断点续训 = 计数器"下一步的位置"） const std::size_t start_epoch = static_cast(state_->counters.epoch); - const std::size_t start_step = - static_cast(state_->counters.step_in_epoch); + const std::size_t start_step = static_cast(state_->counters.step_in_epoch); if (epoch_mode && start_epoch > static_cast(cfg_.epochs)) NN_FAIL("RunDriver::fit: 断点 epoch 超出本次配方 epochs"); @@ -176,7 +170,8 @@ class RunDriver epoch_lr_cfg.warmup_epochs = cfg_.warmup_epochs; epoch_lr_cfg.total_epochs = (cfg_.epochs > 0) ? cfg_.epochs : 1; epoch_lr_cfg.schedule = cfg_.lr_schedule; - for (Scalar v : cfg_.lr_per_epoch) epoch_lr_cfg.lr_per_epoch.push_back(v); + for (Scalar v : cfg_.lr_per_epoch) + epoch_lr_cfg.lr_per_epoch.push_back(v); // lr 调度配置（step 级；step_cosine 用） StepLrScheduleConfig step_lr_cfg; @@ -184,9 +179,9 @@ class RunDriver step_lr_cfg.min_lr = cfg_.min_lr; step_lr_cfg.warmup_steps = static_cast(cfg_.warmup_steps); step_lr_cfg.total_steps = static_cast( - (steps_per_epoch > 0) ? steps_per_epoch * static_cast( - (cfg_.epochs > 0) ? cfg_.epochs : 1) - : cfg_.max_steps); + (steps_per_epoch > 0) + ? steps_per_epoch * static_cast((cfg_.epochs > 0) ? cfg_.epochs : 1) + : cfg_.max_steps); step_lr_cfg.cosine = (cfg_.lr_schedule == "step_cosine"); Scalar current_lr = cfg_.lr; @@ -215,9 +210,9 @@ class RunDriver // 绝对步号 = epoch·steps_per_epoch + step（与 step 级调度同定义； // 跨续训累计，进程无关）。步进驱动退化为 step 本身。 auto abs_global_step = [steps_per_epoch](std::size_t epoch, - std::size_t step) -> std::uint64_t { - return static_cast(epoch) * - static_cast(steps_per_epoch) + + std::size_t step) -> std::uint64_t + { + return static_cast(epoch) * static_cast(steps_per_epoch) + static_cast(step); }; @@ -257,17 +252,16 @@ class RunDriver const std::size_t epoch_first_step = (static_cast(epoch) == start_epoch) ? start_step : 0; - NN_TRY(rw2, source_.rewind(static_cast(epoch), - epoch_first_step)); + NN_TRY(rw2, source_.rewind(static_cast(epoch), epoch_first_step)); for (std::size_t step = epoch_first_step; step < steps_per_epoch; ++step) { if (cfg_.lr_schedule == "step_cosine") { - const int gstep = static_cast( - static_cast(epoch) * - static_cast(steps_per_epoch) + - static_cast(step)); + const int gstep = + static_cast(static_cast(epoch) * + static_cast(steps_per_epoch) + + static_cast(step)); const Scalar slr = compute_step_lr(step_lr_cfg, gstep); if (slr != current_lr) { @@ -283,8 +277,7 @@ class RunDriver UpdateRequest req; req.epoch = static_cast(epoch); req.step = step; - req.global_step = - abs_global_step(static_cast(epoch), step); + req.global_step = abs_global_step(static_cast(epoch), step); req.lr = current_lr; req.is_final_step = is_final; req.steps_per_epoch = steps_per_epoch; @@ -303,8 +296,7 @@ class RunDriver // save_interval 断点（以 update step 计） if (cfg_.save_interval > 0 && - ((step + 1) % cfg_.save_interval == 0 || - step + 1 == steps_per_epoch)) + ((step + 1) % cfg_.save_interval == 0 || step + 1 == steps_per_epoch)) { NN_TRY(ck, do_checkpoint(static_cast(epoch), step + 1)); } @@ -348,8 +340,7 @@ class RunDriver } // --stop-after-epoch：训到该 epoch（1-based，不含）即退出 - if (cfg_.stop_after_epoch > 0 && - epoch + 1 >= cfg_.stop_after_epoch) + if (cfg_.stop_after_epoch > 0 && epoch + 1 >= cfg_.stop_after_epoch) { stop_training = true; last_epoch = static_cast(epoch) + 1; @@ -443,8 +434,8 @@ class RunDriver [[nodiscard]] const RunDriverConfig& config() const noexcept { return cfg_; } private: - RunDriver(TrainingSession& session, BatchSource& source, - std::unique_ptr state, RunDriverConfig cfg) + RunDriver(TrainingSession& session, BatchSource& source, std::unique_ptr state, + RunDriverConfig cfg) : session_(session), source_(source), state_(std::move(state)), cfg_(std::move(cfg)) { } diff --git a/include/neuralnet.cpp/tensor_schema.hpp b/include/neuralnet.cpp/tensor_schema.hpp index ca313c0c..12fa6e4f 100644 --- a/include/neuralnet.cpp/tensor_schema.hpp +++ b/include/neuralnet.cpp/tensor_schema.hpp @@ -27,8 +27,7 @@ namespace nn { -inline constexpr std::size_t kInvalidFieldIndex = - static_cast(-1); +inline constexpr std::size_t kInvalidFieldIndex = static_cast(-1); inline constexpr std::uint32_t kMaxTensorSchemaFields = 65536; // 字段所在设备约束。使用独立枚举而不是 Device，保持本头不依赖 compute_tensor。 @@ -39,14 +38,16 @@ enum class TensorPlacement : std::uint8_t GPU = 2, }; -[[nodiscard]] inline const char* tensor_placement_name( - TensorPlacement p) noexcept +[[nodiscard]] inline const char* tensor_placement_name(TensorPlacement p) noexcept { switch (p) { - case TensorPlacement::CPU: return "cpu"; - case TensorPlacement::GPU: return "gpu"; - default: return "any"; + case TensorPlacement::CPU: + return "cpu"; + case TensorPlacement::GPU: + return "gpu"; + default: + return "any"; } } @@ -57,22 +58,17 @@ struct TensorShapeSpec std::optional cols; [[nodiscard]] static constexpr TensorShapeSpec any() noexcept { return {}; } - [[nodiscard]] static constexpr TensorShapeSpec exact( - std::size_t rows, std::size_t cols) noexcept - { - return TensorShapeSpec{rows, cols}; - } + [[nodiscard]] static constexpr TensorShapeSpec exact(std::size_t rows, + std::size_t cols) noexcept + { return TensorShapeSpec{rows, cols}; } [[nodiscard]] bool accepts(std::size_t r, std::size_t c) const noexcept - { - return (!rows || *rows == r) && (!cols || *cols == c); - } + { return (!rows || *rows == r) && (!cols || *cols == c); } [[nodiscard]] std::string to_string() const { - auto axis = [](const std::optional& v) { - return v ? std::to_string(*v) : std::string("*"); - }; + auto axis = [](const std::optional& v) + { return v ? std::to_string(*v) : std::string("*"); }; return "(" + axis(rows) + "," + axis(cols) + ")"; } }; @@ -86,12 +82,10 @@ struct TensorFieldSpec TensorPlacement placement = TensorPlacement::Any; }; -[[nodiscard]] inline TensorFieldSpec tensor_field( - std::string name, - TensorShapeSpec shape = {}, - bool required = true, - std::optional precision = std::nullopt, - TensorPlacement placement = TensorPlacement::Any) +[[nodiscard]] inline TensorFieldSpec tensor_field(std::string name, TensorShapeSpec shape = {}, + bool required = true, + std::optional precision = std::nullopt, + TensorPlacement placement = TensorPlacement::Any) { TensorFieldSpec field; field.name = std::move(name); @@ -109,9 +103,7 @@ struct FieldHandle std::size_t index = kInvalidFieldIndex; [[nodiscard]] constexpr bool valid() const noexcept - { - return schema_fingerprint != 0 && index != kInvalidFieldIndex; - } + { return schema_fingerprint != 0 && index != kInvalidFieldIndex; } }; // FNV-1a over schema name + each field's semantic metadata. @@ -120,8 +112,8 @@ inline constexpr std::uint64_t kTensorSchemaFnvPrime = 1099511628211ULL; namespace detail { -[[nodiscard]] inline std::uint64_t schema_hash_bytes( - std::uint64_t h, std::string_view bytes) noexcept +[[nodiscard]] inline std::uint64_t schema_hash_bytes(std::uint64_t h, + std::string_view bytes) noexcept { for (unsigned char c : bytes) { @@ -131,8 +123,7 @@ namespace detail return h; } -[[nodiscard]] inline std::uint64_t schema_hash_u64( - std::uint64_t h, std::uint64_t v) noexcept +[[nodiscard]] inline std::uint64_t schema_hash_u64(std::uint64_t h, std::uint64_t v) noexcept { for (unsigned i = 0; i < 8; ++i) { @@ -149,8 +140,8 @@ namespace detail return schema_hash_bytes(h, value); } -[[nodiscard]] inline std::uint64_t schema_hash_optional_size( - std::uint64_t h, const std::optional& v) noexcept +[[nodiscard]] inline std::uint64_t +schema_hash_optional_size(std::uint64_t h, const std::optional& v) noexcept { h = schema_hash_u64(h, v ? 1u : 0u); return v ? schema_hash_u64(h, static_cast(*v)) : h; @@ -164,16 +155,13 @@ class TensorSchema std::vector fields_; std::uint64_t fingerprint_ = 0; - TensorSchema(std::string name, std::vector fields, - std::uint64_t fingerprint) - : name_(std::move(name)), - fields_(std::move(fields)), - fingerprint_(fingerprint) + TensorSchema(std::string name, std::vector fields, std::uint64_t fingerprint) + : name_(std::move(name)), fields_(std::move(fields)), fingerprint_(fingerprint) { } - [[nodiscard]] static std::uint64_t make_fingerprint_( - std::string_view name, const std::vector& fields) noexcept + [[nodiscard]] static std::uint64_t + make_fingerprint_(std::string_view name, const std::vector& fields) noexcept { std::uint64_t h = kTensorSchemaFnvOffset; h = detail::schema_hash_string(h, name); @@ -185,30 +173,27 @@ class TensorSchema h = detail::schema_hash_optional_size(h, f.shape.cols); h = detail::schema_hash_u64(h, f.required ? 1u : 0u); h = detail::schema_hash_u64( - h, f.precision - ? static_cast(*f.precision) + 1u - : 0u); - h = detail::schema_hash_u64( - h, static_cast(f.placement)); + h, f.precision ? static_cast(*f.precision) + 1u : 0u); + h = detail::schema_hash_u64(h, static_cast(f.placement)); } // 0 保留给"未初始化 handle"。 return h == 0 ? 1u : h; } - [[nodiscard]] Result validate_handle_( - const FieldHandle& handle, std::string_view operation) const + [[nodiscard]] Result validate_handle_(const FieldHandle& handle, + std::string_view operation) const { - if (!handle.valid()) - NN_FAIL(std::string(operation) + ": 字段 handle 未初始化"); + if (!handle.valid()) NN_FAIL(std::string(operation) + ": 字段 handle 未初始化"); if (handle.schema_fingerprint != fingerprint_) - NN_FAIL(std::string(operation) + ": 字段 handle 属于另一个 schema（" - "handle=" + std::to_string(handle.schema_fingerprint) + - "，schema=" + std::to_string(fingerprint_) + "；schema 名=" + - name_ + "）"); + NN_FAIL(std::string(operation) + + ": 字段 handle 属于另一个 schema（" + "handle=" + + std::to_string(handle.schema_fingerprint) + + "，schema=" + std::to_string(fingerprint_) + "；schema 名=" + name_ + "）"); if (handle.index >= fields_.size()) NN_FAIL(std::string(operation) + ": 字段 handle 下标越界（" + - std::to_string(handle.index) + " >= " + - std::to_string(fields_.size()) + "；schema 名=" + name_ + "）"); + std::to_string(handle.index) + " >= " + std::to_string(fields_.size()) + + "；schema 名=" + name_ + "）"); return {}; } @@ -220,22 +205,19 @@ class TensorSchema TensorSchema& operator=(TensorSchema&&) noexcept = default; ~TensorSchema() = default; - [[nodiscard]] static Result create( - std::string name, std::vector fields) + [[nodiscard]] static Result create(std::string name, + std::vector fields) { - if (name.empty()) - NN_FAIL("TensorSchema::create: schema 名不能为空"); - if (fields.empty()) - NN_FAIL("TensorSchema::create: schema '" + name + - "' 至少需要一个字段"); + if (name.empty()) NN_FAIL("TensorSchema::create: schema 名不能为空"); + if (fields.empty()) NN_FAIL("TensorSchema::create: schema '" + name + "' 至少需要一个字段"); std::unordered_set field_names; field_names.reserve(fields.size()); for (std::size_t i = 0; i < fields.size(); ++i) { if (fields[i].name.empty()) - NN_FAIL("TensorSchema::create: schema '" + name + "' 的第 " + - std::to_string(i) + " 个字段名为空"); + NN_FAIL("TensorSchema::create: schema '" + name + "' 的第 " + std::to_string(i) + + " 个字段名为空"); if (!field_names.insert(fields[i].name).second) NN_FAIL("TensorSchema::create: schema '" + name + "' 含重复字段 '" + fields[i].name + "'"); @@ -247,19 +229,12 @@ class TensorSchema [[nodiscard]] bool valid() const noexcept { return fingerprint_ != 0; } [[nodiscard]] const std::string& name() const noexcept { return name_; } - [[nodiscard]] std::uint64_t fingerprint() const noexcept - { - return fingerprint_; - } + [[nodiscard]] std::uint64_t fingerprint() const noexcept { return fingerprint_; } [[nodiscard]] std::size_t size() const noexcept { return fields_.size(); } [[nodiscard]] bool empty() const noexcept { return fields_.empty(); } - [[nodiscard]] const std::vector& fields() const noexcept - { - return fields_; - } + [[nodiscard]] const std::vector& fields() const noexcept { return fields_; } - [[nodiscard]] Result field( - const FieldHandle& handle) const + [[nodiscard]] Result field(const FieldHandle& handle) const { NN_TRY_CHECK(validate_handle_(handle, "TensorSchema::field")); return &fields_[handle.index]; @@ -268,49 +243,39 @@ class TensorSchema // setup/反序列化期使用；训练热路径应缓存返回的 handle。 [[nodiscard]] Result field(std::string_view field_name) const { - if (!valid()) - NN_FAIL("TensorSchema::field: schema 未初始化"); + if (!valid()) NN_FAIL("TensorSchema::field: schema 未初始化"); for (std::size_t i = 0; i < fields_.size(); ++i) { - if (fields_[i].name == field_name) - return FieldHandle{fingerprint_, i}; + if (fields_[i].name == field_name) return FieldHandle{fingerprint_, i}; } - NN_FAIL("TensorSchema::field: schema '" + name_ + "' 不含字段 '" + - std::string(field_name) + "'"); + NN_FAIL("TensorSchema::field: schema '" + name_ + "' 不含字段 '" + std::string(field_name) + + "'"); } - [[nodiscard]] Result optional_field( - std::string_view field_name) const + [[nodiscard]] Result optional_field(std::string_view field_name) const { if (!valid()) return FieldHandle{}; for (std::size_t i = 0; i < fields_.size(); ++i) { - if (fields_[i].name == field_name) - return FieldHandle{fingerprint_, i}; + if (fields_[i].name == field_name) return FieldHandle{fingerprint_, i}; } return FieldHandle{}; } [[nodiscard]] Result validate(const FieldHandle& handle) const - { - return validate_handle_(handle, "TensorSchema::validate"); - } + { return validate_handle_(handle, "TensorSchema::validate"); } [[nodiscard]] bool operator==(const TensorSchema& other) const noexcept { - if (fingerprint_ != other.fingerprint_ || - name_ != other.name_ || + if (fingerprint_ != other.fingerprint_ || name_ != other.name_ || fields_.size() != other.fields_.size()) return false; for (std::size_t i = 0; i < fields_.size(); ++i) { const TensorFieldSpec& a = fields_[i]; const TensorFieldSpec& b = other.fields_[i]; - if (a.name != b.name || - a.shape.rows != b.shape.rows || - a.shape.cols != b.shape.cols || - a.required != b.required || - a.precision != b.precision || + if (a.name != b.name || a.shape.rows != b.shape.rows || a.shape.cols != b.shape.cols || + a.required != b.required || a.precision != b.precision || a.placement != b.placement) return false; } @@ -345,36 +310,28 @@ inline void schema_put_u64(std::string& out, std::uint64_t v) [[nodiscard]] inline Result schema_take_u32(std::string_view& in) { - if (in.size() < 4) - NN_FAIL("TensorSchema::decode: 截断（u32）"); + if (in.size() < 4) NN_FAIL("TensorSchema::decode: 截断（u32）"); std::uint32_t v = 0; for (unsigned i = 0; i < 4; ++i) - v |= static_cast( - static_cast(in[i])) - << (i * 8); + v |= static_cast(static_cast(in[i])) << (i * 8); in.remove_prefix(4); return v; } [[nodiscard]] inline Result schema_take_u64(std::string_view& in) { - if (in.size() < 8) - NN_FAIL("TensorSchema::decode: 截断（u64）"); + if (in.size() < 8) NN_FAIL("TensorSchema::decode: 截断（u64）"); std::uint64_t v = 0; for (unsigned i = 0; i < 8; ++i) - v |= static_cast( - static_cast(in[i])) - << (i * 8); + v |= static_cast(static_cast(in[i])) << (i * 8); in.remove_prefix(8); return v; } -[[nodiscard]] inline Result schema_take_string( - std::string_view& in) +[[nodiscard]] inline Result schema_take_string(std::string_view& in) { NN_TRY(n, schema_take_u32(in)); - if (*n > in.size()) - NN_FAIL("TensorSchema::decode: 字符串长度越界"); + if (*n > in.size()) NN_FAIL("TensorSchema::decode: 字符串长度越界"); std::string out(in.substr(0, *n)); in.remove_prefix(*n); return out; @@ -395,29 +352,24 @@ inline void schema_put_string(std::string& out, std::string_view s) detail::schema_put_string(out, schema.name()); if (schema.size() > std::numeric_limits::max()) NN_UNWRAP(false, "TensorSchema::encode: 字段数过多"); - detail::schema_put_u32( - out, static_cast(schema.size())); + detail::schema_put_u32(out, static_cast(schema.size())); constexpr std::uint64_t no_axis = std::numeric_limits::max(); for (const TensorFieldSpec& f : schema.fields()) { detail::schema_put_string(out, f.name); - detail::schema_put_u64( - out, f.shape.rows ? static_cast(*f.shape.rows) - : no_axis); - detail::schema_put_u64( - out, f.shape.cols ? static_cast(*f.shape.cols) - : no_axis); + detail::schema_put_u64(out, + f.shape.rows ? static_cast(*f.shape.rows) : no_axis); + detail::schema_put_u64(out, + f.shape.cols ? static_cast(*f.shape.cols) : no_axis); out.push_back(static_cast(f.required ? 1 : 0)); - out.push_back(static_cast( - f.precision ? static_cast(*f.precision) + 1u : 0u)); - out.push_back(static_cast( - static_cast(f.placement))); + out.push_back( + static_cast(f.precision ? static_cast(*f.precision) + 1u : 0u)); + out.push_back(static_cast(static_cast(f.placement))); } return out; } -[[nodiscard]] inline Result decode_tensor_schema( - std::string_view bytes) +[[nodiscard]] inline Result decode_tensor_schema(std::string_view bytes) { if (bytes.size() < 4 || bytes.substr(0, 4) != "NNS1") NN_FAIL("TensorSchema::decode: 魔数不匹配"); @@ -438,21 +390,15 @@ inline void schema_put_string(std::string& out, std::string_view s) NN_TRY(field_name, detail::schema_take_string(bytes)); NN_TRY(rows, detail::schema_take_u64(bytes)); NN_TRY(cols, detail::schema_take_u64(bytes)); - if (bytes.size() < 3) - NN_FAIL("TensorSchema::decode: 截断（field flags）"); - const unsigned char required = - static_cast(bytes[0]); - const unsigned char precision_tag = - static_cast(bytes[1]); - const unsigned char placement = - static_cast(bytes[2]); + if (bytes.size() < 3) NN_FAIL("TensorSchema::decode: 截断（field flags）"); + const unsigned char required = static_cast(bytes[0]); + const unsigned char precision_tag = static_cast(bytes[1]); + const unsigned char placement = static_cast(bytes[2]); bytes.remove_prefix(3); - if (required > 1) - NN_FAIL("TensorSchema::decode: required 标志非法"); + if (required > 1) NN_FAIL("TensorSchema::decode: required 标志非法"); if (precision_tag > static_cast(Precision::T1_58) + 1u) NN_FAIL("TensorSchema::decode: precision 标志非法"); - if (placement > - static_cast(TensorPlacement::GPU)) + if (placement > static_cast(TensorPlacement::GPU)) NN_FAIL("TensorSchema::decode: placement 标志非法"); TensorFieldSpec field; @@ -470,21 +416,16 @@ inline void schema_put_string(std::string& out, std::string_view s) field.shape.cols = static_cast(*cols); } field.required = required != 0; - if (precision_tag != 0) - field.precision = - static_cast(precision_tag - 1u); - field.placement = - static_cast(placement); + if (precision_tag != 0) field.precision = static_cast(precision_tag - 1u); + field.placement = static_cast(placement); fields.push_back(std::move(field)); } - if (!bytes.empty()) - NN_FAIL("TensorSchema::decode: 尾部存在未消费字节"); + if (!bytes.empty()) NN_FAIL("TensorSchema::decode: 尾部存在未消费字节"); return TensorSchema::create(std::move(*name), std::move(fields)); } // ── 名称列表（output packing 等纯结构元数据）────────────────────────── -[[nodiscard]] inline std::string encode_tensor_name_list( - const std::vector& names) +[[nodiscard]] inline std::string encode_tensor_name_list(const std::vector& names) { std::string out("NNL1", 4); if (names.size() > std::numeric_limits::max()) @@ -514,8 +455,7 @@ decode_tensor_name_list(std::string_view bytes) NN_TRY(name, detail::schema_take_string(bytes)); names.push_back(std::move(*name)); } - if (!bytes.empty()) - NN_FAIL("decode_tensor_name_list: 尾部存在未消费字节"); + if (!bytes.empty()) NN_FAIL("decode_tensor_name_list: 尾部存在未消费字节"); return names; } diff --git a/include/neuralnet.cpp/ternary_pack.hpp b/include/neuralnet.cpp/ternary_pack.hpp index e6e7ba77..791c6173 100644 --- a/include/neuralnet.cpp/ternary_pack.hpp +++ b/include/neuralnet.cpp/ternary_pack.hpp @@ -40,15 +40,11 @@ inline constexpr std::uint32_t kMask = (1u << kBits) - 1u; // 每行需要的字（f32 元素）数：(K + 15) / 16 [[nodiscard]] inline constexpr std::size_t words_per_row(std::size_t k) noexcept -{ - return (k + kPerWord - 1) / kPerWord; -} +{ return (k + kPerWord - 1) / kPerWord; } // 每行打包后的字节数：4 字节/字 × 字数 = 0.25·K（向上取到 4B 对齐） [[nodiscard]] inline constexpr std::size_t bytes_per_row(std::size_t k) noexcept -{ - return words_per_row(k) * sizeof(std::uint32_t); -} +{ return words_per_row(k) * sizeof(std::uint32_t); } // 量化码：码 = 三值 + 2（∈{1,2,3}）；判据 = |w| > 0.5·τ（**严格**大于： // 恰好 |w| == 0.5·τ 归 0，半值向零）。与 BitLinear::quantize_weights_ 的 @@ -63,15 +59,12 @@ inline constexpr std::uint32_t kMask = (1u << kBits) - 1u; // 码 → 三值（float，精确可表示） [[nodiscard]] inline float value_of(std::uint32_t code) noexcept -{ - return static_cast(static_cast(code & kMask) - 2); -} +{ return static_cast(static_cast(code & kMask) - 2); } // 打包一行：w[0..K)（宿主连续）+ 逐行尺度 τ → words_per_row(K) 个 32-bit 字。 // 尾部（idx ≥ K）填"码 2 = 三值 0"而不是 0：码 0 在映射里无定义，填 0 会让 // 误读尾部的内核拿到 -2（静默错值），填 2 最多是多算一个 0。 -inline void pack_row(const float* w, std::size_t k, float tau, - std::uint32_t* out) noexcept +inline void pack_row(const float* w, std::size_t k, float tau, std::uint32_t* out) noexcept { const std::size_t words = words_per_row(k); for (std::size_t wi = 0; wi < words; ++wi) @@ -127,11 +120,11 @@ inline void pack_row(const float* w, std::size_t k, float tau, } // 打包行（2-bit 容器）→ latent 行（k 个 f32；`out` 为宿主连续缓冲） -inline void dequant_row(const std::uint32_t* words, std::size_t k, float tau, - float* out) noexcept +inline void dequant_row(const std::uint32_t* words, std::size_t k, float tau, float* out) noexcept { const float s = latent_row_scale(words, k, tau); - for (std::size_t i = 0; i < k; ++i) out[i] = at(words, i) * s; + for (std::size_t i = 0; i < k; ++i) + out[i] = at(words, i) * s; } // ═══════════════════════════════════════════════════════════════════════════ @@ -157,14 +150,11 @@ inline constexpr std::uint32_t kBase3PerByte = 5; // 5 个三值/字节 // 每行 base-3 字节数：ceil(K/5) [[nodiscard]] inline constexpr std::size_t base3_bytes_per_row(std::size_t k) noexcept -{ - return (k + kBase3PerByte - 1) / kBase3PerByte; -} +{ return (k + kBase3PerByte - 1) / kBase3PerByte; } // 2-bit 容器行 → base-3 字节（部署落盘方向）。尾部不足 5 个的位置补 digit 0 // （解码按 k 截断，补位不承载信息）。 -inline void pack_row_base3(const std::uint32_t* words, std::size_t k, - std::uint8_t* out) noexcept +inline void pack_row_base3(const std::uint32_t* words, std::size_t k, std::uint8_t* out) noexcept { const std::size_t bytes = base3_bytes_per_row(k); for (std::size_t b = 0; b < bytes; ++b) @@ -174,9 +164,9 @@ inline void pack_row_base3(const std::uint32_t* words, std::size_t k, for (std::uint32_t i = 0; i < kBase3PerByte; ++i) { const std::size_t idx = b * kBase3PerByte + i; - const std::uint32_t code = (idx < k) - ? ((words[idx / kPerWord] >> (kBits * (idx % kPerWord))) & kMask) - : 2u; // 补位：码 2 → digit 1（无信息，解码截断） + const std::uint32_t code = + (idx < k) ? ((words[idx / kPerWord] >> (kBits * (idx % kPerWord))) & kMask) + : 2u; // 补位：码 2 → digit 1（无信息，解码截断） byte += (code - 1u) * pow3; pow3 *= 3u; } @@ -185,15 +175,16 @@ inline void pack_row_base3(const std::uint32_t* words, std::size_t k, } // base-3 字节行 → 2-bit 容器行（部署读盘方向）。尾部字按既有约定填"码 2 = 三值 0"。 -inline void unpack_row_base3(const std::uint8_t* in, std::size_t k, - std::uint32_t* words) noexcept +inline void unpack_row_base3(const std::uint8_t* in, std::size_t k, std::uint32_t* words) noexcept { const std::size_t nwords = words_per_row(k); - for (std::size_t wi = 0; wi < nwords; ++wi) words[wi] = 0; + for (std::size_t wi = 0; wi < nwords; ++wi) + words[wi] = 0; for (std::size_t idx = 0; idx < k; ++idx) { std::uint32_t byte = in[idx / kBase3PerByte]; - for (std::uint32_t i = 0; i < idx % kBase3PerByte; ++i) byte /= 3u; + for (std::uint32_t i = 0; i < idx % kBase3PerByte; ++i) + byte /= 3u; const std::uint32_t code = (byte % 3u) + 1u; // digit → 码 words[idx / kPerWord] |= code << (kBits * (idx % kPerWord)); } @@ -206,7 +197,8 @@ inline void unpack_row_base3(const std::uint8_t* in, std::size_t k, [[nodiscard]] inline float at_base3(const std::uint8_t* in, std::size_t idx) noexcept { std::uint32_t byte = in[idx / kBase3PerByte]; - for (std::uint32_t i = 0; i < idx % kBase3PerByte; ++i) byte /= 3u; + for (std::uint32_t i = 0; i < idx % kBase3PerByte; ++i) + byte /= 3u; return value_of((byte % 3u) + 1u); } diff --git a/include/neuralnet.cpp/text_trainer.hpp b/include/neuralnet.cpp/text_trainer.hpp index 72502c99..9de038c2 100644 --- a/include/neuralnet.cpp/text_trainer.hpp +++ b/include/neuralnet.cpp/text_trainer.hpp @@ -170,8 +170,7 @@ make_text_train_state(const TextTrainConfig& cfg, const DatasetIdentity& identit s.recipe.warmup_epochs = static_cast(cfg.warmup_epochs); s.recipe.warmup_steps = static_cast(cfg.warmup_steps); s.recipe.batch_size = static_cast(cfg.batch_size); - s.recipe.micro_batch_size = - static_cast(cfg.batch_size / cfg.accum_steps); + s.recipe.micro_batch_size = static_cast(cfg.batch_size / cfg.accum_steps); s.recipe.batch_size_is_global = true; s.recipe.accum_steps = static_cast(cfg.accum_steps); s.stride = static_cast(cfg.stride); @@ -236,19 +235,17 @@ class TextTrainer if (test && test->subset() != DatasetSubset::Test) NN_FAIL("TextTrainer::create: 测试样本集必须是 DatasetSubset::Test"); if (cfg.seq_len != train.seq_len()) - NN_FAIL("TextTrainer::create: 配方 seq_len 与训练样本集不一致（" - + std::to_string(cfg.seq_len) + " vs " - + std::to_string(train.seq_len()) + "）"); + NN_FAIL("TextTrainer::create: 配方 seq_len 与训练样本集不一致（" + + std::to_string(cfg.seq_len) + " vs " + std::to_string(train.seq_len()) + "）"); if (cfg.stride != 0 && cfg.stride != train.stride()) - NN_FAIL("TextTrainer::create: 配方 stride 与训练样本集不一致（" - + std::to_string(cfg.stride) + " vs " - + std::to_string(train.stride()) + "）"); + NN_FAIL("TextTrainer::create: 配方 stride 与训练样本集不一致（" + + std::to_string(cfg.stride) + " vs " + std::to_string(train.stride()) + "）"); NN_TRY_MSG(tokenizer_r, "加载数据集内嵌词表失败", load_tokenizer_from_kvrec(train.vocab_kvrec())); if ((*tokenizer_r)->pad_id() != train.pad_id()) - NN_FAIL("TextTrainer::create: 样本集 pad_id 与词表不一致（" - + std::to_string(train.pad_id()) + " vs " - + std::to_string((*tokenizer_r)->pad_id()) + "）"); + NN_FAIL("TextTrainer::create: 样本集 pad_id 与词表不一致（" + + std::to_string(train.pad_id()) + " vs " + + std::to_string((*tokenizer_r)->pad_id()) + "）"); return create_impl(std::move(session), std::move(train), std::move(test), std::move(*tokenizer_r), std::move(cfg), std::move(hooks)); } @@ -294,18 +291,16 @@ class TextTrainer TextTrainer(TrainingSession session, TextSampleSet train, std::optional test, std::unique_ptr tokenizer, TextTrainConfig cfg, TextTrainerHooks hooks, - DatasetIdentity identity, std::size_t steps_per_epoch, - std::size_t micro_batch_size) + DatasetIdentity identity, std::size_t steps_per_epoch, std::size_t micro_batch_size) : session_(std::move(session)), train_(std::move(train)), test_(std::move(test)), tokenizer_(std::move(tokenizer)), cfg_(std::move(cfg)), hooks_(std::move(hooks)), identity_(std::move(identity)), steps_per_epoch_(steps_per_epoch), micro_batch_size_(micro_batch_size) { ce_loss_.set_precision_profile(cfg_.precision); - train_state_ = make_text_train_state( - cfg_, identity_, steps_per_epoch_, - static_cast(cfg_.start_epoch), cfg_.start_step, - cfg_.lr, session_.step_count()); + train_state_ = make_text_train_state(cfg_, identity_, steps_per_epoch_, + static_cast(cfg_.start_epoch), + cfg_.start_step, cfg_.lr, session_.step_count()); } [[nodiscard]] static Result @@ -313,20 +308,18 @@ class TextTrainer std::unique_ptr tokenizer, TextTrainConfig cfg, TextTrainerHooks hooks) { // 形状/网格前置条件（fail-fast：训练工厂不接受无法定义 step 的配置） - if (cfg.batch_size == 0) - NN_FAIL("TextTrainer::create: batch_size（global batch）必须 > 0"); + if (cfg.batch_size == 0) NN_FAIL("TextTrainer::create: batch_size（global batch）必须 > 0"); if (cfg.accum_steps == 0) NN_FAIL("TextTrainer::create: accum_steps 必须 > 0"); if (cfg.batch_size % cfg.accum_steps != 0) - NN_FAIL("TextTrainer::create: batch_size 必须能被 accum_steps 整除（global=" - + std::to_string(cfg.batch_size) + ", accum=" - + std::to_string(cfg.accum_steps) + "）"); + NN_FAIL("TextTrainer::create: batch_size 必须能被 accum_steps 整除（global=" + + std::to_string(cfg.batch_size) + ", accum=" + std::to_string(cfg.accum_steps) + + "）"); if (cfg.log_interval == 0) NN_FAIL("TextTrainer::create: log_interval 必须 > 0（0 会在进度打印时除零；" "关闭输出请把 hooks.out 置空）"); if (cfg.epochs < 0) NN_FAIL("TextTrainer::create: epochs 不能为负"); if (cfg.start_epoch < 0) NN_FAIL("TextTrainer::create: start_epoch 不能为负"); - if (cfg.start_epoch > cfg.epochs) - NN_FAIL("TextTrainer::create: start_epoch 超出 epochs"); + if (cfg.start_epoch > cfg.epochs) NN_FAIL("TextTrainer::create: start_epoch 超出 epochs"); if (train.sample_count() == 0) NN_FAIL("TextTrainer::create: 训练样本集为空"); // 模型序列长度校验（发布审查 P1-1）：窗口 seq_len 必须等于模型构建时的 // seq_len。GPTModel::forward 用 input.rows() 解释序列长度，位置编码/ @@ -335,16 +328,14 @@ class TextTrainer // 上存在；手工 build_* 调用没有 spec 时无法在此校验（该缺口记入 // v1.10.0 release note 已知限制）。 const auto& spec = session.model().spec(); - const std::size_t model_seq_len = - spec ? (spec->is_gpt() ? spec->as()->seq_len - : spec->is_rapt() ? spec->as()->seq_len - : 0) - : 0; + const std::size_t model_seq_len = spec ? (spec->is_gpt() ? spec->as()->seq_len + : spec->is_rapt() ? spec->as()->seq_len + : 0) + : 0; if (model_seq_len > 0 && cfg.seq_len != model_seq_len) - NN_FAIL("TextTrainer::create: 配方 seq_len 与模型 spec.seq_len 不一致（" - + std::to_string(cfg.seq_len) + " vs " - + std::to_string(model_seq_len) - + "）；请让 TextTrainConfig.seq_len 与模型构建时的 seq_len 对齐"); + NN_FAIL("TextTrainer::create: 配方 seq_len 与模型 spec.seq_len 不一致（" + + std::to_string(cfg.seq_len) + " vs " + std::to_string(model_seq_len) + + "）；请让 TextTrainConfig.seq_len 与模型构建时的 seq_len 对齐"); if (tokenizer->pad_id() != train.pad_id()) NN_FAIL("TextTrainer::create: 样本集 pad_id 与词表不一致"); // 空 test 子集 = 不评估（旧 CLI 语义：样本数 0 ⇒ 跳过测试评估，不终止训练） @@ -359,40 +350,38 @@ class TextTrainer // 一致，否则会出现"实际用 Session 的 lr、日志/断点写配方 lr"的静默错训。 const TrainingConfig& sc = session.config(); if (sc.optimizer != cfg.optimizer_name) - NN_FAIL("TextTrainer::create: 优化器不一致（Session=" + sc.optimizer - + ", 配方=" + cfg.optimizer_name + "）"); + NN_FAIL("TextTrainer::create: 优化器不一致（Session=" + sc.optimizer + + ", 配方=" + cfg.optimizer_name + "）"); if (sc.lr != cfg.lr) - NN_FAIL("TextTrainer::create: lr 不一致（Session=" + std::to_string(sc.lr) - + ", 配方=" + std::to_string(cfg.lr) + "）"); + NN_FAIL("TextTrainer::create: lr 不一致（Session=" + std::to_string(sc.lr) + + ", 配方=" + std::to_string(cfg.lr) + "）"); if (sc.weight_decay != cfg.weight_decay) - NN_FAIL("TextTrainer::create: weight_decay 不一致（Session=" - + std::to_string(sc.weight_decay) + ", 配方=" - + std::to_string(cfg.weight_decay) + "）"); + NN_FAIL("TextTrainer::create: weight_decay 不一致（Session=" + + std::to_string(sc.weight_decay) + ", 配方=" + std::to_string(cfg.weight_decay) + + "）"); if (sc.beta1 != cfg.beta1 || sc.beta2 != cfg.beta2) - NN_FAIL("TextTrainer::create: beta1/beta2 不一致（Session=" - + std::to_string(sc.beta1) + "/" + std::to_string(sc.beta2) - + ", 配方=" + std::to_string(cfg.beta1) + "/" - + std::to_string(cfg.beta2) + "）"); + NN_FAIL("TextTrainer::create: beta1/beta2 不一致（Session=" + std::to_string(sc.beta1) + + "/" + std::to_string(sc.beta2) + ", 配方=" + std::to_string(cfg.beta1) + "/" + + std::to_string(cfg.beta2) + "）"); if (!same_precision_profile(sc.precision, cfg.precision)) NN_FAIL("TextTrainer::create: precision profile 不一致" "（Session 与配方必须用同一份；见 make_text_train_config）"); if (sc.clip_grad_norm != cfg.max_norm) - NN_FAIL("TextTrainer::create: 梯度裁剪配置不一致（Session=" - + std::to_string(sc.clip_grad_norm) + ", 配方 max_norm=" - + std::to_string(cfg.max_norm) + "）"); + NN_FAIL("TextTrainer::create: 梯度裁剪配置不一致（Session=" + + std::to_string(sc.clip_grad_norm) + + ", 配方 max_norm=" + std::to_string(cfg.max_norm) + "）"); if (sc.grad_accum_steps != cfg.accum_steps) - NN_FAIL("TextTrainer::create: 梯度累积次数不一致（Session=" - + std::to_string(sc.grad_accum_steps) + ", 配方 accum_steps=" - + std::to_string(cfg.accum_steps) + "）"); + NN_FAIL("TextTrainer::create: 梯度累积次数不一致（Session=" + + std::to_string(sc.grad_accum_steps) + + ", 配方 accum_steps=" + std::to_string(cfg.accum_steps) + "）"); const std::size_t steps = train.steps_per_epoch(cfg.batch_size); const std::size_t micro_batch_size = cfg.batch_size / cfg.accum_steps; if (steps == 0) NN_FAIL("TextTrainer::create: 训练步数为 0"); - if (micro_batch_size == 0) - NN_FAIL("TextTrainer::create: 派生 micro_batch_size 为 0"); + if (micro_batch_size == 0) NN_FAIL("TextTrainer::create: 派生 micro_batch_size 为 0"); if (cfg.start_step >= steps) - NN_FAIL("TextTrainer::create: start_step 超出每 epoch 步数（" - + std::to_string(cfg.start_step) + " >= " + std::to_string(steps) + "）"); + NN_FAIL("TextTrainer::create: start_step 超出每 epoch 步数（" + + std::to_string(cfg.start_step) + " >= " + std::to_string(steps) + "）"); // 数据身份：调用方已算好的字段优先（CLI 已为断点门禁算过一次 sha256， // 这里不重复哈希大文件）；未给的由样本集合补齐。 @@ -416,7 +405,7 @@ class TextTrainer // 断点状态的权威装配（计数器 = "下一步的位置"，epoch 边界归一化）。 [[nodiscard]] TextTrainState make_train_state(std::size_t epoch, std::size_t step_done, - Scalar current_lr) const + Scalar current_lr) const { return make_text_train_state(cfg_, identity_, steps_per_epoch_, epoch, step_done, current_lr, session_.step_count()); @@ -480,15 +469,15 @@ inline Result TextTrainer::fit() // 装配回调只产出归一化单位（有效 token）；真正的 batch 装配在 do_update // 里按 micro 子跨度做（global batch 的 micro 切分是文本域语义）。 const std::uint64_t seed = static_cast(std::random_device{}()); - NN_TRY(source_r, IndexedBatchSource::create( - train_.sample_count(), cfg_.batch_size, seed, - [this](std::span ix) - -> Result> - { - NN_TRY_MSG(u, "统计 global batch 有效 token 数失败", - train_.valid_target_count(ix)); - return std::make_pair(std::any{}, static_cast(*u)); - })); + NN_TRY(source_r, + IndexedBatchSource::create( + train_.sample_count(), cfg_.batch_size, seed, + [this](std::span ix) -> Result> + { + NN_TRY_MSG(u, "统计 global batch 有效 token 数失败", + train_.valid_target_count(ix)); + return std::make_pair(std::any{}, static_cast(*u)); + })); IndexedBatchSource& source = *source_r; // ── 非阻塞 loss 回读状态（设备侧异步槽位；host 不等 GPU）───────────── @@ -556,15 +545,15 @@ inline Result TextTrainer::fit() TextSampleBatch batch; // 可复用的装配缓冲（训练 micro + 评估共用） - NN_TRY(driver_r, RunDriver::create( - session_, source, std::make_unique(train_state_), std::move(rcfg))); + NN_TRY(driver_r, + RunDriver::create(session_, source, std::make_unique(train_state_), + std::move(rcfg))); RunDriver& driver = *driver_r; RunHooks hooks; // ── 必选：一步 update = 取 global batch + micro 切分 + train_step ───── - hooks.do_update = [&](TrainingSession&, const UpdateRequest& req) - -> Result + hooks.do_update = [&](TrainingSession&, const UpdateRequest& req) -> Result { // epoch 切换 ⇒ 重置本 epoch 的回读/统计状态（旧 fit 的 epoch 局部量） if (req.epoch != cur_epoch) @@ -590,8 +579,7 @@ inline Result TextTrainer::fit() UpdateOutcome outcome; outcome.samples = global_bs; - const std::size_t micro_count = - (global_bs + micro_batch_size_ - 1) / micro_batch_size_; + const std::size_t micro_count = (global_bs + micro_batch_size_ - 1) / micro_batch_size_; NN_PROF_SCOPE("step"); nn::prof::Phase prof; @@ -601,12 +589,10 @@ inline Result TextTrainer::fit() for (std::size_t micro = 0; micro < micro_count; ++micro) { const std::size_t micro_offset = micro * micro_batch_size_; - const std::size_t this_micro_bs = - std::min(micro_batch_size_, global_bs - micro_offset); + const std::size_t this_micro_bs = std::min(micro_batch_size_, global_bs - micro_offset); const auto micro_indices = global_indices.subspan(micro_offset, this_micro_bs); - NN_TRY_MSG( - batch_r, "装配训练 micro-batch 失败", - train_.make_batch(micro_indices, batch)); + NN_TRY_MSG(batch_r, "装配训练 micro-batch 失败", + train_.make_batch(micro_indices, batch)); model.set_doc_ids(batch.doc_ids); if (hooks_.on_batch_ready) hooks_.on_batch_ready(session_, batch, req.step); @@ -622,52 +608,49 @@ inline Result TextTrainer::fit() options.force_update = (micro + 1 == micro_count) || req.is_final_step; if (hooks_.on_probe_check) { - options.before_update = [&]() -> Result { - return hooks_.on_probe_check(session_, req.step); - }; + options.before_update = [&]() -> Result + { return hooks_.on_probe_check(session_, req.step); }; } if (hooks_.on_update) { - options.after_update = [&]() -> Result { + options.after_update = [&]() -> Result + { hooks_.on_update(session_, req.epoch, req.step, req.lr); return {}; }; } prof.start(micro == 0 ? "forward" : "micro"); - NN_TRY_MSG(step_r, "TrainingSession::train_step 失败", - session_.train_step( - batch, std::move(options), - [&](Model& step_model, ComputeEngine& step_engine, - const TextSampleBatch& micro_batch) - -> Result { - auto logits_r = step_model.forward(*x_tensor_r); - NN_TRY_CHECK(logits_r); - Tensor logits = std::move(*logits_r); - if (logits.rows() != vocab_size) - NN_FAIL("TextTrainer::fit: 模型 logits 行数（" - + std::to_string(logits.rows()) + - "）≠ 数据集词表大小（" + - std::to_string(vocab_size) + - "）；模型与数据集词表不匹配"); - - auto mask_span = - std::span(micro_batch.flat_mask); - std::size_t loss_units = 0; - NN_TRY_MSG( - loss_sum_r, "损失前向（稀疏 CE）失败", - ce_loss_.forward_sparse_sum( - step_engine, logits, micro_batch.flat_targets, - mask_span, vocab_size, loss_units, - /*grad_reuse=*/&logits)); - NN_TRY(grad_r, ce_loss_.backward()); - TrainingStepOutput out; - out.grad = std::move(*grad_r); - out.gradient_units = loss_units; - out.loss_sum = std::move(*loss_sum_r); - out.loss_units = loss_units; - return out; - })); + NN_TRY_MSG( + step_r, "TrainingSession::train_step 失败", + session_.train_step( + batch, std::move(options), + [&](Model& step_model, ComputeEngine& step_engine, + const TextSampleBatch& micro_batch) -> Result + { + auto logits_r = step_model.forward(*x_tensor_r); + NN_TRY_CHECK(logits_r); + Tensor logits = std::move(*logits_r); + if (logits.rows() != vocab_size) + NN_FAIL("TextTrainer::fit: 模型 logits 行数（" + + std::to_string(logits.rows()) + "）≠ 数据集词表大小（" + + std::to_string(vocab_size) + "）；模型与数据集词表不匹配"); + + auto mask_span = std::span(micro_batch.flat_mask); + std::size_t loss_units = 0; + NN_TRY_MSG(loss_sum_r, "损失前向（稀疏 CE）失败", + ce_loss_.forward_sparse_sum(step_engine, logits, + micro_batch.flat_targets, mask_span, + vocab_size, loss_units, + /*grad_reuse=*/&logits)); + NN_TRY(grad_r, ce_loss_.backward()); + TrainingStepOutput out; + out.grad = std::move(*grad_r); + out.gradient_units = loss_units; + out.loss_sum = std::move(*loss_sum_r); + out.loss_units = loss_units; + return out; + })); mem_mark("micro/backward"); NN_TRY_MSG(rel_r, "\n显存回收失败", engine.release_idle_pool_blocks()); @@ -694,8 +677,7 @@ inline Result TextTrainer::fit() pl.slot = loss_slot_next; pl.step = req.step; pl.valid = step_r->loss_units; - pl.inv_num_valid = - Scalar{1} / static_cast(step_r->loss_units); + pl.inv_num_valid = Scalar{1} / static_cast(step_r->loss_units); pl.keepalive = std::move(*loss_sum_f32_r); NN_TRY_MSG(sr, "submit_scalar_readback failed", engine.submit_scalar_readback(pl.slot, pl.keepalive)); @@ -704,8 +686,7 @@ inline Result TextTrainer::fit() updated_this_step = true; break; } - if (!updated_this_step) - NN_FAIL("TextTrainer::fit: global batch 结束时未完成参数更新"); + if (!updated_this_step) NN_FAIL("TextTrainer::fit: global batch 结束时未完成参数更新"); NN_TRY_MSG(hv, "harvest_loss failed", harvest_loss()); if (req.is_final_step) saw_final_step = true; @@ -747,8 +728,8 @@ inline Result TextTrainer::fit() std::vector test_indices(this_bs); for (std::size_t i = 0; i < this_bs; ++i) test_indices[i] = toffset + i; - if (const auto br = test_->make_batch( - std::span(test_indices), batch); + if (const auto br = + test_->make_batch(std::span(test_indices), batch); !br) { err << " 测试装配 batch 失败: " << br.error().message << '\n'; @@ -777,8 +758,8 @@ inline Result TextTrainer::fit() } auto tlogits = std::move(*tf_r); auto tmask = std::span(batch.flat_mask); - auto tl_r = ce_loss_.forward_sparse(engine, tlogits, batch.flat_targets, tmask, - vocab_size); + auto tl_r = + ce_loss_.forward_sparse(engine, tlogits, batch.flat_targets, tmask, vocab_size); if (!tl_r) { err << " 测试评估出错: " << tl_r.error().message << '\n'; @@ -812,8 +793,7 @@ inline Result TextTrainer::fit() { const TrainCounters& c = st.counters; // 权威装配走 make_train_state（配方 + 数据集身份 + 优化器 t 全量重建） - train_state_ = make_train_state(c.epoch, c.step_in_epoch, - static_cast(c.lr)); + train_state_ = make_train_state(c.epoch, c.step_in_epoch, static_cast(c.lr)); if (hooks_.on_checkpoint) { const TextCheckpointRequest req{session_, train_state_}; @@ -832,8 +812,8 @@ inline Result TextTrainer::fit() { NN_TRY_MSG(ep_drain, "loss drain failed", drain_loss()); if (saw_final_step) - out << "\n已达 --max-steps=" << cfg_.max_steps << "（本进程执行 " - << rep.steps_executed << " 步），提前结束训练\n"; + out << "\n已达 --max-steps=" << cfg_.max_steps << "（本进程执行 " << rep.steps_executed + << " 步），提前结束训练\n"; } // Run state = "下一步位置"：正常跑完所有 epoch ⇒ (epochs, 0)； @@ -857,8 +837,7 @@ inline Result TextTrainer::fit() if (rep.abort_code == 0) { report.final_epoch = static_cast(train_state_.counters.epoch); - report.final_step_in_epoch = - static_cast(train_state_.counters.step_in_epoch); + report.final_step_in_epoch = static_cast(train_state_.counters.step_in_epoch); } else { diff --git a/include/neuralnet.cpp/train_state.hpp b/include/neuralnet.cpp/train_state.hpp index aac91527..6490494d 100644 --- a/include/neuralnet.cpp/train_state.hpp +++ b/include/neuralnet.cpp/train_state.hpp @@ -52,14 +52,10 @@ namespace nn { // ── 浮点字段的落盘编码（double 位模式装进 u64）────────────────────────── [[nodiscard]] inline std::uint64_t frac_to_bits(double v) noexcept -{ - return std::bit_cast(v); -} +{ return std::bit_cast(v); } [[nodiscard]] inline double bits_to_frac(std::uint64_t bits) noexcept -{ - return std::bit_cast(bits); -} +{ return std::bit_cast(bits); } // ── 计数器（续训起点）────────────────────────────────────────────────── struct TrainCounters @@ -115,16 +111,15 @@ struct DatasetIdentity // steps_per_epoch 变 ⇒ step 级调度的 total_steps、global_step 的语义、 // epoch 边界**同时**改变 ⇒ "同一断点 + 不同数据集"不存在正确的续训语义。 // 返回空串 = 一致；否则返回人类可读的差异说明（诊断面，逐项列出变了什么）。 -[[nodiscard]] inline std::string dataset_mismatch_reason( - const DatasetIdentity &saved, const DatasetIdentity ¤t) +[[nodiscard]] inline std::string dataset_mismatch_reason(const DatasetIdentity& saved, + const DatasetIdentity& current) { - auto short_sha = [](const std::string &s) { - return s.size() > 12 ? s.substr(0, 12) + "…" : s; - }; + auto short_sha = [](const std::string& s) { return s.size() > 12 ? s.substr(0, 12) + "…" : s; }; auto u64 = [](std::uint64_t v) { return std::to_string(v); }; std::string why; - auto add = [&why](const std::string &line) { + auto add = [&why](const std::string& line) + { if (!why.empty()) why += "\n"; why += " · " + line; }; @@ -137,8 +132,7 @@ struct DatasetIdentity if (saved.num_samples != current.num_samples) add("样本数 : " + u64(saved.num_samples) + " → " + u64(current.num_samples)); if (saved.steps_per_epoch != current.steps_per_epoch) - add("每 epoch 步数: " + u64(saved.steps_per_epoch) + " → " + - u64(current.steps_per_epoch)); + add("每 epoch 步数: " + u64(saved.steps_per_epoch) + " → " + u64(current.steps_per_epoch)); if (saved.loss_scope != current.loss_scope) add("loss 范围 : " + saved.loss_scope + " → " + current.loss_scope); return why; @@ -155,22 +149,17 @@ class TrainState virtual ~TrainState() = default; // 域状态种类（落盘键 `state_kind`；旧文件缺键 ⇒ "text"） - [[nodiscard]] virtual const char *state_kind() const noexcept = 0; + [[nodiscard]] virtual const char* state_kind() const noexcept = 0; // 深拷贝（状态段写盘需要快照；多态值语义的支撑） [[nodiscard]] virtual std::unique_ptr clone() const = 0; // 域字段落盘（基类通用字段之后追加） - virtual void write_extra_kv(KeyValueRecord & /*kv*/) const {} + virtual void write_extra_kv(KeyValueRecord& /*kv*/) const {} // 域字段读回（缺失键 → 保持默认） - [[nodiscard]] virtual Result read_extra_kv(const KeyValueRecord & /*kv*/) - { - return {}; - } + [[nodiscard]] virtual Result read_extra_kv(const KeyValueRecord& /*kv*/) { return {}; } // 域级续训交叉校验（`verify_resume` 在通用校验后调用；返回错误 = 不可续训） - [[nodiscard]] virtual Result verify_extra(const TrainState & /*current*/) const - { - return {}; - } + [[nodiscard]] virtual Result verify_extra(const TrainState& /*current*/) const + { return {}; } TrainCounters counters; TrainRecipe recipe; @@ -182,16 +171,14 @@ class TrainState class TextTrainState : public TrainState { public: - [[nodiscard]] const char *state_kind() const noexcept override { return "text"; } + [[nodiscard]] const char* state_kind() const noexcept override { return "text"; } [[nodiscard]] std::unique_ptr clone() const override - { - return std::make_unique(*this); - } + { return std::make_unique(*this); } std::uint64_t stride = 0; // 滑动窗口步长（0 = 等于 seq_len） DatasetIdentity dataset; // 数据集身份（续训门禁的数据源） - void write_extra_kv(KeyValueRecord &kv) const override + void write_extra_kv(KeyValueRecord& kv) const override { kv.set("stride", stride); kv.set("dataset_path", dataset.path); @@ -203,7 +190,7 @@ class TextTrainState : public TrainState kv.set("dataset_loss_scope", dataset.loss_scope); } - [[nodiscard]] Result read_extra_kv(const KeyValueRecord &kv) override + [[nodiscard]] Result read_extra_kv(const KeyValueRecord& kv) override { std::uint64_t u = 0; if (kv.get("stride", u)) stride = u; @@ -222,19 +209,18 @@ class TextTrainState : public TrainState return {}; } - [[nodiscard]] Result verify_extra(const TrainState ¤t) const override + [[nodiscard]] Result verify_extra(const TrainState& current) const override { - const auto *ts = dynamic_cast(¤t); + const auto* ts = dynamic_cast(¤t); if (ts == nullptr) NN_FAIL("续训交叉校验: 断点是文本状态（text）而本次不是文本状态（kind=" + std::string(current.state_kind()) + "）"); const std::string why = dataset_mismatch_reason(dataset, ts->dataset); if (!why.empty()) - NN_FAIL("续训交叉校验: 数据集身份不一致（内容身份变了 ⇒ 续训不存在正确语义）:\n" + - why); + NN_FAIL("续训交叉校验: 数据集身份不一致（内容身份变了 ⇒ 续训不存在正确语义）:\n" + why); if (stride != ts->stride) - NN_FAIL("续训交叉校验: 滑窗 stride 不一致（" + std::to_string(stride) + - " vs " + std::to_string(ts->stride) + "）"); + NN_FAIL("续训交叉校验: 滑窗 stride 不一致（" + std::to_string(stride) + " vs " + + std::to_string(ts->stride) + "）"); return {}; } }; @@ -244,7 +230,7 @@ class TextTrainState : public TrainState // ═══════════════════════════════════════════════════════════════════════ inline constexpr std::uint64_t TRAIN_STATE_VERSION = 1; -[[nodiscard]] inline KeyValueRecord train_state_to_kv(const TrainState &s) +[[nodiscard]] inline KeyValueRecord train_state_to_kv(const TrainState& s) { KeyValueRecord kv; kv.set("state_version", TRAIN_STATE_VERSION); @@ -280,13 +266,13 @@ inline constexpr std::uint64_t TRAIN_STATE_VERSION = 1; kv.set("warmup_steps", s.recipe.warmup_steps); kv.set("batch_size", s.recipe.batch_size); kv.set("micro_batch_size", s.recipe.micro_batch_size); - kv.set("batch_size_semantics", - std::string(s.recipe.batch_size_is_global ? "global" : "micro")); + kv.set("batch_size_semantics", std::string(s.recipe.batch_size_is_global ? "global" : "micro")); kv.set("accum_steps", s.recipe.accum_steps); { std::vector lrs; lrs.reserve(s.recipe.lr_per_epoch.size()); - for (double v : s.recipe.lr_per_epoch) lrs.push_back(frac_to_bits(v)); + for (double v : s.recipe.lr_per_epoch) + lrs.push_back(frac_to_bits(v)); kv.set("lr_per_epoch_bits", lrs); } @@ -296,8 +282,8 @@ inline constexpr std::uint64_t TRAIN_STATE_VERSION = 1; } // 精度槽读回：只接受**存储精度**与三值（BF16/F64 是保留值，文件里出现即报错） -[[nodiscard]] inline Result precision_from_state_kv( - const KeyValueRecord &kv, const std::string &key, Precision fallback) +[[nodiscard]] inline Result +precision_from_state_kv(const KeyValueRecord& kv, const std::string& key, Precision fallback) { std::uint64_t v = 0; if (!kv.get(key, v)) return fallback; @@ -311,16 +297,16 @@ inline constexpr std::uint64_t TRAIN_STATE_VERSION = 1; // 工厂：kind → 子类实例。旧文件缺 `state_kind` 键 ⇒ "text"（历史断点全部是 // 文本训练；零破坏）。未知 kind ⇒ 硬报错（不用错误类型的字段布局读文件）。 [[nodiscard]] inline Result> -train_state_from_kv(const KeyValueRecord &kv) +train_state_from_kv(const KeyValueRecord& kv) { std::string kind = "text"; if (!kv.get("state_kind", kind)) kind = "text"; std::unique_ptr s; - if (kind == "text") - s = std::make_unique(); + if (kind == "text") s = std::make_unique(); else - NN_FAIL("未知的训练状态 kind=\"" + kind + "\"（本构建只认识 text；" + NN_FAIL("未知的训练状态 kind=\"" + kind + + "\"（本构建只认识 text；" "自定义域请注册对应 TrainState 子类的读回路径）"); std::uint64_t u = 0; @@ -333,17 +319,15 @@ train_state_from_kv(const KeyValueRecord &kv) if (kv.get("optimizer_step", u)) s->optimizer_step = u; if (kv.get("optimizer_blob_count", u)) s->optimizer_blob_count = u; - NN_TRY(p_param, precision_from_state_kv( - kv, "precision_param", s->recipe.precision.param)); + NN_TRY(p_param, precision_from_state_kv(kv, "precision_param", s->recipe.precision.param)); s->recipe.precision.param = *p_param; - NN_TRY(p_compute, precision_from_state_kv( - kv, "precision_compute", s->recipe.precision.compute)); + NN_TRY(p_compute, + precision_from_state_kv(kv, "precision_compute", s->recipe.precision.compute)); s->recipe.precision.compute = *p_compute; - NN_TRY(p_stable, precision_from_state_kv( - kv, "precision_stable", s->recipe.precision.stable)); + NN_TRY(p_stable, precision_from_state_kv(kv, "precision_stable", s->recipe.precision.stable)); s->recipe.precision.stable = *p_stable; - NN_TRY(p_optimizer, precision_from_state_kv( - kv, "precision_optimizer", s->recipe.precision.optimizer)); + NN_TRY(p_optimizer, + precision_from_state_kv(kv, "precision_optimizer", s->recipe.precision.optimizer)); s->recipe.precision.optimizer = *p_optimizer; if (!kv.get("lr_schedule", s->recipe.lr_schedule)) s->recipe.lr_schedule = "fixed"; @@ -375,7 +359,8 @@ train_state_from_kv(const KeyValueRecord &kv) { s->recipe.lr_per_epoch.clear(); s->recipe.lr_per_epoch.reserve(lrs.size()); - for (std::uint64_t b : lrs) s->recipe.lr_per_epoch.push_back(bits_to_frac(b)); + for (std::uint64_t b : lrs) + s->recipe.lr_per_epoch.push_back(bits_to_frac(b)); } } @@ -387,33 +372,35 @@ train_state_from_kv(const KeyValueRecord &kv) // Strict（默认）：kind 相等 + 配方关键维度相等 + 计数器不越界 + 域级 // `verify_extra()`（文本 = 数据集身份 + stride）。Relaxed：只做 kind 检查， // 允许研究期换数据集 / 改配方（代价 = lr 曲线形状可能跨续训改变）。 -enum class ResumePolicy { Strict, Relaxed }; +enum class ResumePolicy +{ + Strict, + Relaxed +}; -[[nodiscard]] inline Result verify_resume(const TrainState &saved, - const TrainState ¤t, +[[nodiscard]] inline Result verify_resume(const TrainState& saved, const TrainState& current, ResumePolicy policy = ResumePolicy::Strict) { if (std::string(saved.state_kind()) != std::string(current.state_kind())) - NN_FAIL("续训交叉校验: 状态 kind 不一致（" + std::string(saved.state_kind()) + - " vs " + std::string(current.state_kind()) + "）"); - if (policy == ResumePolicy::Relaxed) - return {}; + NN_FAIL("续训交叉校验: 状态 kind 不一致（" + std::string(saved.state_kind()) + " vs " + + std::string(current.state_kind()) + "）"); + if (policy == ResumePolicy::Relaxed) return {}; - const TrainRecipe &a = saved.recipe; - const TrainRecipe &b = current.recipe; - auto mismatch = [](const char *name, const std::string &x, const std::string &y) { + const TrainRecipe& a = saved.recipe; + const TrainRecipe& b = current.recipe; + auto mismatch = [](const char* name, const std::string& x, const std::string& y) + { return std::string("续训交叉校验: 配方 ") + name + " 不一致（断点 \"" + x + "\" vs 本次 \"" + y + "\"）；lr 曲线形状必须跨续训不变"; }; - if (a.optimizer != b.optimizer) - NN_FAIL(mismatch("optimizer", a.optimizer, b.optimizer)); + if (a.optimizer != b.optimizer) NN_FAIL(mismatch("optimizer", a.optimizer, b.optimizer)); if (a.lr_schedule != b.lr_schedule) NN_FAIL(mismatch("lr_schedule", a.lr_schedule, b.lr_schedule)); if (a.epochs != b.epochs) NN_FAIL(mismatch("epochs", std::to_string(a.epochs), std::to_string(b.epochs))); if (a.total_steps != b.total_steps) - NN_FAIL(mismatch("total_steps", std::to_string(a.total_steps), - std::to_string(b.total_steps))); + NN_FAIL( + mismatch("total_steps", std::to_string(a.total_steps), std::to_string(b.total_steps))); if (a.warmup_epochs != b.warmup_epochs) NN_FAIL(mismatch("warmup_epochs", std::to_string(a.warmup_epochs), std::to_string(b.warmup_epochs))); @@ -421,15 +408,12 @@ enum class ResumePolicy { Strict, Relaxed }; NN_FAIL(mismatch("warmup_steps", std::to_string(a.warmup_steps), std::to_string(b.warmup_steps))); if (a.batch_size != b.batch_size) - NN_FAIL(mismatch("batch_size", std::to_string(a.batch_size), - std::to_string(b.batch_size))); + NN_FAIL(mismatch("batch_size", std::to_string(a.batch_size), std::to_string(b.batch_size))); if (a.accum_steps != b.accum_steps) - NN_FAIL(mismatch("accum_steps", std::to_string(a.accum_steps), - std::to_string(b.accum_steps))); - if (a.precision.param != b.precision.param || - a.precision.compute != b.precision.compute || - a.precision.stable != b.precision.stable || - a.precision.optimizer != b.precision.optimizer) + NN_FAIL( + mismatch("accum_steps", std::to_string(a.accum_steps), std::to_string(b.accum_steps))); + if (a.precision.param != b.precision.param || a.precision.compute != b.precision.compute || + a.precision.stable != b.precision.stable || a.precision.optimizer != b.precision.optimizer) NN_FAIL(mismatch("precision", "(见断点)", "(见本次)")); // 计数器不越界：断点已训完的 epoch 超过本次配方的 epochs ⇒ 配方缩水 diff --git a/include/neuralnet.cpp/training_group_session.hpp b/include/neuralnet.cpp/training_group_session.hpp index d4f9b5e7..0d7ce7f9 100644 --- a/include/neuralnet.cpp/training_group_session.hpp +++ b/include/neuralnet.cpp/training_group_session.hpp @@ -35,20 +35,12 @@ class TrainingGroupSession std::unique_ptr training; std::unique_ptr readonly; - [[nodiscard]] bool trainable() const noexcept - { - return training != nullptr; - } + [[nodiscard]] bool trainable() const noexcept { return training != nullptr; } - [[nodiscard]] Model& model() - { - return training ? training->model() : *readonly; - } + [[nodiscard]] Model& model() { return training ? training->model() : *readonly; } [[nodiscard]] const Model& model() const - { - return training ? training->model() : *readonly; - } + { return training ? training->model() : *readonly; } }; ComputeEngine* engine_ = nullptr; @@ -68,29 +60,24 @@ class TrainingGroupSession return nullptr; } - [[nodiscard]] Result validate_new_role( - std::string_view name, Model& model) + [[nodiscard]] Result validate_new_role(std::string_view name, Model& model) { - if (name.empty()) - NN_FAIL("TrainingGroupSession: role 名不能为空"); + if (name.empty()) NN_FAIL("TrainingGroupSession: role 名不能为空"); if (find_role(name) != nullptr) - NN_FAIL("TrainingGroupSession: role '" + std::string(name) + - "' 重复"); + NN_FAIL("TrainingGroupSession: role '" + std::string(name) + "' 重复"); if (!model.engine_bound()) - NN_FAIL("TrainingGroupSession: role '" + std::string(name) + - "' 的 Model 未绑定引擎"); + NN_FAIL("TrainingGroupSession: role '" + std::string(name) + "' 的 Model 未绑定引擎"); ComputeEngine& candidate = model.engine(); - if (engine_ == nullptr) - engine_ = &candidate; + if (engine_ == nullptr) engine_ = &candidate; else if (engine_ != &candidate) NN_FAIL("TrainingGroupSession: 所有 role 必须共享同一 ComputeEngine" - "（role='" + std::string(name) + "'）"); + "（role='" + + std::string(name) + "'）"); // 有效性校验；不同算法角色可拥有不同 schema，不强制 fingerprint 相同。 NN_TRY(schema, model.output_schema()); if (!schema->valid()) - NN_FAIL("TrainingGroupSession: role '" + std::string(name) + - "' 的 output schema 无效"); + NN_FAIL("TrainingGroupSession: role '" + std::string(name) + "' 的 output schema 无效"); return {}; } @@ -102,20 +89,17 @@ class TrainingGroupSession TrainingGroupSession& operator=(TrainingGroupSession&&) noexcept = default; ~TrainingGroupSession() = default; - [[nodiscard]] Result add_training_role( - std::string name, TrainingSession session) + [[nodiscard]] Result add_training_role(std::string name, TrainingSession session) { NN_TRY_CHECK(validate_new_role(name, session.model())); auto role = std::make_unique(); role->name = std::move(name); - role->training = - std::make_unique(std::move(session)); + role->training = std::make_unique(std::move(session)); roles_.push_back(std::move(role)); return {}; } - [[nodiscard]] Result add_readonly_role( - std::string name, Model model) + [[nodiscard]] Result add_readonly_role(std::string name, Model model) { NN_TRY_CHECK(validate_new_role(name, model)); auto role = std::make_unique(); @@ -127,14 +111,10 @@ class TrainingGroupSession [[nodiscard]] bool empty() const noexcept { return roles_.empty(); } [[nodiscard]] std::size_t size() const noexcept { return roles_.size(); } - [[nodiscard]] bool engine_bound() const noexcept - { - return engine_ != nullptr; - } + [[nodiscard]] bool engine_bound() const noexcept { return engine_ != nullptr; } [[nodiscard]] ComputeEngine& engine() const { - NN_ASSERT_ALWAYS(engine_ != nullptr, - "TrainingGroupSession::engine: group 为空"); + NN_ASSERT_ALWAYS(engine_ != nullptr, "TrainingGroupSession::engine: group 为空"); return *engine_; } @@ -150,50 +130,38 @@ class TrainingGroupSession [[nodiscard]] Result is_trainable(std::string_view name) const { const Role* role = find_role(name); - if (role == nullptr) - NN_FAIL("TrainingGroupSession: 未知 role '" + std::string(name) + - "'"); + if (role == nullptr) NN_FAIL("TrainingGroupSession: 未知 role '" + std::string(name) + "'"); return role->trainable(); } [[nodiscard]] Result model(std::string_view name) { Role* role = find_role(name); - if (role == nullptr) - NN_FAIL("TrainingGroupSession: 未知 role '" + std::string(name) + - "'"); + if (role == nullptr) NN_FAIL("TrainingGroupSession: 未知 role '" + std::string(name) + "'"); return &role->model(); } [[nodiscard]] Result model(std::string_view name) const { const Role* role = find_role(name); - if (role == nullptr) - NN_FAIL("TrainingGroupSession: 未知 role '" + std::string(name) + - "'"); + if (role == nullptr) NN_FAIL("TrainingGroupSession: 未知 role '" + std::string(name) + "'"); return &role->model(); } - [[nodiscard]] Result training_session( - std::string_view name) + [[nodiscard]] Result training_session(std::string_view name) { Role* role = find_role(name); - if (role == nullptr) - NN_FAIL("TrainingGroupSession: 未知 role '" + std::string(name) + - "'"); + if (role == nullptr) NN_FAIL("TrainingGroupSession: 未知 role '" + std::string(name) + "'"); if (role->training == nullptr) NN_FAIL("TrainingGroupSession: role '" + std::string(name) + "' 是只读 Model，不是 TrainingSession"); return role->training.get(); } - [[nodiscard]] Result training_session( - std::string_view name) const + [[nodiscard]] Result training_session(std::string_view name) const { const Role* role = find_role(name); - if (role == nullptr) - NN_FAIL("TrainingGroupSession: 未知 role '" + std::string(name) + - "'"); + if (role == nullptr) NN_FAIL("TrainingGroupSession: 未知 role '" + std::string(name) + "'"); if (role->training == nullptr) NN_FAIL("TrainingGroupSession: role '" + std::string(name) + "' 是只读 Model，不是 TrainingSession"); diff --git a/include/neuralnet.cpp/training_session.hpp b/include/neuralnet.cpp/training_session.hpp index dd9f9168..4eaa4522 100644 --- a/include/neuralnet.cpp/training_session.hpp +++ b/include/neuralnet.cpp/training_session.hpp @@ -128,8 +128,7 @@ struct TrainingStepReport class TrainingSession { public: - [[nodiscard]] static Result create(Model model, - TrainingConfig cfg = {}) + [[nodiscard]] static Result create(Model model, TrainingConfig cfg = {}) { if (!model.engine_bound()) NN_FAIL("TrainingSession::create: Model 未绑定引擎" @@ -137,8 +136,7 @@ class TrainingSession if (cfg.grad_accum_steps == 0) NN_FAIL("TrainingSession::create: grad_accum_steps 必须 > 0"); - if (cfg.checkpoint_every > 0) - model.set_checkpoint_every(cfg.checkpoint_every); + if (cfg.checkpoint_every > 0) model.set_checkpoint_every(cfg.checkpoint_every); if (cfg.activation_offload) { if (model.engine().device() != Device::GPU) @@ -148,18 +146,17 @@ class TrainingSession } ComputeEngine& engine = model.engine(); - auto opt = create_optimizer(cfg.optimizer, engine, model.parameters(), - model.param_gradients(), cfg.lr, - cfg.weight_decay, cfg.precision, cfg.beta1, - cfg.beta2); + auto opt = + create_optimizer(cfg.optimizer, engine, model.parameters(), model.param_gradients(), + cfg.lr, cfg.weight_decay, cfg.precision, cfg.beta1, cfg.beta2); if (!opt) - NN_FAIL("TrainingSession::create: 未知优化器 '" + cfg.optimizer + "'" + NN_FAIL("TrainingSession::create: 未知优化器 '" + cfg.optimizer + + "'" "（可用: sgd / sgd_momentum / adam / adamw / adamwquant / muon）"); TrainingSessionState state; state.current_lr = cfg.lr; - return TrainingSession(std::move(model), std::move(opt), std::move(cfg), - std::move(state)); + return TrainingSession(std::move(model), std::move(opt), std::move(cfg), std::move(state)); } TrainingSession(TrainingSession&&) noexcept = default; @@ -176,8 +173,7 @@ class TrainingSession [[nodiscard]] const TrainingConfig& config() const noexcept { return cfg_; } [[nodiscard]] const TrainingSessionState& state() const noexcept { return state_; } [[nodiscard]] Scalar current_lr() const noexcept { return state_.current_lr; } - [[nodiscard]] std::uint64_t step_count() const noexcept - { return state_.optimizer_step; } + [[nodiscard]] std::uint64_t step_count() const noexcept { return state_.optimizer_step; } // 核心单步入口。`step_fn` 只在本次调用期间使用，不进入 Session 状态。 template @@ -225,8 +221,7 @@ class TrainingSession NN_FAIL("TrainingSession::train_step: Step 回调返回了空梯度"); } - if (out.gradient_units > 0 && - options.normalization_units != out.gradient_units) + if (out.gradient_units > 0 && options.normalization_units != out.gradient_units) { if (out.gradient_units > options.normalization_units) { @@ -236,9 +231,8 @@ class TrainingSession std::to_string(out.gradient_units) + " > " + std::to_string(options.normalization_units) + "）"); } - const Scalar scale = - static_cast(out.gradient_units) / - static_cast(options.normalization_units); + const Scalar scale = static_cast(out.gradient_units) / + static_cast(options.normalization_units); auto sr = engine().scale_inplace(out.grad, scale); if (!sr) { @@ -280,8 +274,7 @@ class TrainingSession } ++state_.microsteps_in_update; - const bool update_ready = - state_.microsteps_in_update >= cfg_.grad_accum_steps; + const bool update_ready = state_.microsteps_in_update >= cfg_.grad_accum_steps; if (!update_ready && !options.force_update) { TrainingStepReport report; @@ -326,8 +319,7 @@ class TrainingSession // 是 step_fn 返回 GradientBundle 并走 Model::backward_named。 template [[nodiscard]] Result - train_named_step(const Batch& batch, TrainingStepOptions options, - StepFn&& step_fn) + train_named_step(const Batch& batch, TrainingStepOptions options, StepFn&& step_fn) { if (state_.accumulation_dirty) NN_FAIL("TrainingSession::train_named_step: 上一次失败留下了不完整" @@ -351,8 +343,7 @@ class TrainingSession if (options.normalization_units != state_.update_normalization_units) NN_FAIL("TrainingSession::train_named_step: 同一 update 的 " "normalization_units 发生变化（" + - std::to_string(state_.update_normalization_units) + - " vs " + + std::to_string(state_.update_normalization_units) + " vs " + std::to_string(options.normalization_units) + "）"); } @@ -373,8 +364,7 @@ class TrainingSession NN_FAIL("TrainingSession::train_named_step: Step 回调返回了空梯度包"); } - if (out.gradient_units > 0 && - options.normalization_units != out.gradient_units) + if (out.gradient_units > 0 && options.normalization_units != out.gradient_units) { if (out.gradient_units > options.normalization_units) { @@ -384,13 +374,11 @@ class TrainingSession std::to_string(out.gradient_units) + " > " + std::to_string(options.normalization_units) + "）"); } - const Scalar scale = - static_cast(out.gradient_units) / - static_cast(options.normalization_units); + const Scalar scale = static_cast(out.gradient_units) / + static_cast(options.normalization_units); for (std::size_t i = 0; i < out.gradients.schema().size(); ++i) { - const FieldHandle handle{ - out.gradients.schema().fingerprint(), i}; + const FieldHandle handle{out.gradients.schema().fingerprint(), i}; auto grad_r = out.gradients.require(handle); NN_TRY_CHECK(grad_r); if (*grad_r == nullptr) continue; @@ -418,8 +406,7 @@ class TrainingSession } else { - auto ar = engine().accumulate(state_.pending_loss_sum, - out.loss_sum); + auto ar = engine().accumulate(state_.pending_loss_sum, out.loss_sum); if (!ar) { state_.accumulation_dirty = true; @@ -437,8 +424,7 @@ class TrainingSession } ++state_.microsteps_in_update; - const bool update_ready = - state_.microsteps_in_update >= cfg_.grad_accum_steps; + const bool update_ready = state_.microsteps_in_update >= cfg_.grad_accum_steps; if (!update_ready && !options.force_update) { TrainingStepReport report; @@ -543,8 +529,7 @@ class TrainingSession // 状态 → 以 Checkpoint 角色落盘（blob 计数由 save_model 按实参回填， // 调用方不再需要手动同步）。`state` 的域字段（数据集身份等）由调用方 // 预填；本入口不覆盖它们。 - [[nodiscard]] Result save_checkpoint(const std::string& path, - const ModelSpec& spec, + [[nodiscard]] Result save_checkpoint(const std::string& path, const ModelSpec& spec, TrainState& state, const std::string& tokenizer_data = {}) { @@ -555,8 +540,7 @@ class TrainingSession std::vector blobs; NN_TRY(er, opt_->save_state(blobs)); const CheckpointState ckpt{&state, blobs}; - return save_model(path, model_, spec, tokenizer_data, - ModelRole::Checkpoint, ckpt); + return save_model(path, model_, spec, tokenizer_data, ModelRole::Checkpoint, ckpt); } // 一站式续训：载入权重（同一 .nnckpt 文件）→ 读训练状态段 → 可选交叉 @@ -564,22 +548,17 @@ class TrainingSession // Relaxed 只查 kind）→ 恢复优化器状态（t 先于/后于张量无序依赖，本入口 // 固定先 t 后张量）→ 恢复会话计数器。返回加载出的状态（域读取数据集身份 // 等字段用）。 - [[nodiscard]] Result> resume( - const std::string& path, const TrainState* current = nullptr, - ResumePolicy policy = ResumePolicy::Strict) + [[nodiscard]] Result> + resume(const std::string& path, const TrainState* current = nullptr, + ResumePolicy policy = ResumePolicy::Strict) { - NN_TRY(lm, load_model(path, model_, nullptr, nullptr, - WeightLoadIntent::Trainable)); + NN_TRY(lm, load_model(path, model_, nullptr, nullptr, WeightLoadIntent::Trainable)); NN_TRY(loaded, load_train_state(engine(), path)); - if (current != nullptr) - { - NN_TRY(vr, verify_resume(*loaded->state, *current, policy)); - } + if (current != nullptr) { NN_TRY(vr, verify_resume(*loaded->state, *current, policy)); } opt_->set_state_step(loaded->state->optimizer_step); NN_TRY(ls, opt_->load_state(loaded->optimizer_blobs)); state_.optimizer_step = loaded->state->counters.global_step; - if (loaded->state->counters.lr > 0) - set_lr(static_cast(loaded->state->counters.lr)); + if (loaded->state->counters.lr > 0) set_lr(static_cast(loaded->state->counters.lr)); clear_accumulation_state_(); return std::move(loaded->state); } @@ -590,8 +569,8 @@ class TrainingSession { public: explicit RecordingGuard(ComputeEngine& engine) noexcept : engine_(&engine) {} - RecordingGuard(RecordingGuard&& other) noexcept - : engine_(other.engine_) { other.engine_ = nullptr; } + RecordingGuard(RecordingGuard&& other) noexcept : engine_(other.engine_) + { other.engine_ = nullptr; } RecordingGuard& operator=(RecordingGuard&&) = delete; RecordingGuard(const RecordingGuard&) = delete; ~RecordingGuard() diff --git a/research/f16_weight_decay/decay_probe.cpp b/research/f16_weight_decay/decay_probe.cpp index 4e1c5aa6..ef460840 100644 --- a/research/f16_weight_decay/decay_probe.cpp +++ b/research/f16_weight_decay/decay_probe.cpp @@ -23,11 +23,12 @@ using nn::Matrix; using nn::Precision; using nn::Scalar; -namespace { +namespace +{ // f16 网格：给定 f32 值与相邻 f16 网格点，打印“掉了几个 ULP” -void report(const char* label, const std::vector& before, - const std::vector& after, int steps) +void report(const char* label, const std::vector& before, const std::vector& after, + int steps) { std::printf(" [%s] %d 步：\n", label, steps); for (std::size_t i = 0; i < before.size(); ++i) @@ -39,16 +40,18 @@ void report(const char* label, const std::vector& before, continue; } // ULP 位移 = |1/p - 1/a| 不便；直接用 half 网格步长推：按 f16 尾数 - auto ulp = [](Scalar v) -> Scalar { + auto ulp = [](Scalar v) -> Scalar + { v = std::fabs(v); if (v == 0) return std::numeric_limits::min(); - int e; std::frexp(v, &e); // v = f * 2^e, 0.5<=f<1 + int e; + std::frexp(v, &e); // v = f * 2^e, 0.5<=f<1 return std::ldexp(1.0f, e - 11); // f16 尾数 10 bit + 符号位对齐 }; const Scalar u = ulp(b); const Scalar steps_down = (b - a) / u; // 正 = 掉网格 - std::printf(" p0=%-12.7g → %-12.7g Δ/ULP=%-8.3f 相对=%-10.4g\n", - b, a, steps_down, (a - b) / b); + std::printf(" p0=%-12.7g → %-12.7g Δ/ULP=%-8.3f 相对=%-10.4g\n", b, a, steps_down, + (a - b) / b); } } @@ -69,24 +72,31 @@ void probe_decay(nn::ComputeEngine& eng, const char* label, Precision P) Matrix m(1, 6); const Scalar init[6] = {1.0f, 0.5f, 0.077f, -0.077f, 0.036f, 2.0f}; - for (std::size_t i = 0; i < 6; ++i) m.span()[i] = init[i]; + for (std::size_t i = 0; i < 6; ++i) + m.span()[i] = init[i]; auto p_r = eng.from_matrix(m, P); - if (!p_r) { std::printf(" [%s] from_matrix 失败\n", label); return; } + if (!p_r) + { + std::printf(" [%s] from_matrix 失败\n", label); + return; + } auto p = std::move(*p_r); const auto before = download(eng, p); for (int s = 0; s < steps; ++s) { - auto r = nn::dsl::compute_into(eng, - nn::dsl::leaf(p) * nn::dsl::rparam(factor), p); - if (!r) { std::printf(" [%s] decay compute_into 失败: %s\n", - label, r.error().message.c_str()); return; } + auto r = nn::dsl::compute_into(eng, nn::dsl::leaf(p) * nn::dsl::rparam(factor), p); + if (!r) + { + std::printf(" [%s] decay compute_into 失败: %s\n", label, r.error().message.c_str()); + return; + } } const auto after = download(eng, p); - std::printf("[decay] %s target=%s factor=%.7f（应为“不变”：|Δ| 远小于 1 ULP）\n", - label, P == Precision::F16 ? "f16" : "f32", factor); + std::printf("[decay] %s target=%s factor=%.7f（应为“不变”：|Δ| 远小于 1 ULP）\n", label, + P == Precision::F16 ? "f16" : "f32", factor); report("decay", before, after, steps); } @@ -98,10 +108,15 @@ void probe_add(nn::ComputeEngine& eng, const char* label, Precision P) Matrix m(1, 6); const Scalar init[6] = {1.0f, 0.5f, 0.077f, -0.077f, 0.036f, 2.0f}; - for (std::size_t i = 0; i < 6; ++i) m.span()[i] = init[i]; + for (std::size_t i = 0; i < 6; ++i) + m.span()[i] = init[i]; auto p_r = eng.from_matrix(m, P); - if (!p_r) { std::printf(" [%s] from_matrix 失败\n", label); return; } + if (!p_r) + { + std::printf(" [%s] from_matrix 失败\n", label); + return; + } auto p = std::move(*p_r); const auto before = download(eng, p); @@ -110,18 +125,25 @@ void probe_add(nn::ComputeEngine& eng, const char* label, Precision P) Matrix dm(1, 6); // 交替 ±d：偶数步 +d，奇数步 -d（模拟 Adam 双向更新） const Scalar sign = (s % 2 == 0) ? d : -d; - for (std::size_t i = 0; i < 6; ++i) dm.span()[i] = sign; + for (std::size_t i = 0; i < 6; ++i) + dm.span()[i] = sign; auto d_r = eng.from_matrix(dm, P); - if (!d_r) { std::printf(" [%s] delta from_matrix 失败\n", label); return; } - auto r = nn::dsl::compute_into(eng, - nn::dsl::leaf(p) + nn::dsl::leaf(*d_r), p); - if (!r) { std::printf(" [%s] add compute_into 失败: %s\n", - label, r.error().message.c_str()); return; } + if (!d_r) + { + std::printf(" [%s] delta from_matrix 失败\n", label); + return; + } + auto r = nn::dsl::compute_into(eng, nn::dsl::leaf(p) + nn::dsl::leaf(*d_r), p); + if (!r) + { + std::printf(" [%s] add compute_into 失败: %s\n", label, r.error().message.c_str()); + return; + } } const auto after = download(eng, p); - std::printf("[p+delta] %s target=%s（±%.0e 交替：应为小幅随机游走、无系统性收缩）\n", - label, P == Precision::F16 ? "f16" : "f32", d); + std::printf("[p+delta] %s target=%s（±%.0e 交替：应为小幅随机游走、无系统性收缩）\n", label, + P == Precision::F16 ? "f16" : "f32", d); report("add", before, after, steps); } @@ -143,10 +165,13 @@ int main(int argc, char* argv[]) for (int i = 1; i < argc; ++i) { std::string arg = argv[i]; - if (auto dev = nn::cli::parse_gpu_option(argc, argv, i, true)) - gpu_device = *dev; + if (auto dev = nn::cli::parse_gpu_option(argc, argv, i, true)) gpu_device = *dev; else if (arg == "--gpu") gpu_device.clear(); - else { std::fprintf(stderr, "未知参数: %s\n", arg.c_str()); return 2; } + else + { + std::fprintf(stderr, "未知参数: %s\n", arg.c_str()); + return 2; + } } std::printf("decay_probe：f16 权重衰减乘法 / 加法写回的舍入行为\n"); @@ -160,7 +185,7 @@ int main(int argc, char* argv[]) #ifdef NN_HAS_VULKAN { auto backend_owner = nn::GpuBackend::shared(); - nn::GpuBackend& backend = *backend_owner; + nn::GpuBackend& backend = *backend_owner; auto init_r = backend.initialize(gpu_device); if (!init_r) { @@ -173,7 +198,8 @@ int main(int argc, char* argv[]) run_all(adapter, "gpu"); } #else - (void)argc; (void)argv; + (void)argc; + (void)argv; std::printf("[SKIP] 需要 Vulkan SDK。\n"); return 77; #endif diff --git a/research/ternary_scale/probe_scale.cpp b/research/ternary_scale/probe_scale.cpp index 278f98bb..1cf27eef 100644 --- a/research/ternary_scale/probe_scale.cpp +++ b/research/ternary_scale/probe_scale.cpp @@ -32,8 +32,19 @@ using nn::Scalar; namespace { -enum class Rule { AbsMean, AbsMax }; -enum class Gran { Tensor, Row, Blk256, Blk128, Blk64 }; +enum class Rule +{ + AbsMean, + AbsMax +}; +enum class Gran +{ + Tensor, + Row, + Blk256, + Blk128, + Blk64 +}; constexpr std::size_t kOutCols = 4; // 输出误差用的随机 x 列数 constexpr int kNumRules = 2; @@ -43,11 +54,16 @@ const char* gran_name(Gran g) { switch (g) { - case Gran::Tensor: return "tensor"; - case Gran::Row: return "row"; - case Gran::Blk256: return "blk256"; - case Gran::Blk128: return "blk128"; - case Gran::Blk64: return "blk64"; + case Gran::Tensor: + return "tensor"; + case Gran::Row: + return "row"; + case Gran::Blk256: + return "blk256"; + case Gran::Blk128: + return "blk128"; + case Gran::Blk64: + return "blk64"; } return "?"; } @@ -56,10 +72,14 @@ std::size_t gran_block(Gran g) { switch (g) { - case Gran::Blk256: return 256; - case Gran::Blk128: return 128; - case Gran::Blk64: return 64; - default: return 0; + case Gran::Blk256: + return 256; + case Gran::Blk128: + return 128; + case Gran::Blk64: + return 64; + default: + return 0; } } @@ -90,8 +110,9 @@ Raw evaluate(const Matrix& m, const std::vector& x, Rule rule, Gran gran const auto w = m.span(); const std::size_t n = w.size(); const std::size_t cols = m.cols(); - const std::size_t block = (gran == Gran::Row) ? cols - : (gran == Gran::Tensor) ? n : gran_block(gran); + const std::size_t block = (gran == Gran::Row) ? cols + : (gran == Gran::Tensor) ? n + : gran_block(gran); std::vector wq(n, Scalar{0}); Raw r; @@ -103,7 +124,8 @@ Raw evaluate(const Matrix& m, const std::vector& x, Rule rule, Gran gran if (rule == Rule::AbsMean) { double s = 0.0; - for (std::size_t i = begin; i < end; ++i) s += std::fabs(static_cast(w[i])); + for (std::size_t i = begin; i < end; ++i) + s += std::fabs(static_cast(w[i])); tau = s / static_cast(end - begin); } else @@ -133,7 +155,7 @@ Raw evaluate(const Matrix& m, const std::vector& x, Rule rule, Gran gran for (std::size_t c = 0; c < cols; ++c) { const double xv = static_cast(x[c * kOutCols + j]); - y += static_cast(w[rr * cols + c]) * xv; + y += static_cast(w[rr * cols + c]) * xv; yh += static_cast(wq[rr * cols + c]) * xv; } const double d = yh - y; @@ -144,13 +166,12 @@ Raw evaluate(const Matrix& m, const std::vector& x, Rule rule, Gran gran return r; } -constexpr Gran kGrans[kNumGrans] = {Gran::Tensor, Gran::Row, Gran::Blk256, Gran::Blk128, Gran::Blk64}; +constexpr Gran kGrans[kNumGrans] = {Gran::Tensor, Gran::Row, Gran::Blk256, Gran::Blk128, + Gran::Blk64}; constexpr Rule kRules[kNumRules] = {Rule::AbsMean, Rule::AbsMax}; inline double pct(double num, double den) -{ - return den > 0.0 ? 100.0 * std::sqrt(num / den) : 0.0; -} +{ return den > 0.0 ? 100.0 * std::sqrt(num / den) : 0.0; } // 三值分布的信息熵（bit/权重）：把"多少权重存活"翻译成"实际携带多少信息" inline double entropy_bits(const Raw& r) @@ -184,33 +205,50 @@ int main(int argc, char** argv) nn::CpuEngine eng; auto spec_r = nn::peek_model_spec(path); - if (!spec_r) { std::printf("peek 失败: %s\n", spec_r.error().message.c_str()); continue; } + if (!spec_r) + { + std::printf("peek 失败: %s\n", spec_r.error().message.c_str()); + continue; + } const nn::ModelSpec& spec = *spec_r; std::printf("spec: type=%u d_model=%zu d_ff=%zu layers=%zu vocab=%zu layer_dims=%zu\n", - static_cast(spec.type), spec.d_model, spec.d_ff, - spec.num_layers, spec.vocab_size, spec.layer_dims.size()); + static_cast(spec.type), spec.d_model, spec.d_ff, spec.num_layers, + spec.vocab_size, spec.layer_dims.size()); nn::Result built = (spec.is_gpt() || spec.is_alibi_gpt()) - ? nn::build_gpt_model_from_spec(eng, spec) - : nn::build_mnist_model_from_spec(eng, spec); - if (!built) { std::printf("build 失败: %s\n", built.error().message.c_str()); continue; } + ? nn::build_gpt_model_from_spec(eng, spec) + : nn::build_mnist_model_from_spec(eng, spec); + if (!built) + { + std::printf("build 失败: %s\n", built.error().message.c_str()); + continue; + } nn::Model model = std::move(*built); auto loaded = nn::load_model(path, model); - if (!loaded) { std::printf("load 失败: %s\n", loaded.error().message.c_str()); continue; } + if (!loaded) + { + std::printf("load 失败: %s\n", loaded.error().message.c_str()); + continue; + } auto params = model.parameters(); Raw agg[kNumRules][kNumGrans]; std::size_t total_w = 0, n_mat = 0, n_skip = 0; - std::printf("\n%-4s %-13s %-11s %-10s %-9s [absmean@row] [absmax@row] [absmean@256] [absmax@256]\n", + std::printf("\n%-4s %-13s %-11s %-10s %-9s [absmean@row] [absmax@row] [absmean@256] " + "[absmax@256]\n", "#", "shape", "mean|w|", "max|w|", "max/mean"); for (std::size_t pi = 0; pi < params.size(); ++pi) { const nn::Tensor& t = params[pi].get(); - if (t.cols() < 2) { ++n_skip; continue; } // 跳过 bias（(out,1)） + if (t.cols() < 2) + { + ++n_skip; + continue; + } // 跳过 bias（(out,1)） auto mr = eng.to_matrix(t, Precision::F32); if (!mr) continue; const Matrix& m = *mr; @@ -218,7 +256,8 @@ int main(int argc, char** argv) std::mt19937 rng(20261010u + static_cast(pi)); std::normal_distribution nd(0.f, 1.f); std::vector x(m.cols() * kOutCols); - for (auto& v : x) v = nd(rng); + for (auto& v : x) + v = nd(rng); double s = 0.0, mx = 0.0; for (auto wv : m.span()) @@ -234,7 +273,7 @@ int main(int argc, char** argv) for (int gi = 0; gi < kNumGrans; ++gi) { const Raw r = evaluate(m, x, kRules[ri], kGrans[gi]); - z[ri][gi] = 100.0 * r.zero / static_cast(m.size()); + z[ri][gi] = 100.0 * r.zero / static_cast(m.size()); rl[ri][gi] = pct(r.wnum, r.wden); agg[ri][gi].zero += r.zero; agg[ri][gi].pos += r.pos; @@ -248,29 +287,25 @@ int main(int argc, char** argv) total_w += m.size(); ++n_mat; - const std::string shape = - std::to_string(m.rows()) + "x" + std::to_string(m.cols()); + const std::string shape = std::to_string(m.rows()) + "x" + std::to_string(m.cols()); std::printf("%-4zu %-13s %-11.5f %-10.5f %-9.2f " "%5.1f%%/%5.1f%% %5.1f%%/%5.1f%% %5.1f%%/%5.1f%% %5.1f%%/%5.1f%%\n", - pi, shape.c_str(), mean, mx, mean > 0.0 ? mx / mean : 0.0, - z[0][1], rl[0][1], z[1][1], rl[1][1], - z[0][2], rl[0][2], z[1][2], rl[1][2]); + pi, shape.c_str(), mean, mx, mean > 0.0 ? mx / mean : 0.0, z[0][1], + rl[0][1], z[1][1], rl[1][1], z[0][2], rl[0][2], z[1][2], rl[1][2]); } std::printf("\n-- 汇总（按元素数加权；矩阵 %zu 个 / 权重 %zu 个，跳过 bias %zu 个）--\n", n_mat, total_w, n_skip); - std::printf("%-8s %-30s %-30s\n", "粒度", - "absmean zero%/relerr%/bits", "absmax zero%/relerr%/bits"); + std::printf("%-8s %-30s %-30s\n", "粒度", "absmean zero%/relerr%/bits", + "absmax zero%/relerr%/bits"); for (int gi = 0; gi < kNumGrans; ++gi) { std::printf("%-8s %8.1f%% / %8.1f%% / %5.3f %8.1f%% / %8.1f%% / %5.3f\n", gran_name(kGrans[gi]), agg[0][gi].zero * 100.0 / static_cast(total_w), - pct(agg[0][gi].wnum, agg[0][gi].wden), - entropy_bits(agg[0][gi]), + pct(agg[0][gi].wnum, agg[0][gi].wden), entropy_bits(agg[0][gi]), agg[1][gi].zero * 100.0 / static_cast(total_w), - pct(agg[1][gi].wnum, agg[1][gi].wden), - entropy_bits(agg[1][gi])); + pct(agg[1][gi].wnum, agg[1][gi].wden), entropy_bits(agg[1][gi])); } std::printf("\n-- 层输出相对误差（%%，全模型加权）--\n"); diff --git a/research/ternary_ste/poc_binary_ste.cpp b/research/ternary_ste/poc_binary_ste.cpp index 7f1e4362..1c9dc6b4 100644 --- a/research/ternary_ste/poc_binary_ste.cpp +++ b/research/ternary_ste/poc_binary_ste.cpp @@ -41,7 +41,8 @@ int main() Wt.set_value(r, c, (td(rng) == 0 ? -0.5f : 0.5f)); Matrix X(IN, NB), Y(OUT, NB); - for (auto& v : X.span()) v = nd(rng); + for (auto& v : X.span()) + v = nd(rng); Wt.multiply_to(Y, X); Tensor Xt = NN_CHECK(eng.from_matrix(X)); @@ -54,18 +55,17 @@ int main() { const Matrix Wh = NN_CHECK(eng.to_matrix(W, Precision::F32)); Scalar gamma = 0.f; - for (auto v : Wh.span()) gamma += std::fabs(v); + for (auto v : Wh.span()) + gamma += std::fabs(v); gamma /= static_cast(IN * OUT); // Wq = sign(W) —— 纯 1-bit（二值）量化 - Tensor Wq = NN_CHECK(nn::dsl::compute(eng, - nn::dsl::select(nn::dsl::leaf(W) > nn::dsl::rparam(0.f), - Scalar{1}, Scalar{-1}), + Tensor Wq = NN_CHECK(nn::dsl::compute( + eng, nn::dsl::select(nn::dsl::leaf(W) > nn::dsl::rparam(0.f), Scalar{1}, Scalar{-1}), OUT, IN, Precision::F32)); - Tensor Yhat = NN_CHECK(nn::dsl::compute(eng, - nn::dsl::matmul(Wq, Xt) * nn::dsl::rparam(gamma), - OUT, NB, Precision::F32)); + Tensor Yhat = NN_CHECK(nn::dsl::compute( + eng, nn::dsl::matmul(Wq, Xt) * nn::dsl::rparam(gamma), OUT, NB, Precision::F32)); const Matrix Yh = NN_CHECK(eng.to_matrix(Yhat, Precision::F32)); Scalar loss = 0.f; @@ -79,13 +79,14 @@ int main() last_loss = loss; Tensor dY = NN_CHECK(nn::dsl::compute(eng, - (nn::dsl::leaf(Yhat) - nn::dsl::leaf(Yt)) * nn::dsl::rparam(Scalar{2} / Scalar(OUT * NB)), - OUT, NB, Precision::F32)); - Tensor dW = NN_CHECK(nn::dsl::compute(eng, - nn::dsl::matmul(dY, Xt, false, true) * nn::dsl::rparam(gamma), - OUT, IN, Precision::F32)); - NN_CHECK(nn::dsl::compute_into(eng, - nn::dsl::leaf(W) - nn::dsl::leaf(dW) * nn::dsl::rparam(lr), W)); + (nn::dsl::leaf(Yhat) - nn::dsl::leaf(Yt)) * + nn::dsl::rparam(Scalar{2} / Scalar(OUT * NB)), + OUT, NB, Precision::F32)); + Tensor dW = NN_CHECK( + nn::dsl::compute(eng, nn::dsl::matmul(dY, Xt, false, true) * nn::dsl::rparam(gamma), + OUT, IN, Precision::F32)); + NN_CHECK(nn::dsl::compute_into( + eng, nn::dsl::leaf(W) - nn::dsl::leaf(dW) * nn::dsl::rparam(lr), W)); if (step % 500 == 0 || step == STEPS - 1) std::printf("step %4d mse=%.6f gamma=%.5f\n", step, loss, gamma); @@ -95,12 +96,11 @@ int main() std::size_t match = 0; for (std::size_t r = 0; r < OUT; ++r) for (std::size_t c = 0; c < IN; ++c) - if ((Wh.at(r, c) > 0.f) == (Wt.at(r, c) > 0.f)) - ++match; + if ((Wh.at(r, c) > 0.f) == (Wt.at(r, c) > 0.f)) ++match; - std::printf("\nmse: first=%.6f final=%.6f (降 %.1f%%)\n", - first_loss, last_loss, 100.0 * (1.0 - last_loss / first_loss)); - std::printf("二值符号吻合率: %zu / %zu = %.1f%%\n", - match, IN * OUT, 100.0 * static_cast(match) / double(IN * OUT)); + std::printf("\nmse: first=%.6f final=%.6f (降 %.1f%%)\n", first_loss, last_loss, + 100.0 * (1.0 - last_loss / first_loss)); + std::printf("二值符号吻合率: %zu / %zu = %.1f%%\n", match, IN * OUT, + 100.0 * static_cast(match) / double(IN * OUT)); return 0; } diff --git a/research/ternary_ste/poc_ternary_ste.cpp b/research/ternary_ste/poc_ternary_ste.cpp index 0ba11e4c..7e0ca027 100644 --- a/research/ternary_ste/poc_ternary_ste.cpp +++ b/research/ternary_ste/poc_ternary_ste.cpp @@ -44,7 +44,8 @@ int main() Wt.set_value(r, c, static_cast(td(rng) - 1) * 0.5f); Matrix X(IN, NB), Y(OUT, NB); - for (auto& v : X.span()) v = nd(rng); + for (auto& v : X.span()) + v = nd(rng); Wt.multiply_to(Y, X); Tensor Xt = NN_CHECK(eng.from_matrix(X)); @@ -61,23 +62,23 @@ int main() // Layer 化时改为 dsl 的 row_reduce_sum + row_broadcast（逐行 absmean）。 const Matrix Wh = NN_CHECK(eng.to_matrix(W, Precision::F32)); Scalar gamma = 0.f; - for (auto v : Wh.span()) gamma += std::fabs(v); + for (auto v : Wh.span()) + gamma += std::fabs(v); gamma /= static_cast(IN * OUT); const Scalar thr = 0.5f * gamma; // Wq = select(W > thr, +1, select(W < -thr, -1, 0)) —— 纯 DSL 三值量化 // （DSL 没有 round/floor；在 [-1,1] 区间内这条 select 链等价于 RoundClip(w/γ, -1, 1)） - Tensor Wq = NN_CHECK(nn::dsl::compute(eng, - nn::dsl::select(nn::dsl::leaf(W) > nn::dsl::rparam(thr), - Scalar{1}, - nn::dsl::select(nn::dsl::leaf(W) < nn::dsl::rparam(-thr), - Scalar{-1}, Scalar{0})), + Tensor Wq = NN_CHECK(nn::dsl::compute( + eng, + nn::dsl::select( + nn::dsl::leaf(W) > nn::dsl::rparam(thr), Scalar{1}, + nn::dsl::select(nn::dsl::leaf(W) < nn::dsl::rparam(-thr), Scalar{-1}, Scalar{0})), OUT, IN, Precision::F32)); // Ŷ = γ · (Wq · X)（matmul 段 + 缩放融合为单 kernel） - Tensor Yhat = NN_CHECK(nn::dsl::compute(eng, - nn::dsl::matmul(Wq, Xt) * nn::dsl::rparam(gamma), - OUT, NB, Precision::F32)); + Tensor Yhat = NN_CHECK(nn::dsl::compute( + eng, nn::dsl::matmul(Wq, Xt) * nn::dsl::rparam(gamma), OUT, NB, Precision::F32)); // MSE + 其雅可比 dY = 2(Ŷ - Y)/(OUT·NB) const Matrix Yh = NN_CHECK(eng.to_matrix(Yhat, Precision::F32)); @@ -92,16 +93,17 @@ int main() last_loss = loss; Tensor dY = NN_CHECK(nn::dsl::compute(eng, - (nn::dsl::leaf(Yhat) - nn::dsl::leaf(Yt)) * nn::dsl::rparam(Scalar{2} / Scalar(OUT * NB)), - OUT, NB, Precision::F32)); + (nn::dsl::leaf(Yhat) - nn::dsl::leaf(Yt)) * + nn::dsl::rparam(Scalar{2} / Scalar(OUT * NB)), + OUT, NB, Precision::F32)); // STE：把 Wq 当常数，梯度直接穿回 latent W（dW = γ · dY · Xᵀ） - Tensor dW = NN_CHECK(nn::dsl::compute(eng, - nn::dsl::matmul(dY, Xt, false, true) * nn::dsl::rparam(gamma), - OUT, IN, Precision::F32)); + Tensor dW = NN_CHECK( + nn::dsl::compute(eng, nn::dsl::matmul(dY, Xt, false, true) * nn::dsl::rparam(gamma), + OUT, IN, Precision::F32)); - NN_CHECK(nn::dsl::compute_into(eng, - nn::dsl::leaf(W) - nn::dsl::leaf(dW) * nn::dsl::rparam(lr), W)); + NN_CHECK(nn::dsl::compute_into( + eng, nn::dsl::leaf(W) - nn::dsl::leaf(dW) * nn::dsl::rparam(lr), W)); if (step % 500 == 0 || step == STEPS - 1) std::printf("step %4d mse=%.6f gamma=%.5f\n", step, loss, gamma); @@ -110,7 +112,8 @@ int main() // ── 验收：量化后的符号/零帽与教师三元模式的重合度 ── const Matrix Wh = NN_CHECK(eng.to_matrix(W, Precision::F32)); Scalar gamma = 0.f; - for (auto v : Wh.span()) gamma += std::fabs(v); + for (auto v : Wh.span()) + gamma += std::fabs(v); gamma /= static_cast(IN * OUT); std::size_t nz = 0, match = 0; @@ -124,13 +127,12 @@ int main() Scalar q = 0.f; if (w > 0.5f * gamma) q = 1.f; else if (w < -0.5f * gamma) q = -1.f; - if ((q > 0.f && t > 0.f) || (q < 0.f && t < 0.f) || q == 0.f) - ++match; + if ((q > 0.f && t > 0.f) || (q < 0.f && t < 0.f) || q == 0.f) ++match; } - std::printf("\nmse: first=%.6f final=%.6f (降 %.1f%%)\n", - first_loss, last_loss, 100.0 * (1.0 - last_loss / first_loss)); - std::printf("教师非零项符号/零帽吻合率: %zu / %zu = %.1f%%\n", - match, nz, 100.0 * static_cast(match) / static_cast(nz)); + std::printf("\nmse: first=%.6f final=%.6f (降 %.1f%%)\n", first_loss, last_loss, + 100.0 * (1.0 - last_loss / first_loss)); + std::printf("教师非零项符号/零帽吻合率: %zu / %zu = %.1f%%\n", match, nz, + 100.0 * static_cast(match) / static_cast(nz)); return 0; } diff --git a/src/adam_quant_test.cpp b/src/adam_quant_test.cpp index 39df3d99..47520f9f 100644 --- a/src/adam_quant_test.cpp +++ b/src/adam_quant_test.cpp @@ -49,8 +49,7 @@ void test_state_bytes(nn::ComputeEngine& eng, const char* tag) const std::size_t q_state = q.state_bytes(); const std::size_t a_state = elems * 4u * 2u; // m/v 各 f32 std::printf(" AdamWQuant 状态 = %zu B（%.2f B/权重）; f32 AdamW = %zu B（8.00 B/权重）\n", - q_state, static_cast(q_state) / static_cast(elems), - a_state); + q_state, static_cast(q_state) / static_cast(elems), a_state); CHECK(q_state < a_state, "量化状态比 f32 状态小"); // 上界：2 B/权重（f16 m）+ 2 B/权重（f16 v）+ 尺度（4 B/行 × 2 个） const std::size_t expect_max = elems * 4u + 64u * 4u * 2u; @@ -111,24 +110,22 @@ void test_scales(nn::ComputeEngine& eng, const char* tag) CHECK(row1_nonzero, "小量级行的状态未被下溢吃掉（动态尺度生效）"); // 尺度的量级应随行的梯度量级走（行 0 比行 1 大约 1e3 倍） CHECK(sm.at(0, 0) > 100.0f * sm.at(1, 0), "行尺度随梯度量级（行0 ≫ 行1）"); - std::printf(" scale_m: row0=%.3g row1=%.3g（比值 %.0f）\n", - static_cast(sm.at(0, 0)), static_cast(sm.at(1, 0)), - static_cast(sm.at(0, 0) / sm.at(1, 0))); + std::printf(" scale_m: row0=%.3g row1=%.3g（比值 %.0f）\n", static_cast(sm.at(0, 0)), + static_cast(sm.at(1, 0)), static_cast(sm.at(0, 0) / sm.at(1, 0))); } // ── [3][4] 端到端：与 f32 AdamW 同量级收敛 + 参数移动 ──────────────────── -bool train_mlp(nn::ComputeEngine& eng, const char* opt_name, int steps, - nn::Scalar& loss_out, std::vector& w_before_out, - std::vector& w_after_out) +bool train_mlp(nn::ComputeEngine& eng, const char* opt_name, int steps, nn::Scalar& loss_out, + std::vector& w_before_out, std::vector& w_after_out) { nn::PrecisionProfile prof; // 默认 f32 - nn::Result m_r = nn::build_mnist_mlp_model( - eng, nn::MNIST_LAYER_DIMS, nn::NormType::LayerNorm, prof); + nn::Result m_r = + nn::build_mnist_mlp_model(eng, nn::MNIST_LAYER_DIMS, nn::NormType::LayerNorm, prof); if (!m_r) return false; nn::Model& model = *m_r; - auto os = nn::create_optimizer(opt_name, eng, model.parameters(), - model.param_gradients(), 0.02f, 0.0f, prof); + auto os = nn::create_optimizer(opt_name, eng, model.parameters(), model.param_gradients(), + 0.02f, 0.0f, prof); if (!os) return false; std::unique_ptr opt = std::move(os); @@ -157,12 +154,21 @@ bool train_mlp(nn::ComputeEngine& eng, const char* opt_name, int steps, { if (auto r = model.zero_grad(); !r) return false; auto pred = model.forward(Xt); - if (!pred) { std::printf(" forward 失败: %s\n", pred.error().message.c_str()); return false; } + if (!pred) + { + std::printf(" forward 失败: %s\n", pred.error().message.c_str()); + return false; + } std::vector labels(B); - for (std::size_t b = 0; b < B; ++b) labels[b] = b % C; + for (std::size_t b = 0; b < B; ++b) + labels[b] = b % C; std::vector mask(B, 1.0f); auto ls = loss.forward_sparse(eng, *pred, labels, mask, C); - if (!ls) { std::printf(" loss 失败: %s\n", ls.error().message.c_str()); return false; } + if (!ls) + { + std::printf(" loss 失败: %s\n", ls.error().message.c_str()); + return false; + } last = *ls; if (!std::isfinite(last)) return false; auto gl = loss.backward(); @@ -188,8 +194,8 @@ void test_convergence(nn::ComputeEngine& eng, const char* tag) CHECK(ok_a, "f32 AdamW 训练跑通"); CHECK(ok_q, "AdamWQuant 训练跑通"); if (!ok_a || !ok_q) return; - std::printf(" 最终 loss: AdamW=%.6f AdamWQuant=%.6f\n", - static_cast(loss_adam), static_cast(loss_quant)); + std::printf(" 最终 loss: AdamW=%.6f AdamWQuant=%.6f\n", static_cast(loss_adam), + static_cast(loss_quant)); CHECK(std::isfinite(loss_quant), "AdamWQuant 的 loss 有限"); // 收敛量级：能显著下降，且与 f32 AdamW 同量级（不要求逐位：状态量化必然有差） CHECK(loss_quant < 0.5f, "AdamWQuant 收敛到小 loss"); @@ -200,12 +206,16 @@ void test_convergence(nn::ComputeEngine& eng, const char* tag) for (std::size_t i = 0; i < wa_q.size(); ++i) moved += std::fabs(static_cast(wa_q[i]) - static_cast(wb_q[i])); CHECK(moved > 1e-3, "AdamWQuant 下参数确实移动（非静默冻结）"); - std::printf(" 首层权重 L1 位移: AdamW=%.4f AdamWQuant=%.4f\n", [&] { - double s = 0.0; - for (std::size_t i = 0; i < wa_a.size(); ++i) - s += std::fabs(static_cast(wa_a[i]) - static_cast(wb_a[i])); - return s; - }(), moved); + std::printf( + " 首层权重 L1 位移: AdamW=%.4f AdamWQuant=%.4f\n", + [&] + { + double s = 0.0; + for (std::size_t i = 0; i < wa_a.size(); ++i) + s += std::fabs(static_cast(wa_a[i]) - static_cast(wb_a[i])); + return s; + }(), + moved); } // ── [5] 首步零梯度不产生 NaN ────────────────────────────────────────────── @@ -257,7 +267,6 @@ int main(int argc, char* argv[]) test_convergence(eng, "CPU"); } - if (g_failures == 0) - std::puts("adam_quant_test: 全部通过"); + if (g_failures == 0) std::puts("adam_quant_test: 全部通过"); return g_failures == 0 ? 0 : 1; } diff --git a/src/assert_always_test.cpp b/src/assert_always_test.cpp index ad1dc2a8..cbf26c51 100644 --- a/src/assert_always_test.cpp +++ b/src/assert_always_test.cpp @@ -39,7 +39,8 @@ int g_fail = 0; -namespace { +namespace +{ constexpr int kSkip = -1; // std::system 无法启动子进程（区别于任何真实退出码） @@ -47,7 +48,7 @@ std::string g_self; // argv[0]：子进程自举用 const char* kLogPath = "nn_assert_always_child.log"; // 子进程 stderr 落地 const char* kTagAssert = "NN_ASSERT_ALWAYS failed at"; -const char* kTagFatal = "NN_FATAL failed at"; +const char* kTagFatal = "NN_FATAL failed at"; int run_self(const char* arg) { @@ -116,8 +117,7 @@ bool has_fatal_line(const std::string& log, const char* tag, const char* msg) } // 一个负例：非零退出 + 打印 tag/file:line/消息 -void check_crash_case(const char* arg, const char* tag, const char* msg, - const char* what) +void check_crash_case(const char* arg, const char* tag, const char* msg, const char* what) { const int rc = normalize_rc(run_self(arg)); if (rc == kSkip) return; @@ -126,8 +126,7 @@ void check_crash_case(const char* arg, const char* tag, const char* msg, const std::string log = read_log(); const bool ok = has_fatal_line(log, tag, msg); CHECK(ok, (std::string(what) + "：stderr 含 tag/file:line/消息").c_str()); - if (!ok) - std::printf(" 子进程 stderr 捕获：<<<%s>>>\n", log.c_str()); + if (!ok) std::printf(" 子进程 stderr 捕获：<<<%s>>>\n", log.c_str()); } } // namespace @@ -197,17 +196,14 @@ int main(int argc, char* argv[]) CHECK(true, "NN_ASSERT_ALWAYS 条件成立时不终止"); // 负例三件（详见文件头） - check_crash_case("--crash-matrix-at", kTagFatal, - "MatrixT::at index out of range", + check_crash_case("--crash-matrix-at", kTagFatal, "MatrixT::at index out of range", "MatrixT::at 越界 → 非零退出"); check_crash_case("--crash-model-unbound", kTagFatal, "未绑定引擎（构造时引擎必须由 shared_ptr 管理）", "未绑定 Model::engine() → 非零退出"); - check_crash_case("--crash-empty-logits", kTagAssert, - "sample_next_token_: empty logits", + check_crash_case("--crash-empty-logits", kTagAssert, "sample_next_token_: empty logits", "sample_next_token_ 空 logits → 非零退出"); - check_crash_case("--crash-layer-uninit", kTagFatal, - "未绑定引擎（先调 init(engine)）", + check_crash_case("--crash-layer-uninit", kTagFatal, "未绑定引擎（先调 init(engine)）", "未 init 的层 forward → 非零退出"); std::remove(kLogPath); diff --git a/src/attn_consistency_test.cpp b/src/attn_consistency_test.cpp index 5676cdce..dce04be2 100644 --- a/src/attn_consistency_test.cpp +++ b/src/attn_consistency_test.cpp @@ -28,18 +28,18 @@ #include #include -using nn::Scalar; -using nn::Matrix; -using nn::Tensor; using nn::ComputeEngine; using nn::GPTModel; +using nn::Matrix; using nn::PosEncodingType; +using nn::Scalar; +using nn::Tensor; -namespace { +namespace +{ // ── 辅助：从 (vocab, seq) logits 矩阵取第 col 列 → 向量 ── -std::vector extract_logits_column( - const Matrix& logits, std::size_t col) +std::vector extract_logits_column(const Matrix& logits, std::size_t col) { const std::size_t vocab = logits.rows(); std::vector out(vocab); @@ -50,9 +50,7 @@ std::vector extract_logits_column( // ── 辅助：从 (vocab, 1) logits 矩阵取第 0 列 → 向量 ── std::vector extract_logits_single(const Matrix& logits) -{ - return extract_logits_column(logits, 0); -} +{ return extract_logits_column(logits, 0); } // ── 辅助：两个向量的最大绝对差 ── Scalar max_abs_diff(const std::vector& a, const std::vector& b) @@ -83,35 +81,32 @@ void print_vec_head(const std::vector& v, std::size_t n = 8) // ── 单引擎一致性测试 ── // 返回 true 表示通过（forward 与 forward_step 输出一致） -bool test_consistency( - ComputeEngine& engine, - const std::string& engine_name, - std::size_t seq_len, - std::size_t num_layers, - PosEncodingType pos_enc, - const std::string& pos_enc_name, - bool use_batch = true) +bool test_consistency(ComputeEngine& engine, const std::string& engine_name, std::size_t seq_len, + std::size_t num_layers, PosEncodingType pos_enc, + const std::string& pos_enc_name, bool use_batch = true) { // 模型超参（小规模，加速测试） const std::size_t vocab_size = 100; - const std::size_t d_model = 32; - const std::size_t num_heads = 4; - const std::size_t d_ff = 64; + const std::size_t d_model = 32; + const std::size_t num_heads = 4; + const std::size_t d_ff = 64; std::cout << "\n────────────────────────────────────────────────────────\n" - << "引擎: " << engine_name - << " | 位置编码: " << pos_enc_name - << " | seq=" << seq_len - << " | layers=" << num_layers - << " | batch=" << (use_batch ? "on" : "off") << "\n" + << "引擎: " << engine_name << " | 位置编码: " << pos_enc_name << " | seq=" << seq_len + << " | layers=" << num_layers << " | batch=" << (use_batch ? "on" : "off") << "\n" << "────────────────────────────────────────────────────────\n"; // ── 1. 构建模型 ── std::cout << "[1/4] 构建 GPTModel..." << std::flush; - GPTModel model( - vocab_size, d_model, seq_len, - num_heads, d_ff, num_layers, pos_enc); - { auto r = model.init(engine); if (!r) { std::cerr << "GPTModel init 失败: " << r.error().message << "\n"; return false; } } + GPTModel model(vocab_size, d_model, seq_len, num_heads, d_ff, num_layers, pos_enc); + { + auto r = model.init(engine); + if (!r) + { + std::cerr << "GPTModel init 失败: " << r.error().message << "\n"; + return false; + } + } std::cout << " 完成\n"; // ── 2. 生成随机 token 序列 ── @@ -119,7 +114,8 @@ bool test_consistency( std::mt19937_64 rng(12345); std::uniform_int_distribution tok_dist(0, vocab_size - 1); std::vector tokens(seq_len); - for (auto& t : tokens) t = tok_dist(rng); + for (auto& t : tokens) + t = tok_dist(rng); std::cout << " 完成 (tokens:"; for (std::size_t i = 0; i < std::min(8, seq_len); ++i) std::cout << " " << tokens[i]; @@ -169,8 +165,8 @@ bool test_consistency( std::cout << "\n ❌ to_matrix 失败: " << fwd_logits_m.error().message << "\n"; return false; } - std::cout << " 完成 (logits: " << fwd_logits_m->rows() - << " × " << fwd_logits_m->cols() << ")\n"; + std::cout << " 完成 (logits: " << fwd_logits_m->rows() << " × " << fwd_logits_m->cols() + << ")\n"; // ── 4. 逐 token forward_step ── std::cout << "[4/4] 运行 forward_step（增量推理）并对比..." << std::flush; @@ -184,9 +180,17 @@ bool test_consistency( k_caches.push_back(engine.create_tensor(seq_len, d_model)); v_caches.push_back(engine.create_tensor(seq_len, d_model)); auto r1 = engine.zero(k_caches.back()); - if (!r1) { std::cout << "\n ❌ zero(k_cache) 失败\n"; return false; } + if (!r1) + { + std::cout << "\n ❌ zero(k_cache) 失败\n"; + return false; + } auto r2 = engine.zero(v_caches.back()); - if (!r2) { std::cout << "\n ❌ zero(v_cache) 失败\n"; return false; } + if (!r2) + { + std::cout << "\n ❌ zero(v_cache) 失败\n"; + return false; + } } Scalar total_max_diff = Scalar{0}; @@ -199,22 +203,29 @@ bool test_consistency( if (use_batch) { auto br = engine.begin_batch(); - if (!br) { std::cout << "\n ❌ begin_batch 失败\n"; return false; } + if (!br) + { + std::cout << "\n ❌ begin_batch 失败\n"; + return false; + } } - auto step_r = model.forward_step( - engine, tokens[i], i, k_caches, v_caches, i); + auto step_r = model.forward_step(engine, tokens[i], i, k_caches, v_caches, i); if (!step_r) { - std::cout << "\n ❌ forward_step 在 pos=" << i - << " 失败: " << step_r.error().message << "\n"; + std::cout << "\n ❌ forward_step 在 pos=" << i << " 失败: " << step_r.error().message + << "\n"; return false; } if (use_batch) { auto er = engine.end_batch(); - if (!er) { std::cout << "\n ❌ end_batch 失败\n"; return false; } + if (!er) + { + std::cout << "\n ❌ end_batch 失败\n"; + return false; + } } auto step_logits_m = engine.to_matrix(*step_r); @@ -241,8 +252,8 @@ bool test_consistency( if (diff > tolerance) { all_ok = false; - std::cout << "\n ❌ pos=" << i << " 不一致 (max_diff=" - << std::scientific << std::setprecision(4) << diff << ")\n"; + std::cout << "\n ❌ pos=" << i << " 不一致 (max_diff=" << std::scientific + << std::setprecision(4) << diff << ")\n"; std::cout << " forward: "; print_vec_head(fwd_col); std::cout << "\n forward_step: "; @@ -252,8 +263,8 @@ bool test_consistency( } std::cout << "\n"; - std::cout << " 最大绝对误差: " << std::scientific << std::setprecision(6) - << total_max_diff << " (位置 " << worst_pos << ")\n"; + std::cout << " 最大绝对误差: " << std::scientific << std::setprecision(6) << total_max_diff + << " (位置 " << worst_pos << ")\n"; std::cout << " 结果: " << (all_ok ? "✅ 通过" : "❌ 失败") << "\n"; return all_ok; @@ -317,28 +328,20 @@ int main(int argc, char* argv[]) bool all_pass = true; // 测试 1: Learned 位置编码（GPT 默认） - all_pass &= test_consistency( - *engine, use_gpu ? "GPU" : "CPU", - seq_len, num_layers, - PosEncodingType::Learned, "Learned", use_batch); + all_pass &= test_consistency(*engine, use_gpu ? "GPU" : "CPU", seq_len, num_layers, + PosEncodingType::Learned, "Learned", use_batch); // 测试 2: Sinusoidal 位置编码 - all_pass &= test_consistency( - *engine, use_gpu ? "GPU" : "CPU", - seq_len, num_layers, - PosEncodingType::Sinusoidal, "Sinusoidal", use_batch); + all_pass &= test_consistency(*engine, use_gpu ? "GPU" : "CPU", seq_len, num_layers, + PosEncodingType::Sinusoidal, "Sinusoidal", use_batch); // 测试 3: ALiBi（无位置嵌入，因果 + 线性偏置） - all_pass &= test_consistency( - *engine, use_gpu ? "GPU" : "CPU", - seq_len, num_layers, - PosEncodingType::ALiBi, "ALiBi", use_batch); + all_pass &= test_consistency(*engine, use_gpu ? "GPU" : "CPU", seq_len, num_layers, + PosEncodingType::ALiBi, "ALiBi", use_batch); // 测试 4: RoPE（旋转位置编码，Q/K 上施加） - all_pass &= test_consistency( - *engine, use_gpu ? "GPU" : "CPU", - seq_len, num_layers, - PosEncodingType::RoPE, "RoPE", use_batch); + all_pass &= test_consistency(*engine, use_gpu ? "GPU" : "CPU", seq_len, num_layers, + PosEncodingType::RoPE, "RoPE", use_batch); std::cout << "\n════════════════════════════════════════\n"; std::cout << " 总体结果: " << (all_pass ? "✅ 全部通过" : "❌ 存在失败"); diff --git a/src/attn_gradcheck.cpp b/src/attn_gradcheck.cpp index 77ce9c6b..64704a0e 100644 --- a/src/attn_gradcheck.cpp +++ b/src/attn_gradcheck.cpp @@ -20,33 +20,37 @@ #include #include "test_common.hpp" -using nn::Scalar; +using nn::CausalSelfAttention; +using nn::ComputeEngine; using nn::Matrix; +using nn::PosEncodingType; +using nn::Scalar; using nn::Tensor; using nn::TensorRef; -using nn::ComputeEngine; -using nn::CausalSelfAttention; -using nn::PosEncodingType; -namespace { +namespace +{ -Scalar eval_loss(ComputeEngine& engine, CausalSelfAttention& attn, - const Tensor& x, const Tensor& go) +Scalar eval_loss(ComputeEngine& engine, CausalSelfAttention& attn, const Tensor& x, + const Tensor& go) { auto y = attn.forward(x); - if (!y) { std::cerr << " [eval] forward failed: " << y.error().message << "\n"; std::abort(); } + if (!y) + { + std::cerr << " [eval] forward failed: " << y.error().message << "\n"; + std::abort(); + } auto y_m = engine.to_matrix(*y); auto go_m = engine.to_matrix(go); return dot(*y_m, *go_m); } -bool check_grad_tensor( - ComputeEngine& engine, CausalSelfAttention& attn, - const Tensor& x, const Tensor& go, - Tensor& param, const Matrix& base, const Matrix& grad_analytical, - const std::string& name, Scalar eps, Scalar tol) +bool check_grad_tensor(ComputeEngine& engine, CausalSelfAttention& attn, const Tensor& x, + const Tensor& go, Tensor& param, const Matrix& base, + const Matrix& grad_analytical, const std::string& name, Scalar eps, + Scalar tol) { bool ok = true; std::size_t bad = 0; @@ -76,14 +80,14 @@ bool check_grad_tensor( { ok = false; if (bad < 5) - std::cout << " [FAIL] " << name << "(" << r << "," << c - << ") num=" << num << " ana=" << ana << " err=" << err << "\n"; + std::cout << " [FAIL] " << name << "(" << r << "," << c << ") num=" << num + << " ana=" << ana << " err=" << err << "\n"; ++bad; } } } - std::cout << " " << name << ": " << (ok ? "OK" : "FAIL") - << " max_err=" << max_err << " bad=" << bad << "\n"; + std::cout << " " << name << ": " << (ok ? "OK" : "FAIL") << " max_err=" << max_err + << " bad=" << bad << "\n"; return ok; } @@ -101,14 +105,16 @@ int main(int argc, char* argv[]) { std::string a = argv[i]; if (a == "--tol" && i + 1 < argc) tol = static_cast(std::atof(argv[++i])); - else if (a == "--batch" && i + 1 < argc) batch = static_cast(std::atoi(argv[++i])); + else if (a == "--batch" && i + 1 < argc) + batch = static_cast(std::atoi(argv[++i])); else if (a == "--seq" && i + 1 < argc) seq = static_cast(std::atoi(argv[++i])); else if (a == "--pos-enc" && i + 1 < argc) pos_enc_name = argv[++i]; else if (a == "--doc") use_doc = true; else if (a == "--gpu") use_gpu = true; else if (a == "--help") { - std::cout << "用法: attn_gradcheck [--gpu] [--batch N] [--seq N] [--tol ] [--pos-enc learned|sinusoidal|alibi|rope] [--doc]\n"; + std::cout << "用法: attn_gradcheck [--gpu] [--batch N] [--seq N] [--tol ] " + "[--pos-enc learned|sinusoidal|alibi|rope] [--doc]\n"; return 0; } } @@ -131,10 +137,9 @@ int main(int argc, char* argv[]) std::cout << "========================================\n"; std::cout << " CausalSelfAttention 数值梯度检查 (gradcheck)\n"; std::cout << "========================================\n"; - std::cout << " d_model=" << d_model << " heads=" << num_heads - << " seq=" << seq << " batch=" << batch << " tol=" << tol - << " pos_enc=" << pos_enc_name << " doc=" << (use_doc ? 1 : 0) - << "\n"; + std::cout << " d_model=" << d_model << " heads=" << num_heads << " seq=" << seq + << " batch=" << batch << " tol=" << tol << " pos_enc=" << pos_enc_name + << " doc=" << (use_doc ? 1 : 0) << "\n"; CausalSelfAttention attn(d_model, num_heads, seq, seq, pos_enc); NN_EXIT(attn.init(eng), 1, "CausalSelfAttention init 失败: "); @@ -147,8 +152,7 @@ int main(int argc, char* argv[]) std::vector ids(batch * seq, 0); for (std::size_t b = 0; b < batch; ++b) for (std::size_t t = 0; t < seq; ++t) - ids[b * seq + t] = b * 2 + (t < seq / 2 ? std::size_t{0} - : std::size_t{1}); + ids[b * seq + t] = b * 2 + (t < seq / 2 ? std::size_t{0} : std::size_t{1}); attn.set_doc_ids(ids); } @@ -158,8 +162,10 @@ int main(int argc, char* argv[]) const std::size_t total = seq * batch; Matrix x_m(d_model, total); Matrix 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); + 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); NN_EXIT(x, 1, "from_matrix(x) failed\n"); @@ -170,7 +176,7 @@ int main(int argc, char* argv[]) NN_EXIT(y_fwd, 1, "forward failed: "); auto params = attn.parameters(); - auto grads = attn.param_gradients(); + auto grads = attn.param_gradients(); for (auto& g : grads) { auto rz = eng.zero(g.get()); @@ -181,13 +187,15 @@ int main(int argc, char* argv[]) const Scalar eps = 1e-3f; bool all_pass = true; - static constexpr std::array tags{"wq.w","wq.b","wk.w","wk.b","wv.w","wv.b","wo.w","wo.b"}; + static constexpr std::array tags{"wq.w", "wq.b", "wk.w", "wk.b", + "wv.w", "wv.b", "wo.w", "wo.b"}; for (std::size_t p = 0; p < params.size(); ++p) { auto pm = eng.to_matrix(params[p].get()); auto gm = eng.to_matrix(grads[p].get()); std::string name = (p < 8) ? std::string(tags[p]) : ("param[" + std::to_string(p) + "]"); - all_pass &= check_grad_tensor(eng, attn, *x, *go, params[p].get(), *pm, *gm, name, eps, tol); + all_pass &= + check_grad_tensor(eng, attn, *x, *go, params[p].get(), *pm, *gm, name, eps, tol); } // 输入梯度 diff --git a/src/attn_test.cpp b/src/attn_test.cpp index e1d72286..7cebc64f 100644 --- a/src/attn_test.cpp +++ b/src/attn_test.cpp @@ -43,19 +43,21 @@ static int mha_bidirectional_probe() std::uniform_real_distribution dist(-1, 1); nn::Matrix x1(d_model, total); - for (std::size_t i = 0; i < x1.size(); ++i) x1.span()[i] = dist(rng); + for (std::size_t i = 0; i < x1.size(); ++i) + x1.span()[i] = dist(rng); nn::Matrix x2(d_model, total); // 仅扰动 (b0, t=seq-1) 一列 - for (std::size_t i = 0; i < x2.size(); ++i) x2.span()[i] = x1.span()[i]; + for (std::size_t i = 0; i < x2.size(); ++i) + x2.span()[i] = x1.span()[i]; for (std::size_t r = 0; r < d_model; ++r) x2.span()[r * total + (seq - 1)] = dist(rng); // 位置 0 输出的扰动敏感度（d_model 行 × 第 0 列） - const auto col0_delta = [](const nn::Matrix& a, const nn::Matrix& b) { + const auto col0_delta = [](const nn::Matrix& a, const nn::Matrix& b) + { nn::Scalar m = 0; for (std::size_t r = 0; r < a.rows(); ++r) { - const nn::Scalar d = std::fabs( - a.span()[r * a.cols()] - b.span()[r * b.cols()]); + const nn::Scalar d = std::fabs(a.span()[r * a.cols()] - b.span()[r * b.cols()]); if (d > m) m = d; } return m; @@ -66,14 +68,26 @@ static int mha_bidirectional_probe() NN_EXIT(attn_c.init(cpu), 1, "MHA cpu init: "); auto t1 = cpu.from_matrix(x1); auto t2 = cpu.from_matrix(x2); - if (!t1 || !t2) { std::cerr << "MHA probe: from_matrix failed\n"; return 1; } + if (!t1 || !t2) + { + std::cerr << "MHA probe: from_matrix failed\n"; + return 1; + } auto y1 = attn_c.forward(*t1); auto y2 = attn_c.forward(*t2); - if (!y1 || !y2) { std::cerr << "MHA probe forward: " - << (!y1 ? y1.error().message : y2.error().message) << "\n"; return 1; } + if (!y1 || !y2) + { + std::cerr << "MHA probe forward: " << (!y1 ? y1.error().message : y2.error().message) + << "\n"; + return 1; + } auto m1 = cpu.to_matrix(*y1); auto m2 = cpu.to_matrix(*y2); - if (!m1 || !m2) { std::cerr << "MHA probe: to_matrix failed\n"; return 1; } + if (!m1 || !m2) + { + std::cerr << "MHA probe: to_matrix failed\n"; + return 1; + } const nn::Scalar d_mha = col0_delta(*m1, *m2); const bool mha_ok = d_mha > 1e-5f; std::cout << " MHA 双向探针(改未来列→位置0须变): delta=" << d_mha @@ -84,10 +98,18 @@ static int mha_bidirectional_probe() NN_EXIT(attn_cs_c.init(cpu), 1, "CSA cpu init: "); auto cy1 = attn_cs_c.forward(*t1); auto cy2 = attn_cs_c.forward(*t2); - if (!cy1 || !cy2) { std::cerr << "CSA probe forward failed\n"; return 1; } + if (!cy1 || !cy2) + { + std::cerr << "CSA probe forward failed\n"; + return 1; + } auto cm1 = cpu.to_matrix(*cy1); auto cm2 = cpu.to_matrix(*cy2); - if (!cm1 || !cm2) { std::cerr << "CSA probe: to_matrix failed\n"; return 1; } + if (!cm1 || !cm2) + { + std::cerr << "CSA probe: to_matrix failed\n"; + return 1; + } const nn::Scalar d_csa = col0_delta(*cm1, *cm2); const bool csa_ok = d_csa < 1e-5f; std::cout << " CSA 因果阴性对照(位置0须不变): delta=" << d_csa @@ -119,7 +141,11 @@ static int mha_bidirectional_probe() } auto g1 = gpu.from_matrix(x1); auto g2 = gpu.from_matrix(x2); - if (!g1 || !g2) { std::cerr << "MHA probe gpu from_matrix failed\n"; return 1; } + if (!g1 || !g2) + { + std::cerr << "MHA probe gpu from_matrix failed\n"; + return 1; + } auto gy1 = attn_g.forward(*g1); // 覆盖 GPU Plain fold key auto gy2 = attn_g.forward(*g2); // GPU 双向探针 if (!gy1 || !gy2) @@ -130,7 +156,11 @@ static int mha_bidirectional_probe() } auto gm1 = gpu.to_matrix(*gy1); auto gm2 = gpu.to_matrix(*gy2); - if (!gm1 || !gm2) { std::cerr << "MHA gpu to_matrix failed\n"; return 1; } + if (!gm1 || !gm2) + { + std::cerr << "MHA gpu to_matrix failed\n"; + return 1; + } const nn::Scalar d_gpu = col0_delta(*gm1, *gm2); const bool gpu_ok = d_gpu > 1e-5f; std::cout << " MHA GPU 双向探针: delta=" << d_gpu @@ -144,31 +174,44 @@ static int mha_bidirectional_probe() if (d > max_abs) max_abs = d; } const bool cons_ok = max_abs <= 1e-3f; - std::cout << " MHA CPU/GPU forward 一致性: max_abs=" << max_abs - << (cons_ok ? " ✅" : " ❌") << "\n"; + std::cout << " MHA CPU/GPU forward 一致性: max_abs=" << max_abs << (cons_ok ? " ✅" : " ❌") + << "\n"; if (!cons_ok) ++failures; // backward 一致性：两侧最后 forward 对齐为同一输入后反传同一梯度 // （GPU 侧触发 MHA 裸 S 重算表达式——scan 闭合世界注册的验针） auto b1c = attn_c.forward(*t1); auto b1g = attn_g.forward(*g1); - if (!b1c || !b1g) { std::cerr << "MHA re-forward failed\n"; return 1; } + if (!b1c || !b1g) + { + std::cerr << "MHA re-forward failed\n"; + return 1; + } nn::Matrix gm(d_model, total); - for (std::size_t i = 0; i < gm.size(); ++i) gm.span()[i] = dist(rng); + for (std::size_t i = 0; i < gm.size(); ++i) + gm.span()[i] = dist(rng); auto tg_c = cpu.from_matrix(gm); auto tg_g = gpu.from_matrix(gm); - if (!tg_c || !tg_g) { std::cerr << "grad from_matrix failed\n"; return 1; } + if (!tg_c || !tg_g) + { + std::cerr << "grad from_matrix failed\n"; + return 1; + } auto bc = attn_c.backward(*tg_c); auto bg = attn_g.backward(*tg_g); if (!bc || !bg) { - std::cerr << "MHA backward FAILED: " - << (!bc ? bc.error().message : bg.error().message) << "\n"; + std::cerr << "MHA backward FAILED: " << (!bc ? bc.error().message : bg.error().message) + << "\n"; return 1; } auto bcm = cpu.to_matrix(*bc); auto bgm = gpu.to_matrix(*bg); - if (!bcm || !bgm) { std::cerr << "backward to_matrix failed\n"; return 1; } + if (!bcm || !bgm) + { + std::cerr << "backward to_matrix failed\n"; + return 1; + } nn::Scalar bmax = 0; for (std::size_t i = 0; i < bcm->size(); ++i) { @@ -176,8 +219,7 @@ static int mha_bidirectional_probe() if (d > bmax) bmax = d; } const bool bok = bmax <= 1e-3f; - std::cout << " MHA CPU/GPU backward 一致性: max_abs=" << bmax - << (bok ? " ✅" : " ❌") << "\n"; + std::cout << " MHA CPU/GPU backward 一致性: max_abs=" << bmax << (bok ? " ✅" : " ❌") << "\n"; if (!bok) ++failures; return failures; @@ -187,7 +229,10 @@ int main(int argc, char* argv[]) { int failures = 0; // 77 = 子测试因 GPU 不可用而跳过（ctest SKIP 约定），不计入失败 - const auto add = [&failures](int r) { if (r != 77) failures += r; }; + const auto add = [&failures](int r) + { + if (r != 77) failures += r; + }; std::puts("=== attn_gradcheck (CausalSelfAttention numerical) ==="); add(test_attn_gradcheck(argc, argv)); @@ -195,9 +240,8 @@ int main(int argc, char* argv[]) // CausalAlibiDocScoreMask 分支的梯度（仅 forward 通过不足以验证—— // 梯度跨文档串扰 forward 抓不住）；默认 learned+doc 与 // alibi+doc 各跑一遍数值梯度 - char n0[] = "attn_gradcheck", n1[] = "--doc", n2[] = "--pos-enc", - n3[] = "alibi"; - char* av_doc[] = {n0, n1}; + char n0[] = "attn_gradcheck", n1[] = "--doc", n2[] = "--pos-enc", n3[] = "alibi"; + char* av_doc[] = {n0, n1}; char* av_adoc[] = {n0, n2, n3, n1}; std::puts("--- gradcheck: doc variant ---"); add(test_attn_gradcheck(2, av_doc)); diff --git a/src/attn_w_batch_test.cpp b/src/attn_w_batch_test.cpp index bd4ad80d..cb56ecdf 100644 --- a/src/attn_w_batch_test.cpp +++ b/src/attn_w_batch_test.cpp @@ -21,8 +21,8 @@ #include #include -using nn::Scalar; using nn::Matrix; +using nn::Scalar; int main(int argc, char** argv) { @@ -37,7 +37,12 @@ int main(int argc, char** argv) const std::size_t d_model = 16; std::size_t batch = 0, heads = 0, seq = 5; - if (argc >= 4) { batch = std::stoul(argv[1]); heads = std::stoul(argv[2]); seq = std::stoul(argv[3]); } + if (argc >= 4) + { + batch = std::stoul(argv[1]); + heads = std::stoul(argv[2]); + seq = std::stoul(argv[3]); + } std::mt19937_64 rng(123); std::uniform_real_distribution dist(-1, 1); @@ -48,8 +53,22 @@ int main(int argc, char** argv) << " seq=" << seq << " batch=" << b << " (BH=" << b * h << ")\n"; nn::CausalSelfAttention attn_cpu(d_model, h, seq, seq, nn::PosEncodingType::Learned); nn::CausalSelfAttention attn_gpu(d_model, h, seq, seq, nn::PosEncodingType::Learned); - { auto r = attn_cpu.init(cpu); if (!r) { std::cerr << "cpu init: " << r.error().message << "\n"; return false; } } - { auto r = attn_gpu.init(gpu); if (!r) { std::cerr << "gpu init: " << r.error().message << "\n"; return false; } } + { + auto r = attn_cpu.init(cpu); + if (!r) + { + std::cerr << "cpu init: " << r.error().message << "\n"; + return false; + } + } + { + auto r = attn_gpu.init(gpu); + if (!r) + { + std::cerr << "gpu init: " << r.error().message << "\n"; + return false; + } + } // 权重同步：CPU 随机 → GPU { @@ -58,28 +77,53 @@ int main(int argc, char** argv) for (std::size_t i = 0; i < pc.size(); ++i) { auto cm = cpu.to_matrix(pc[i].get()); - if (!cm) { std::cerr << "to_matrix cpu failed\n"; return false; } + if (!cm) + { + std::cerr << "to_matrix cpu failed\n"; + return false; + } auto cp = gpu.write(pg[i].get(), cm->span()); - if (!cp) { std::cerr << "write gpu failed: " << cp.error().message << "\n"; return false; } + if (!cp) + { + std::cerr << "write gpu failed: " << cp.error().message << "\n"; + return false; + } } } const std::size_t total = seq * b; Matrix x_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 < x_m.size(); ++i) + x_m.span()[i] = dist(rng); auto xc = cpu.from_matrix(x_m); auto xg = gpu.from_matrix(x_m); - if (!xc || !xg) { std::cerr << "from_matrix failed\n"; return false; } + if (!xc || !xg) + { + std::cerr << "from_matrix failed\n"; + return false; + } auto yc = attn_cpu.forward(*xc); auto yg = attn_gpu.forward(*xg); - if (!yc) { std::cout << "CPU forward 失败: " << yc.error().message << "\n"; return false; } - if (!yg) { std::cout << "GPU forward 失败: " << yg.error().message << "\n"; return false; } + if (!yc) + { + std::cout << "CPU forward 失败: " << yc.error().message << "\n"; + return false; + } + if (!yg) + { + std::cout << "GPU forward 失败: " << yg.error().message << "\n"; + return false; + } auto ycm = cpu.to_matrix(*yc); auto ygm = gpu.to_matrix(*yg); - if (!ycm || !ygm) { std::cerr << "to_matrix failed\n"; return false; } + if (!ycm || !ygm) + { + std::cerr << "to_matrix failed\n"; + return false; + } Scalar max_abs = 0; std::size_t bad = 0; @@ -91,9 +135,8 @@ int main(int argc, char** argv) if (d > max_abs) max_abs = d; if (d > 1e-3f) ++bad; } - std::cout << " max_abs_diff=" << max_abs - << " bad_cells(>1e-3)=" << bad << " / " << ycm->size() - << " " << (max_abs <= 1e-3f ? "✅ 一致" : "❌ 不一致") << "\n"; + std::cout << " max_abs_diff=" << max_abs << " bad_cells(>1e-3)=" << bad << " / " + << ycm->size() << " " << (max_abs <= 1e-3f ? "✅ 一致" : "❌ 不一致") << "\n"; return max_abs <= 1e-3f; }; diff --git a/src/batch_source_test.cpp b/src/batch_source_test.cpp index 0273596d..c762ec95 100644 --- a/src/batch_source_test.cpp +++ b/src/batch_source_test.cpp @@ -21,12 +21,14 @@ #include static int g_fail = 0; -#define CHECK(cond, msg) \ - do { \ - if (!(cond)) { \ - std::cout << " [FAIL] " << (msg) << " (" << #cond << ")\n"; \ - ++g_fail; \ - } \ +#define CHECK(cond, msg) \ + do \ + { \ + if (!(cond)) \ + { \ + std::cout << " [FAIL] " << (msg) << " (" << #cond << ")\n"; \ + ++g_fail; \ + } \ } while (0) namespace @@ -42,7 +44,8 @@ void test_step_grid() { // 载荷 = 下标和（便于断言）；单位 = 下标个数 std::size_t sum = 0; - for (std::size_t i : ix) sum += i; + for (std::size_t i : ix) + sum += i; return std::make_pair(std::any(sum), ix.size()); }); CHECK(src_r.has_value(), "create 合法参数成功"); @@ -68,8 +71,8 @@ void test_step_grid() void test_shuffle_determinism() { std::cout << "\n── [2] 每 epoch shuffle：同 seed/epoch 复现，跨 epoch 变化 ──\n"; - auto assemble = [](std::span ix) - -> nn::Result> + auto assemble = + [](std::span ix) -> nn::Result> { return std::make_pair(std::any(), ix.size()); }; auto a_r = nn::IndexedBatchSource::create(32, 8, 42, assemble); @@ -109,8 +112,7 @@ void test_rewind_validation() std::cout << "\n── [3] rewind：start_step 越界报错 ──\n"; auto src_r = nn::IndexedBatchSource::create( 10, 4, 1, - [](std::span ix) - -> nn::Result> + [](std::span ix) -> nn::Result> { return std::make_pair(std::any(), ix.size()); }); CHECK(src_r.has_value(), "create 成功"); if (!src_r) return; @@ -123,16 +125,13 @@ void test_rewind_validation() void test_ctor_validation() { std::cout << "\n── [4] 构造期校验：batch_size=0 / sample=0 / 空回调 ──\n"; - auto assemble = [](std::span ix) - -> nn::Result> + auto assemble = + [](std::span ix) -> nn::Result> { return std::make_pair(std::any(), ix.size()); }; - CHECK(!nn::IndexedBatchSource::create(10, 0, 1, assemble).has_value(), - "batch_size=0 → 报错"); - CHECK(!nn::IndexedBatchSource::create(0, 4, 1, assemble).has_value(), - "sample_count=0 → 报错"); - CHECK(!nn::IndexedBatchSource::create(10, 4, 1, nullptr).has_value(), - "空装配回调 → 报错"); + CHECK(!nn::IndexedBatchSource::create(10, 0, 1, assemble).has_value(), "batch_size=0 → 报错"); + CHECK(!nn::IndexedBatchSource::create(0, 4, 1, assemble).has_value(), "sample_count=0 → 报错"); + CHECK(!nn::IndexedBatchSource::create(10, 4, 1, nullptr).has_value(), "空装配回调 → 报错"); auto ok = nn::IndexedBatchSource::create(5, 2, 0, assemble); CHECK(ok && ok->steps_per_epoch() == 3, "5/2 → 3 步"); @@ -144,8 +143,7 @@ void test_payload() std::cout << "\n── [5] 载荷：any 指针形式（-fno-exceptions 安全）──\n"; auto src_r = nn::IndexedBatchSource::create( 6, 3, 9, - [](std::span ix) - -> nn::Result> + [](std::span ix) -> nn::Result> { std::string tag = "batch:" + std::to_string(ix.size()); return std::make_pair(std::any(std::move(tag)), ix.size() * 10); @@ -158,8 +156,7 @@ void test_payload() { const auto* tag = std::any_cast(&b->payload); CHECK(tag != nullptr && *tag == "batch:3", "指针形式读出域载荷"); - CHECK(std::any_cast(&b->payload) == nullptr, - "类型不符 → nullptr（不抛）"); + CHECK(std::any_cast(&b->payload) == nullptr, "类型不符 → nullptr（不抛）"); CHECK(b->units == 30, "域回调给出归一化单位"); } } @@ -187,7 +184,8 @@ void test_pull_source() CHECK(!src.steps_per_epoch().has_value(), "无固定 epoch → nullopt"); auto b = src.next(0); CHECK(b && !b->epoch_done && b->samples == 4, "步进驱动：无 epoch 完成语义"); - CHECK(b && std::any_cast(&b->payload) && *std::any_cast(&b->payload) == 0, + CHECK(b && std::any_cast(&b->payload) && + *std::any_cast(&b->payload) == 0, "载荷可透传任意域数据"); CHECK(src.rewind(0, 0).has_value(), "默认 rewind = 空实现（合法）"); } diff --git a/src/batchnorm_test.cpp b/src/batchnorm_test.cpp index 40b5887e..e98c880a 100644 --- a/src/batchnorm_test.cpp +++ b/src/batchnorm_test.cpp @@ -26,23 +26,23 @@ #include "test_common.hpp" -using nn::Scalar; +using nn::ComputeEngine; using nn::Matrix; +using nn::Scalar; using nn::Tensor; using nn::TensorRef; -using nn::ComputeEngine; int g_fail = 0; // CHECK 计数器（test_common.hpp 约定） -namespace { +namespace +{ constexpr Scalar EPS = 1e-5f; constexpr Scalar MOMENTUM = 0.1f; // 与 BatchNorm 默认一致（新 batch 权重） // ── 宿主参考：单次 forward（训练态用 batch 统计，推理态用 running 统计）── -Matrix ref_forward(bool training, const Matrix &x, const Matrix &gamma, - const Matrix &beta, const Matrix &run_mean, - const Matrix &run_var) +Matrix ref_forward(bool training, const Matrix& x, const Matrix& gamma, const Matrix& beta, + const Matrix& run_mean, const Matrix& run_var) { const std::size_t F = x.rows(), B = x.cols(); Matrix out(F, B); @@ -51,7 +51,8 @@ Matrix ref_forward(bool training, const Matrix &x, const Matrix &gamma, Scalar mean = 0, var = 0; if (training) { - for (std::size_t b = 0; b < B; ++b) mean += x.at_unchecked(f, b); + for (std::size_t b = 0; b < B; ++b) + mean += x.at_unchecked(f, b); mean /= static_cast(B); for (std::size_t b = 0; b < B; ++b) { @@ -63,36 +64,38 @@ Matrix ref_forward(bool training, const Matrix &x, const Matrix &gamma, else { mean = run_mean.at_unchecked(f, 0); - var = run_var.at_unchecked(f, 0); + var = run_var.at_unchecked(f, 0); } const Scalar inv = 1.0f / std::sqrt(var + EPS); for (std::size_t b = 0; b < B; ++b) out.set_value_unchecked(f, b, - gamma.at_unchecked(f, 0) * (x.at_unchecked(f, b) - mean) * inv - + beta.at_unchecked(f, 0)); + gamma.at_unchecked(f, 0) * (x.at_unchecked(f, b) - mean) * inv + + beta.at_unchecked(f, 0)); } return out; } // 随机 (F,1) 列向量（宿主） -Matrix rand_col(std::size_t F, std::mt19937_64 &rng, Scalar lo, Scalar hi) +Matrix rand_col(std::size_t F, std::mt19937_64& rng, Scalar lo, Scalar hi) { Matrix m(F, 1); std::uniform_real_distribution dist(lo, hi); - for (std::size_t i = 0; i < F; ++i) m.span()[i] = dist(rng); + for (std::size_t i = 0; i < F; ++i) + m.span()[i] = dist(rng); return m; } // 随机 (F,B) -Matrix rand_mat(std::size_t F, std::size_t B, std::mt19937_64 &rng) +Matrix rand_mat(std::size_t F, std::size_t B, std::mt19937_64& rng) { Matrix m(F, B); std::uniform_real_distribution dist(-1, 1); - for (std::size_t i = 0; i < m.size(); ++i) m.span()[i] = dist(rng); + for (std::size_t i = 0; i < m.size(); ++i) + m.span()[i] = dist(rng); return m; } -bool matrices_equal(const Matrix &a, const Matrix &b) +bool matrices_equal(const Matrix& a, const Matrix& b) { if (a.rows() != b.rows() || a.cols() != b.cols()) return false; const auto sa = a.span(), sb = b.span(); @@ -102,8 +105,7 @@ bool matrices_equal(const Matrix &a, const Matrix &b) } // ── gradcheck 辅助（rmsnorm_gradcheck 同款：中心差分）──────────────────── -Scalar eval_loss(ComputeEngine &engine, nn::Layer &norm, - const Tensor &x, const Tensor &go) +Scalar eval_loss(ComputeEngine& engine, nn::Layer& norm, const Tensor& x, const Tensor& go) { auto y = norm.forward(x); NN_ASSERT_ALWAYS(y, "eval_loss: forward failed"); @@ -113,11 +115,10 @@ Scalar eval_loss(ComputeEngine &engine, nn::Layer &norm, return dot(*y_m, *go_m); } -bool check_grad_tensor( - ComputeEngine &engine, nn::Layer &norm, - const Tensor &fwd_input, const Tensor &go, - Tensor ¶m, const Matrix &base, const Matrix &grad_analytical, - const std::string &name, Scalar eps, Scalar tol) +bool check_grad_tensor(ComputeEngine& engine, nn::Layer& norm, const Tensor& fwd_input, + const Tensor& go, Tensor& param, const Matrix& base, + const Matrix& grad_analytical, const std::string& name, Scalar eps, + Scalar tol) { bool ok = true; std::size_t bad = 0; @@ -131,17 +132,29 @@ bool check_grad_tensor( Matrix pp = base; pp.set_value_unchecked(r, c, orig + eps); auto r_p = engine.write(param, pp.span()); - if (!r_p) { std::cerr << "write(+) failed\n"; return false; } + if (!r_p) + { + std::cerr << "write(+) failed\n"; + return false; + } const Scalar lp = eval_loss(engine, norm, fwd_input, go); Matrix pm = base; pm.set_value_unchecked(r, c, orig - eps); auto r_m = engine.write(param, pm.span()); - if (!r_m) { std::cerr << "write(-) failed\n"; return false; } + if (!r_m) + { + std::cerr << "write(-) failed\n"; + return false; + } const Scalar lm = eval_loss(engine, norm, fwd_input, go); auto r_r = engine.write(param, base.span()); - if (!r_r) { std::cerr << "write(restore) failed\n"; return false; } + if (!r_r) + { + std::cerr << "write(restore) failed\n"; + return false; + } const Scalar num = (lp - lm) / (Scalar{2} * eps); const Scalar ana = grad_analytical.at_unchecked(r, c); @@ -151,14 +164,14 @@ bool check_grad_tensor( { ok = false; if (bad < 10) - std::cout << " [FAIL] " << name << "(" << r << "," << c - << ") num=" << num << " ana=" << ana << "\n"; + std::cout << " [FAIL] " << name << "(" << r << "," << c << ") num=" << num + << " ana=" << ana << "\n"; ++bad; } } } - std::cout << " " << name << ": " << (ok ? "OK" : "FAIL") - << " max_err=" << max_err << " bad=" << bad << "\n"; + std::cout << " " << name << ": " << (ok ? "OK" : "FAIL") << " max_err=" << max_err + << " bad=" << bad << "\n"; return ok; } @@ -168,30 +181,33 @@ void test_factory() std::cout << "\n── 1. 工厂分支（禁止静默回落）──\n"; { auto l = nn::make_norm_layer(8, nn::NormType::BatchNorm); - CHECK(l && dynamic_cast(l.get()) != nullptr, + CHECK(l && dynamic_cast(l.get()) != nullptr, "make_norm_layer(BatchNorm) 应返回真 BatchNorm"); } { auto l = nn::make_norm_layer(8, nn::NormType::LayerNorm); - CHECK(l && dynamic_cast(l.get()) != nullptr, + CHECK(l && dynamic_cast(l.get()) != nullptr, "make_norm_layer(LayerNorm) 应返回 LayerNorm"); } { auto l = nn::make_norm_layer(8, nn::NormType::RMSNorm); - CHECK(l && dynamic_cast(l.get()) != nullptr, + CHECK(l && dynamic_cast(l.get()) != nullptr, "make_norm_layer(RMSNorm) 应返回 RMSNorm"); } } // ── 2/3. 训练态 & 推理态 forward 对拍 ──────────────────────────────────── -void test_forward(ComputeEngine &eng) +void test_forward(ComputeEngine& eng) { std::cout << "\n── 2/3. 训练态 / 推理态 forward vs 宿主参考 ──\n"; const std::size_t F = 8, B = 5; std::mt19937_64 rng(42); nn::BatchNorm bn(F); - { auto r = bn.init(eng); CHECK(r, "BatchNorm init"); } + { + auto r = bn.init(eng); + CHECK(r, "BatchNorm init"); + } // 随机 γ/β（覆盖参数广播路径） const Matrix g_m = rand_col(F, rng, 0.5f, 1.5f); @@ -211,7 +227,8 @@ void test_forward(ComputeEngine &eng) if (y) { Matrix zeros(F, 1), ones(F, 1); - for (std::size_t i = 0; i < F; ++i) ones.span()[i] = 1; + for (std::size_t i = 0; i < F; ++i) + ones.span()[i] = 1; const Matrix ref = ref_forward(true, x_m, g_m, b_m, zeros, ones); CHECK(close_to(download(eng, *y), ref, 2e-4f, "训练态 forward", 0), "训练态 forward 与宿主参考一致"); @@ -230,8 +247,7 @@ void test_forward(ComputeEngine &eng) const Matrix rm = download(eng, extras[0].get()); const Matrix rv = download(eng, extras[1].get()); const Matrix ref = ref_forward(false, x_m, g_m, b_m, rm, rv); - CHECK(close_to(download(eng, *y_eval), ref, 2e-4f, - "推理态 forward(running 更新后)", 1), + CHECK(close_to(download(eng, *y_eval), ref, 2e-4f, "推理态 forward(running 更新后)", 1), "推理态 forward 与宿主参考一致"); } } @@ -251,8 +267,7 @@ void test_forward(ComputeEngine &eng) const Matrix rm = download(eng, extras[0].get()); const Matrix rv = download(eng, extras[1].get()); const Matrix ref = ref_forward(false, x1, g_m, b_m, rm, rv); - CHECK(close_to(download(eng, *ye), ref, 2e-4f, - "推理态 forward batch=1", 2), + CHECK(close_to(download(eng, *ye), ref, 2e-4f, "推理态 forward batch=1", 2), "推理态 batch=1 与宿主参考一致"); } } @@ -260,19 +275,24 @@ void test_forward(ComputeEngine &eng) } // ── 4. running 统计 EMA ────────────────────────────────────────────────── -void test_running_stats(ComputeEngine &eng) +void test_running_stats(ComputeEngine& eng) { std::cout << "\n── 4. running 统计 EMA（多批）──\n"; const std::size_t F = 6, B = 7; std::mt19937_64 rng(7); nn::BatchNorm bn(F); - { auto r = bn.init(eng); CHECK(r, "BatchNorm init"); } + { + auto r = bn.init(eng); + CHECK(r, "BatchNorm init"); + } Matrix run_mean(F, 1), run_var(F, 1); // 宿主 EMA：初值 mean=0 / var=1 - for (std::size_t i = 0; i < F; ++i) run_var.span()[i] = 1; + for (std::size_t i = 0; i < F; ++i) + run_var.span()[i] = 1; Matrix zeros(F, 1), ones(F, 1); - for (std::size_t i = 0; i < F; ++i) ones.span()[i] = 1; + for (std::size_t i = 0; i < F; ++i) + ones.span()[i] = 1; bn.set_training(true); for (int step = 0; step < 4; ++step) @@ -285,7 +305,8 @@ void test_running_stats(ComputeEngine &eng) for (std::size_t f = 0; f < F; ++f) { Scalar mean = 0, var = 0; - for (std::size_t b = 0; b < B; ++b) mean += x_m.at_unchecked(f, b); + for (std::size_t b = 0; b < B; ++b) + mean += x_m.at_unchecked(f, b); mean /= static_cast(B); for (std::size_t b = 0; b < B; ++b) { @@ -293,10 +314,10 @@ void test_running_stats(ComputeEngine &eng) var += d * d; } var /= static_cast(B); - run_mean.set_value_unchecked(f, 0, - (1 - MOMENTUM) * run_mean.at_unchecked(f, 0) + MOMENTUM * mean); - run_var.set_value_unchecked(f, 0, - (1 - MOMENTUM) * run_var.at_unchecked(f, 0) + MOMENTUM * var); + run_mean.set_value_unchecked( + f, 0, (1 - MOMENTUM) * run_mean.at_unchecked(f, 0) + MOMENTUM * mean); + run_var.set_value_unchecked( + f, 0, (1 - MOMENTUM) * run_var.at_unchecked(f, 0) + MOMENTUM * var); } } @@ -304,11 +325,9 @@ void test_running_stats(ComputeEngine &eng) CHECK(extras.size() == 2, "extra_state 应为 2 个张量"); if (extras.size() == 2) { - CHECK(close_to(download(eng, extras[0].get()), run_mean, 2e-4f, - "running_mean EMA", 3), + CHECK(close_to(download(eng, extras[0].get()), run_mean, 2e-4f, "running_mean EMA", 3), "running_mean 与宿主 EMA 一致"); - CHECK(close_to(download(eng, extras[1].get()), run_var, 2e-4f, - "running_var EMA", 4), + CHECK(close_to(download(eng, extras[1].get()), run_var, 2e-4f, "running_var EMA", 4), "running_var 与宿主 EMA 一致"); } @@ -329,31 +348,34 @@ void test_running_stats(ComputeEngine &eng) } // ── 5. gradcheck（训练态：γ / β / 输入）────────────────────────────────── -void test_gradcheck(ComputeEngine &eng, Scalar tol) +void test_gradcheck(ComputeEngine& eng, Scalar tol) { std::cout << "\n── 5. gradcheck（训练态中心差分）──\n"; const std::size_t F = 8, B = 6; std::mt19937_64 rng(123); nn::BatchNorm bn(F); - { auto r = bn.init(eng); CHECK(r, "BatchNorm init"); } + { + auto r = bn.init(eng); + CHECK(r, "BatchNorm init"); + } bn.set_training(true); const Matrix x_m = rand_mat(F, B, rng); const Matrix go_m = rand_mat(F, B, rng); - auto x = upload(eng, x_m); + auto x = upload(eng, x_m); auto go = upload(eng, go_m); // 先 forward 一次填缓存，再 backward（铁律 #8 前置：缓存非空） auto y_fwd = bn.forward(x); CHECK(y_fwd, "gradcheck forward"); - for (auto &g : bn.param_gradients()) + for (auto& g : bn.param_gradients()) CHECK(eng.zero(g.get()), "zero_grad"); auto gx = bn.backward(go); CHECK(gx, "gradcheck backward"); auto params = bn.parameters(); // [gamma, beta] - auto grads = bn.param_gradients(); // [grad_gamma, grad_beta] + auto grads = bn.param_gradients(); // [grad_gamma, grad_beta] bool all_pass = true; const Scalar eps = 1e-3f; for (std::size_t pi = 0; pi < params.size(); ++pi) @@ -363,8 +385,7 @@ void test_gradcheck(ComputeEngine &eng, Scalar tol) CHECK(p && g, "to_matrix(param/grad)"); if (!p || !g) return; const std::string name = (pi == 0) ? "gamma" : "beta"; - all_pass &= check_grad_tensor(eng, bn, x, go, params[pi].get(), - *p, *g, name, eps, tol); + all_pass &= check_grad_tensor(eng, bn, x, go, params[pi].get(), *p, *g, name, eps, tol); } // 输入梯度（扰动输入副本，前向用副本） { @@ -374,21 +395,23 @@ void test_gradcheck(ComputeEngine &eng, Scalar tol) auto xt = eng.clone(x); CHECK(xt, "clone(x)"); if (!xt) return; - all_pass &= check_grad_tensor(eng, bn, *xt, go, *xt, x_m, *gxm, - "grad_x", eps, tol); + all_pass &= check_grad_tensor(eng, bn, *xt, go, *xt, x_m, *gxm, "grad_x", eps, tol); } CHECK(all_pass, "训练态 gradcheck 全部通过"); } // ── 6. 推理态 backward vs 宿主参考 ─────────────────────────────────────── -void test_eval_backward(ComputeEngine &eng) +void test_eval_backward(ComputeEngine& eng) { std::cout << "\n── 6. 推理态 backward vs 宿主参考 ──\n"; const std::size_t F = 8, B = 5; std::mt19937_64 rng(99); nn::BatchNorm bn(F); - { auto r = bn.init(eng); CHECK(r, "BatchNorm init"); } + { + auto r = bn.init(eng); + CHECK(r, "BatchNorm init"); + } // 非平凡 running 统计：先跑两次训练态 forward（EMA 从 0/1 起步） bn.set_training(true); @@ -404,11 +427,11 @@ void test_eval_backward(ComputeEngine &eng) bn.set_training(false); const Matrix x_m = rand_mat(F, B, rng); const Matrix go_m = rand_mat(F, B, rng); - auto x = upload(eng, x_m); + auto x = upload(eng, x_m); auto go = upload(eng, go_m); auto y = bn.forward(x); CHECK(y, "推理态 forward"); - for (auto &g : bn.param_gradients()) + for (auto& g : bn.param_gradients()) CHECK(eng.zero(g.get()), "zero_grad"); auto gx = bn.backward(go); CHECK(gx, "推理态 backward"); @@ -433,7 +456,7 @@ void test_eval_backward(ComputeEngine &eng) const Scalar inv = 1.0f / std::sqrt(rv.at_unchecked(f, 0) + EPS); for (std::size_t b = 0; b < B; ++b) ref_gx.set_value_unchecked(f, b, - go_m.at_unchecked(f, b) * g_m.at_unchecked(f, 0) * inv); + go_m.at_unchecked(f, b) * g_m.at_unchecked(f, 0) * inv); } CHECK(close_to(download(eng, *gx), ref_gx, 2e-4f, "推理态 grad_x", 6), "推理态 grad_x 与宿主参考一致"); @@ -449,21 +472,18 @@ void test_eval_backward(ComputeEngine &eng) { const Scalar n = (x_m.at_unchecked(f, b) - rm.at_unchecked(f, 0)) * inv; ref_gg.set_value_unchecked(f, 0, - ref_gg.at_unchecked(f, 0) + go_m.at_unchecked(f, b) * n); - ref_gb.set_value_unchecked(f, 0, - ref_gb.at_unchecked(f, 0) + go_m.at_unchecked(f, b)); + ref_gg.at_unchecked(f, 0) + go_m.at_unchecked(f, b) * n); + ref_gb.set_value_unchecked(f, 0, ref_gb.at_unchecked(f, 0) + go_m.at_unchecked(f, b)); } } - CHECK(close_to(download(eng, grads[0].get()), ref_gg, 2e-4f, - "推理态 grad_gamma", 7), + CHECK(close_to(download(eng, grads[0].get()), ref_gg, 2e-4f, "推理态 grad_gamma", 7), "推理态 grad_gamma 与宿主参考一致"); - CHECK(close_to(download(eng, grads[1].get()), ref_gb, 2e-4f, - "推理态 grad_beta", 8), + CHECK(close_to(download(eng, grads[1].get()), ref_gb, 2e-4f, "推理态 grad_beta", 8), "推理态 grad_beta 与宿主参考一致"); } // ── 7. Model 级：extra_state / set_training 转发 / save-load 往返 ──────── -void test_model_level(ComputeEngine &eng) +void test_model_level(ComputeEngine& eng) { std::cout << "\n── 7. Model 级（extra_state / set_training / 序列化）──\n"; const std::vector dims = {12, 8, 4}; @@ -472,17 +492,16 @@ void test_model_level(ComputeEngine &eng) auto model_r = nn::build_mnist_mlp_model(eng, dims, nn::NormType::BatchNorm); CHECK(model_r, "build_mnist_mlp_model(BatchNorm)"); if (!model_r) return; - nn::Model &model = *model_r; + nn::Model& model = *model_r; // extra_state 非空 = 工厂真的造了 BatchNorm（LayerNorm/RMSNorm 为空） CHECK(model.extra_state().size() == 2 * n_norms, "Model::extra_state = 每个 BatchNorm 2 个 running 张量"); // set_training 转发：定位 BatchNorm 层（层序 = Linear, BN, GeLU, Linear…） - nn::BatchNorm *bn_layer = nullptr; + nn::BatchNorm* bn_layer = nullptr; for (std::size_t i = 0; i < model.num_layers(); ++i) - if (auto *p = dynamic_cast(&model.layer_at(i))) - bn_layer = p; + if (auto* p = dynamic_cast(&model.layer_at(i))) bn_layer = p; CHECK(bn_layer != nullptr, "模型中应含 BatchNorm 层"); if (!bn_layer) return; @@ -519,7 +538,8 @@ void test_model_level(ComputeEngine &eng) for (std::size_t f = 0; f < F && ok; ++f) { Scalar mean = 0; - for (std::size_t b = 0; b < B; ++b) mean += got.at_unchecked(f, b); + for (std::size_t b = 0; b < B; ++b) + mean += got.at_unchecked(f, b); mean /= static_cast(B); if (std::fabs(mean) > 2e-4f) ok = false; } @@ -528,7 +548,7 @@ void test_model_level(ComputeEngine &eng) // ── save/load 往返：参数 + running 统计逐位一致 ── nn::ModelSpec spec; - auto &sp = spec.emplace(); + auto& sp = spec.emplace(); sp.layer_dims = dims; sp.norm_type = nn::NormType::BatchNorm; @@ -556,19 +576,16 @@ void test_model_level(ComputeEngine &eng) CHECK(a.size() == b.size(), "参数个数一致"); bool same = a.size() == b.size(); for (std::size_t i = 0; i < a.size() && same; ++i) - same = matrices_equal(download(eng, a[i].get()), - download(eng, b[i].get())); + same = matrices_equal(download(eng, a[i].get()), download(eng, b[i].get())); CHECK(same, "参数 save/load 逐位一致"); } // running 统计逐位一致（extra_state 往返） { auto a = model.extra_state(), b = model2_r->extra_state(); - CHECK(a.size() == b.size() && a.size() == 2 * n_norms, - "extra_state 个数一致"); + CHECK(a.size() == b.size() && a.size() == 2 * n_norms, "extra_state 个数一致"); bool same = a.size() == b.size(); for (std::size_t i = 0; i < a.size() && same; ++i) - same = matrices_equal(download(eng, a[i].get()), - download(eng, b[i].get())); + same = matrices_equal(download(eng, a[i].get()), download(eng, b[i].get())); CHECK(same, "running 统计 save/load 逐位一致"); } std::remove(file.c_str()); @@ -576,7 +593,7 @@ void test_model_level(ComputeEngine &eng) } // namespace -int main(int argc, char *argv[]) +int main(int argc, char* argv[]) { // --gpu 预检不可用 → 77（ctest SKIP 约定，与 cnn_test 同款） bool use_gpu = false; @@ -585,8 +602,7 @@ int main(int argc, char *argv[]) { std::string a = argv[i]; if (a == "--gpu") use_gpu = true; - else if (a == "--tol" && i + 1 < argc) - tol = static_cast(std::atof(argv[++i])); + else if (a == "--tol" && i + 1 < argc) tol = static_cast(std::atof(argv[++i])); else if (a == "--help") { std::cout << "用法: batchnorm_test [--gpu] [--tol ]\n"; @@ -606,7 +622,7 @@ int main(int argc, char *argv[]) auto engine_res = nn::cli::create_engine(ecfg, std::cout); NN_EXIT(engine_res, 1, "引擎创建失败: "); auto engine = std::move(*engine_res); - ComputeEngine &eng = *engine; + ComputeEngine& eng = *engine; std::cout << "========================================\n"; std::cout << " BatchNorm 测试 " << (use_gpu ? "(GPU)" : "(CPU)") << "\n"; diff --git a/src/bitnet_struct_test.cpp b/src/bitnet_struct_test.cpp index 1a2ac7e9..2858fdca 100644 --- a/src/bitnet_struct_test.cpp +++ b/src/bitnet_struct_test.cpp @@ -42,14 +42,16 @@ int g_failures = 0; std::size_t count_params(const std::vector& ps) { std::size_t n = 0; - for (const auto& p : ps) n += p.get().rows() * p.get().cols(); + for (const auto& p : ps) + n += p.get().rows() * p.get().cols(); return n; } double max_abs(const nn::Matrix& m) { double r = 0.0; - for (auto v : m.span()) r = std::max(r, std::fabs(static_cast(v))); + for (auto v : m.span()) + r = std::max(r, std::fabs(static_cast(v))); return r; } @@ -75,26 +77,26 @@ void test_subln_param_count() const auto P = nn::PosEncodingType::Learned; const auto A = nn::ActivationType::SwiGLU; - auto build = [&](nn::NormType nt, bool subln) { - auto m = std::make_unique( - V, D, S, H, FF, L, P, A, nt, nn::PrecisionProfile{}, subln); + auto build = [&](nn::NormType nt, bool subln) + { + auto m = std::make_unique(V, D, S, H, FF, L, P, A, nt, nn::PrecisionProfile{}, + subln); NN_EXIT(m->init(eng), 1, "GPTModel init 失败: "); return m; }; auto off = build(nn::NormType::RMSNorm, false); - auto on = build(nn::NormType::RMSNorm, true); + auto on = build(nn::NormType::RMSNorm, true); const std::size_t n_off = count_params(off->parameters()); - const std::size_t n_on = count_params(on->parameters()); - std::printf(" RMSNorm: off=%zu on=%zu Δ=%zu（期望 %zu）\n", - n_off, n_on, n_on - n_off, L * (D + FF)); + const std::size_t n_on = count_params(on->parameters()); + std::printf(" RMSNorm: off=%zu on=%zu Δ=%zu（期望 %zu）\n", n_off, n_on, n_on - n_off, + L * (D + FF)); CHECK(n_on == n_off + L * (D + FF), "RMSNorm 下 SubLN 增量 = layers×(d_model+d_ff)"); // LayerNorm（gamma + beta 两个张量）→ 增量翻倍：证明类型跟随 norm_type auto off_ln = build(nn::NormType::LayerNorm, false); - auto on_ln = build(nn::NormType::LayerNorm, true); - const std::size_t d_ln = count_params(on_ln->parameters()) - - count_params(off_ln->parameters()); + auto on_ln = build(nn::NormType::LayerNorm, true); + const std::size_t d_ln = count_params(on_ln->parameters()) - count_params(off_ln->parameters()); std::printf(" LayerNorm: Δ=%zu（期望 %zu）\n", d_ln, L * 2 * (D + FF)); CHECK(d_ln == L * 2 * (D + FF), "LayerNorm 下 SubLN 增量 = layers×2×(d_model+d_ff)"); @@ -115,8 +117,7 @@ void test_subln_grad_wiring() // FeedForward：参数序 = [fc1.w, fc1.b, fc2.w, fc2.b, sub_norm.gamma] { - nn::FeedForward ff(D, FF, nn::ActivationType::GeLU, /*subln=*/true, - nn::NormType::RMSNorm); + nn::FeedForward ff(D, FF, nn::ActivationType::GeLU, /*subln=*/true, nn::NormType::RMSNorm); NN_EXIT(ff.init(eng), 1, "FeedForward init 失败: "); auto ps = ff.parameters(); CHECK(ps.size() == 5, "FF 参数条数 = 5（含 sub_norm.gamma）"); @@ -139,9 +140,8 @@ void test_subln_grad_wiring() // CausalSelfAttention：参数序 = [w_q, w_k, w_v, w_o, attn_sub_norm.gamma] { - nn::CausalSelfAttention at(D, H, /*max_len=*/8, /*seq_len=*/S, - nn::PosEncodingType::RoPE, /*subln=*/true, - nn::NormType::RMSNorm); + nn::CausalSelfAttention at(D, H, /*max_len=*/8, /*seq_len=*/S, nn::PosEncodingType::RoPE, + /*subln=*/true, nn::NormType::RMSNorm); NN_EXIT(at.init(eng), 1, "CausalSelfAttention init 失败: "); auto ps = at.parameters(); // 参数序 = [w_q.w, w_q.b, w_k.w, w_k.b, w_v.w, w_v.b, w_o.w, w_o.b, @@ -177,8 +177,7 @@ void test_subln_changes_forward() nn::FeedForward off(D, FF, nn::ActivationType::GeLU, /*subln=*/false); NN_EXIT(off.init(eng), 1, "FeedForward(off) init 失败: "); - nn::FeedForward on(D, FF, nn::ActivationType::GeLU, /*subln=*/true, - nn::NormType::RMSNorm); + nn::FeedForward on(D, FF, nn::ActivationType::GeLU, /*subln=*/true, nn::NormType::RMSNorm); NN_EXIT(on.init(eng), 1, "FeedForward(on) init 失败: "); auto po = off.parameters(); // 4 张 @@ -203,7 +202,8 @@ void test_subln_changes_forward() // 反向消融：把 on 的 sub_norm.gamma 改成 0 → 中间激活整片归零 → fc2 偏置存活 { nn::Matrix zero(FF, 1); - for (auto& v : zero.span()) v = 0.0f; + for (auto& v : zero.span()) + v = 0.0f; NN_EXIT(eng.copy_from(pn[4].get(), zero), 1, "清零 sub_norm.gamma 失败: "); auto y0 = on.forward(x); NN_EXIT(y0, 1, "forward(gamma=0) 失败: "); @@ -223,19 +223,24 @@ void test_subln_spec_roundtrip() nn::CpuEngine& eng = *eng_owner; nn::GptConfig cfg; - cfg.vocab_size = V; cfg.d_model = D; cfg.seq_len = S; - cfg.num_heads = H; cfg.d_ff = FF; cfg.num_layers = L; + cfg.vocab_size = V; + cfg.d_model = D; + cfg.seq_len = S; + cfg.num_heads = H; + cfg.d_ff = FF; + cfg.num_layers = L; cfg.pos_enc = nn::PosEncodingType::Learned; cfg.activation = nn::ActivationType::SwiGLU; cfg.norm_type = nn::NormType::RMSNorm; cfg.subln = true; - const nn::ModelSpec spec = nn::make_gpt_spec( - V, D, S, H, FF, L, cfg.pos_enc, cfg.activation, cfg.norm_type, cfg.subln); + const nn::ModelSpec spec = + nn::make_gpt_spec(V, D, S, H, FF, L, cfg.pos_enc, cfg.activation, cfg.norm_type, cfg.subln); CHECK(spec.as()->subln, "make_gpt_spec 记录 subln"); nn::Matrix x_m(S, B); - for (std::size_t i = 0; i < x_m.size(); ++i) x_m.span()[i] = static_cast(i % V); + for (std::size_t i = 0; i < x_m.size(); ++i) + x_m.span()[i] = static_cast(i % V); nn::Matrix y_ref; { @@ -274,7 +279,8 @@ void test_subln_spec_roundtrip() void test_subln_spec_matches() { std::puts("[5] spec_matches：subln 不同 → 不一致"); - const auto base = [](bool subln) { + const auto base = [](bool subln) + { return nn::make_gpt_spec(64, 32, 8, 4, 64, 2, nn::PosEncodingType::Learned, nn::ActivationType::SwiGLU, nn::NormType::RMSNorm, subln); }; @@ -285,8 +291,8 @@ void test_subln_spec_matches() // ── S2/GQA 辅助：把 (n_kv·d_k, C) 的 KV 权重按**块映射**复制成 (n_head·d_k, C) ── // 与实现同一约定：MHA 的第 h 头 = GQA 的第 (h / n_rep) 头。 -nn::Matrix repeat_kv_rows(const nn::Matrix& src, std::size_t d_k, - std::size_t n_head, std::size_t n_kv) +nn::Matrix repeat_kv_rows(const nn::Matrix& src, std::size_t d_k, std::size_t n_head, + std::size_t n_kv) { const std::size_t n_rep = n_head / n_kv; nn::Matrix out(n_head * d_k, src.cols()); @@ -324,8 +330,8 @@ void test_gqa_mha_equivalence(nn::ComputeEngine& eng, const char* tag) } double wdiff = 0.0; for (std::size_t i = 0; i < pm.size(); ++i) - wdiff = std::max(wdiff, max_abs_diff(download(eng, pm[i].get()), - download(eng, pg[i].get()))); + wdiff = + std::max(wdiff, max_abs_diff(download(eng, pm[i].get()), download(eng, pg[i].get()))); CHECK(wdiff == 0.0, "显式拷贝后两层权重逐位相同"); auto x = make_tensor(eng, D, S * B, 0.15f, 0.02f); @@ -371,9 +377,8 @@ void test_gqa_head_mapping(nn::ComputeEngine& eng, const char* tag) CHECK(pg.size() == 8 && pm.size() == 8, "参数条数 8"); if (pg.size() != 8 || pm.size() != 8) return; for (std::size_t i = 0; i < 8; ++i) - std::printf(" param[%zu] gqa=%zux%zu mha=%zux%zu\n", i, - pg[i].get().rows(), pg[i].get().cols(), - pm[i].get().rows(), pm[i].get().cols()); + std::printf(" param[%zu] gqa=%zux%zu mha=%zux%zu\n", i, pg[i].get().rows(), + pg[i].get().cols(), pm[i].get().rows(), pm[i].get().cols()); // Q/O 直接拷；K/V（w 与 b）按块映射复制行 for (std::size_t i : {std::size_t{0}, std::size_t{1}, std::size_t{6}, std::size_t{7}}) @@ -381,8 +386,7 @@ void test_gqa_head_mapping(nn::ComputeEngine& eng, const char* tag) for (std::size_t i : {std::size_t{2}, std::size_t{3}, std::size_t{4}, std::size_t{5}}) { const nn::Matrix Wg = download(eng, pg[i].get()); - NN_EXIT(eng.copy_from(pm[i].get(), repeat_kv_rows(Wg, d_k, H, n_kv)), - 1, "拷贝 K/V 失败: "); + NN_EXIT(eng.copy_from(pm[i].get(), repeat_kv_rows(Wg, d_k, H, n_kv)), 1, "拷贝 K/V 失败: "); } auto x = make_tensor(eng, D, S * B, 0.15f, 0.02f); @@ -408,7 +412,8 @@ void test_gqa_head_mapping(nn::ComputeEngine& eng, const char* tag) { const nn::Matrix gg = download(eng, gqa.param_gradients()[i].get()); const nn::Matrix gm_full = download(eng, mha.param_gradients()[i].get()); - const nn::Matrix gm_fold = [&] { + const nn::Matrix gm_fold = [&] + { nn::Matrix o(gg.rows(), gg.cols()); const std::size_t n_rep = H / n_kv; for (std::size_t r = 0; r < gg.rows(); ++r) @@ -429,13 +434,12 @@ void test_gqa_head_mapping(nn::ComputeEngine& eng, const char* tag) void test_gqa_spec() { std::puts("[9] GQA 规格往返 + spec_matches"); - const auto mk = [](std::size_t n_kv) { + const auto mk = [](std::size_t n_kv) + { return nn::make_gpt_spec(64, 32, 8, 4, 64, 2, nn::PosEncodingType::Learned, - nn::ActivationType::SwiGLU, nn::NormType::RMSNorm, - false, n_kv); + nn::ActivationType::SwiGLU, nn::NormType::RMSNorm, false, n_kv); }; - CHECK(!mk(0).as()->n_head_kv && - mk(0).as()->n_head_kv_or_heads() == 4, + CHECK(!mk(0).as()->n_head_kv && mk(0).as()->n_head_kv_or_heads() == 4, "0 = num_heads（旧语义）"); CHECK(nn::spec_matches(mk(0), mk(4)), "0 与显式 num_heads 语义等价"); CHECK(!nn::spec_matches(mk(2), mk(4)), "n_head_kv 不同 → 不匹配"); @@ -446,19 +450,23 @@ void test_gqa_spec() auto eng_owner = nn::CpuEngine::create(); nn::CpuEngine& eng = *eng_owner; nn::GptConfig cfg; - cfg.vocab_size = V; cfg.d_model = D; cfg.seq_len = S; - cfg.num_heads = H; cfg.d_ff = FF; cfg.num_layers = L; + cfg.vocab_size = V; + cfg.d_model = D; + cfg.seq_len = S; + cfg.num_heads = H; + cfg.d_ff = FF; + cfg.num_layers = L; cfg.pos_enc = nn::PosEncodingType::Learned; cfg.activation = nn::ActivationType::SwiGLU; cfg.norm_type = nn::NormType::RMSNorm; cfg.subln = true; cfg.n_head_kv = 2; - const nn::ModelSpec spec = nn::make_gpt_spec( - V, D, S, H, FF, L, cfg.pos_enc, cfg.activation, cfg.norm_type, - cfg.subln, cfg.n_head_kv); + const nn::ModelSpec spec = nn::make_gpt_spec(V, D, S, H, FF, L, cfg.pos_enc, cfg.activation, + cfg.norm_type, cfg.subln, cfg.n_head_kv); nn::Matrix x_m(S, B); - for (std::size_t i = 0; i < x_m.size(); ++i) x_m.span()[i] = static_cast(i % V); + for (std::size_t i = 0; i < x_m.size(); ++i) + x_m.span()[i] = static_cast(i % V); nn::Matrix y_ref; { @@ -473,7 +481,11 @@ void test_gqa_spec() auto peeked = nn::peek_model_spec(file); CHECK(peeked.has_value() && peeked->as()->n_head_kv == 2, "规格记住 n_head_kv=2"); - if (!peeked.has_value()) { std::remove(file.c_str()); return; } + if (!peeked.has_value()) + { + std::remove(file.c_str()); + return; + } auto m2 = nn::build_gpt_model_from_spec(eng, *peeked); NN_EXIT(m2, 1, "build_gpt_model_from_spec(GQA) 失败: "); NN_EXIT(nn::load_model(file, *m2), 1, "load_model 失败: "); @@ -494,21 +506,22 @@ void test_tied_embedding() nn::CpuEngine& eng = *eng_owner; const std::size_t V = 64, D = 32, S = 8, H = 4, FF = 64, L = 2, B = 2, N = S * B; - auto mk = [&](bool tie) { - auto m = std::make_unique( - V, D, S, H, FF, L, nn::PosEncodingType::Learned, - nn::ActivationType::SwiGLU, nn::NormType::RMSNorm, - nn::PrecisionProfile{}, /*subln=*/false, /*n_head_kv=*/0, /*tie=*/tie); + auto mk = [&](bool tie) + { + auto m = std::make_unique(V, D, S, H, FF, L, nn::PosEncodingType::Learned, + nn::ActivationType::SwiGLU, nn::NormType::RMSNorm, + nn::PrecisionProfile{}, /*subln=*/false, + /*n_head_kv=*/0, /*tie=*/tie); NN_EXIT(m->init(eng), 1, "GPTModel init 失败: "); return m; }; auto off = mk(false); - auto on = mk(true); + auto on = mk(true); const std::size_t n_off = count_params(off->parameters()); - const std::size_t n_on = count_params(on->parameters()); - std::printf(" off=%zu on=%zu Δ=%zu（期望 %zu = V*D + V）\n", - n_off, n_on, n_off - n_on, V * D + V); + const std::size_t n_on = count_params(on->parameters()); + std::printf(" off=%zu on=%zu Δ=%zu（期望 %zu = V*D + V）\n", n_off, n_on, n_off - n_on, + V * D + V); CHECK(n_off - n_on == V * D + V, "tied 少 head 的 (w,b) 两张"); auto po = on->parameters(); @@ -522,11 +535,12 @@ void test_tied_embedding() const nn::Matrix w = download(eng, po[i].get()); NN_EXIT(eng.copy_from(pf[i].get(), w), 1, "拷贝参数失败: "); } - NN_EXIT(eng.copy_from(pf[po.size()].get(), download(eng, po[0].get())), - 1, "head.w ← token_emb 失败: "); + NN_EXIT(eng.copy_from(pf[po.size()].get(), download(eng, po[0].get())), 1, + "head.w ← token_emb 失败: "); { nn::Matrix zb(V, 1); - for (auto& v : zb.span()) v = 0.0f; + for (auto& v : zb.span()) + v = 0.0f; NN_EXIT(eng.copy_from(pf[po.size() + 1].get(), zb), 1, "head.b ← 0 失败: "); } @@ -570,10 +584,10 @@ void test_tied_embedding() void test_tied_spec() { std::puts("[11] tied embedding 规格往返 + spec_matches"); - const auto mk = [](bool tie) { + const auto mk = [](bool tie) + { return nn::make_gpt_spec(64, 32, 8, 4, 64, 2, nn::PosEncodingType::Learned, - nn::ActivationType::SwiGLU, nn::NormType::RMSNorm, - true, 2, tie); + nn::ActivationType::SwiGLU, nn::NormType::RMSNorm, true, 2, tie); }; CHECK(!mk(false).as()->tie_embeddings, "缺省 = false（旧文件语义）"); CHECK(!nn::spec_matches(mk(false), mk(true)), "tie 不同 → 不匹配"); @@ -584,20 +598,25 @@ void test_tied_spec() auto eng_owner = nn::CpuEngine::create(); nn::CpuEngine& eng = *eng_owner; nn::GptConfig cfg; - cfg.vocab_size = V; cfg.d_model = D; cfg.seq_len = S; - cfg.num_heads = H; cfg.d_ff = FF; cfg.num_layers = L; + cfg.vocab_size = V; + cfg.d_model = D; + cfg.seq_len = S; + cfg.num_heads = H; + cfg.d_ff = FF; + cfg.num_layers = L; cfg.pos_enc = nn::PosEncodingType::Learned; cfg.activation = nn::ActivationType::SwiGLU; cfg.norm_type = nn::NormType::RMSNorm; cfg.subln = true; cfg.n_head_kv = 2; cfg.tie_embeddings = true; - const nn::ModelSpec spec = nn::make_gpt_spec( - V, D, S, H, FF, L, cfg.pos_enc, cfg.activation, cfg.norm_type, - cfg.subln, cfg.n_head_kv, cfg.tie_embeddings); + const nn::ModelSpec spec = + nn::make_gpt_spec(V, D, S, H, FF, L, cfg.pos_enc, cfg.activation, cfg.norm_type, cfg.subln, + cfg.n_head_kv, cfg.tie_embeddings); nn::Matrix x_m(S, B); - for (std::size_t i = 0; i < x_m.size(); ++i) x_m.span()[i] = static_cast(i % V); + for (std::size_t i = 0; i < x_m.size(); ++i) + x_m.span()[i] = static_cast(i % V); nn::Matrix y_ref; { @@ -612,7 +631,11 @@ void test_tied_spec() auto peeked = nn::peek_model_spec(file); CHECK(peeked.has_value() && peeked->as()->tie_embeddings, "规格记住 tie_embeddings"); - if (!peeked.has_value()) { std::remove(file.c_str()); return; } + if (!peeked.has_value()) + { + std::remove(file.c_str()); + return; + } auto m2 = nn::build_gpt_model_from_spec(eng, *peeked); NN_EXIT(m2, 1, "build_gpt_model_from_spec(tied) 失败: "); NN_EXIT(nn::load_model(file, *m2), 1, "load_model 失败: "); @@ -654,8 +677,7 @@ void test_relu2_activation() const double g = x.at(r, c); const double u = x.at(FF + r, c); const double relu = g > 0.0 ? g : 0.0; - fmax = std::max(fmax, std::fabs(static_cast(ym.at(r, c)) - - relu * relu * u)); + fmax = std::max(fmax, std::fabs(static_cast(ym.at(r, c)) - relu * relu * u)); } std::printf(" forward vs 宿主 relu²(gate)·up 最大差 = %.6g\n", fmax); CHECK(fmax < 1e-6, "前向 = relu²(gate)·up"); @@ -675,8 +697,8 @@ void test_relu2_activation() const double g = x.at(o, c); const double u = x.at(FF + o, c); const double relu = g > 0.0 ? g : 0.0; - const double expect = (r < FF) ? go.at(o, c) * u * 2.0 * relu - : go.at(o, c) * relu * relu; + const double expect = + (r < FF) ? go.at(o, c) * u * 2.0 * relu : go.at(o, c) * relu * relu; bmax = std::max(bmax, std::fabs(static_cast(gim.at(r, c)) - expect)); } std::printf(" backward vs 宿主解析式 最大差 = %.6g\n", bmax); @@ -692,8 +714,12 @@ void test_relu2_gpt() const std::size_t V = 48, D = 32, S = 8, H = 4, FF = 64, L = 2, B = 2, N = S * B; nn::GptConfig cfg; - cfg.vocab_size = V; cfg.d_model = D; cfg.seq_len = S; - cfg.num_heads = H; cfg.d_ff = FF; cfg.num_layers = L; + cfg.vocab_size = V; + cfg.d_model = D; + cfg.seq_len = S; + cfg.num_heads = H; + cfg.d_ff = FF; + cfg.num_layers = L; cfg.pos_enc = nn::PosEncodingType::RoPE; cfg.activation = nn::ActivationType::ReLU2; cfg.norm_type = nn::NormType::RMSNorm; @@ -702,7 +728,8 @@ void test_relu2_gpt() cfg.tie_embeddings = true; // = 2B4T 的完整结构组合 nn::Matrix x_m(S, B); - for (std::size_t i = 0; i < x_m.size(); ++i) x_m.span()[i] = static_cast(i % V); + for (std::size_t i = 0; i < x_m.size(); ++i) + x_m.span()[i] = static_cast(i % V); auto m = nn::build_gpt_model(eng, cfg); NN_EXIT(m, 1, "build_gpt_model(relu2) 失败: "); @@ -711,7 +738,8 @@ void test_relu2_gpt() const nn::Matrix ym = download(eng, *y); CHECK(ym.rows() == V && ym.cols() == N, "logits 形状"); double mx = 0.0; - for (auto v : ym.span()) mx = std::max(mx, std::fabs(static_cast(v))); + for (auto v : ym.span()) + mx = std::max(mx, std::fabs(static_cast(v))); CHECK(std::isfinite(mx) && mx > 0.0, "logits 有限且非全零"); auto g = make_tensor(eng, V, N, 0.01f, 0.0002f); @@ -719,9 +747,9 @@ void test_relu2_gpt() NN_EXIT(bi, 1, "relu2 GPT backward 失败: "); // 规格：activation 进 spec_matches - const nn::ModelSpec spec = nn::make_gpt_spec( - V, D, S, H, FF, L, cfg.pos_enc, cfg.activation, cfg.norm_type, - cfg.subln, cfg.n_head_kv, cfg.tie_embeddings); + const nn::ModelSpec spec = + nn::make_gpt_spec(V, D, S, H, FF, L, cfg.pos_enc, cfg.activation, cfg.norm_type, cfg.subln, + cfg.n_head_kv, cfg.tie_embeddings); auto other = spec; other.as()->activation = nn::ActivationType::SwiGLU; CHECK(!nn::spec_matches(spec, other), "activation 不同 → 不匹配"); @@ -736,19 +764,21 @@ bool test_training(nn::ComputeEngine& eng, const char* tag) const std::size_t N = S * B; std::printf("[6] SubLN 端到端小训练（%s，L2 on logits）\n", tag); - nn::GPTModel model(V, D, S, H, FF, L, nn::PosEncodingType::Learned, - nn::ActivationType::SwiGLU, nn::NormType::RMSNorm, - nn::PrecisionProfile{}, /*subln=*/true); + nn::GPTModel model(V, D, S, H, FF, L, nn::PosEncodingType::Learned, nn::ActivationType::SwiGLU, + nn::NormType::RMSNorm, nn::PrecisionProfile{}, /*subln=*/true); if (auto r = model.init(eng); !r) { std::printf(" init 失败: %s\n", r.error().message.c_str()); return false; } - auto opt_r = nn::create_optimizer("adamw", eng, model.parameters(), - model.param_gradients(), + auto opt_r = nn::create_optimizer("adamw", eng, model.parameters(), model.param_gradients(), nn::Scalar{5e-3}, nn::Scalar{0.0}); - if (!opt_r) { std::printf(" create_optimizer 失败\n"); return false; } + if (!opt_r) + { + std::printf(" create_optimizer 失败\n"); + return false; + } nn::Optimizer& opt = *opt_r; // 固定输入 token ids（seq, batch）与固定目标 logits @@ -759,7 +789,8 @@ bool test_training(nn::ComputeEngine& eng, const char* tag) { std::mt19937 rng(11); std::uniform_real_distribution d(-1.0f, 1.0f); - for (auto& v : tgt.span()) v = d(rng); + for (auto& v : tgt.span()) + v = d(rng); } const nn::Tensor x = upload(eng, x_m); @@ -773,15 +804,18 @@ bool test_training(nn::ComputeEngine& eng, const char* tag) return false; } auto y = model.forward(x); - if (!y) { std::printf(" forward 失败: %s\n", y.error().message.c_str()); return false; } + if (!y) + { + std::printf(" forward 失败: %s\n", y.error().message.c_str()); + return false; + } const nn::Matrix ym = download(eng, *y); nn::Matrix g(V, N); double loss = 0.0; for (std::size_t i = 0; i < V * N; ++i) { - const double d = static_cast(ym.span()[i]) - - static_cast(tgt.span()[i]); + const double d = static_cast(ym.span()[i]) - static_cast(tgt.span()[i]); loss += d * d; g.span()[i] = static_cast(2.0 * d / static_cast(V * N)); } @@ -815,8 +849,8 @@ struct BitnetData std::vector labels; // batch-major：i = b*seq + t }; -[[nodiscard]] BitnetData make_bitnet_data(std::size_t V, std::size_t S, - std::size_t B, unsigned seed) +[[nodiscard]] BitnetData make_bitnet_data(std::size_t V, std::size_t S, std::size_t B, + unsigned seed) { BitnetData d; d.x = nn::Matrix(S, B); @@ -829,8 +863,7 @@ struct BitnetData for (std::size_t b = 0; b < B; ++b) for (std::size_t t = 0; t < S; ++t) { - const std::size_t cur = - static_cast(d.x.span()[t * B + b]); + const std::size_t cur = static_cast(d.x.span()[t * B + b]); d.labels[b * S + t] = (cur + 1) % V; } return d; @@ -851,11 +884,11 @@ void test_gpt_ternary_wiring() nn::PrecisionProfile t1 = f32p; t1.param = nn::Precision::T1_58; - const auto build = [&](const nn::PrecisionProfile& pp) { - return nn::build_gpt_model(eng, V, D, S, H, FF, L, - nn::PosEncodingType::RoPE, - nn::ActivationType::ReLU2, nn::NormType::RMSNorm, - pp, /*subln=*/true, /*n_head_kv=*/KV, + const auto build = [&](const nn::PrecisionProfile& pp) + { + return nn::build_gpt_model(eng, V, D, S, H, FF, L, nn::PosEncodingType::RoPE, + nn::ActivationType::ReLU2, nn::NormType::RMSNorm, pp, + /*subln=*/true, /*n_head_kv=*/KV, /*tie_embeddings=*/false); }; auto base = build(f32p); @@ -864,8 +897,7 @@ void test_gpt_ternary_wiring() "三值 GPT 构建+init 成功（4 投影/fc1/fc2/LM head 均 BitLinear）"); if (!base.has_value() || !tern.has_value()) { - if (!tern.has_value()) - std::printf(" init 失败: %s\n", tern.error().message.c_str()); + if (!tern.has_value()) std::printf(" init 失败: %s\n", tern.error().message.c_str()); return; } @@ -873,10 +905,10 @@ void test_gpt_ternary_wiring() const auto& pt = tern->parameters(); bool same = (pb.size() == pt.size()); for (std::size_t i = 0; same && i < pb.size(); ++i) - same = (pb[i].get().rows() == pt[i].get().rows() && - pb[i].get().cols() == pt[i].get().cols()); - std::printf(" 参数：%zu 张量 / %zu 标量（f32 基线 %zu / %zu）\n", - pt.size(), count_params(pt), pb.size(), count_params(pb)); + same = + (pb[i].get().rows() == pt[i].get().rows() && pb[i].get().cols() == pt[i].get().cols()); + std::printf(" 参数：%zu 张量 / %zu 标量（f32 基线 %zu / %zu）\n", pt.size(), + count_params(pt), pb.size(), count_params(pb)); CHECK(same, "逐张量形状与 f32 基线一致（BitLinear latent 恒 f32）"); CHECK(count_params(pb) == count_params(pt), "参数量一致"); @@ -888,8 +920,7 @@ void test_gpt_ternary_wiring() NN_EXIT(tern->zero_grad(), 1, "zero_grad 失败: "); auto logits = tern->forward(xt); NN_EXIT(logits, 1, "三值 GPT forward 失败: "); - auto loss = ce.forward_sparse(eng, *logits, - std::span(data.labels), + auto loss = ce.forward_sparse(eng, *logits, std::span(data.labels), std::span{}, V); NN_EXIT(loss, 1, "sparse CE 失败: "); CHECK(std::isfinite(*loss), "三值 GPT 前向 loss 有限"); @@ -918,19 +949,16 @@ double test_bitnet_training(nn::ComputeEngine& eng, const char* tag) { const std::size_t V = 32, D = 32, S = 8, H = 4, KV = 2, FF = 64, L = 2, B = 4; constexpr int STEPS = 200; - std::printf("[15] 真实 BitNet（2B4T 形状）端到端训练（%s，CE + AdamW，%d 步）\n", - tag, STEPS); + std::printf("[15] 真实 BitNet（2B4T 形状）端到端训练（%s，CE + AdamW，%d 步）\n", tag, STEPS); const BitnetData data = make_bitnet_data(V, S, B, 7u); const nn::Tensor xt = upload(eng, data.x); const nn::PrecisionProfile f32p; - const auto train_one = [&](const nn::PrecisionProfile& pp, bool tie, - const char* name) -> double { - auto m_r = nn::build_gpt_model(eng, V, D, S, H, FF, L, - nn::PosEncodingType::RoPE, - nn::ActivationType::ReLU2, - nn::NormType::RMSNorm, pp, + const auto train_one = [&](const nn::PrecisionProfile& pp, bool tie, const char* name) -> double + { + auto m_r = nn::build_gpt_model(eng, V, D, S, H, FF, L, nn::PosEncodingType::RoPE, + nn::ActivationType::ReLU2, nn::NormType::RMSNorm, pp, /*subln=*/true, /*n_head_kv=*/KV, tie); if (!m_r.has_value()) { @@ -939,10 +967,13 @@ double test_bitnet_training(nn::ComputeEngine& eng, const char* tag) return -1.0; } nn::Model& model = *m_r; - auto opt = nn::create_optimizer("adamw", eng, model.parameters(), - model.param_gradients(), + auto opt = nn::create_optimizer("adamw", eng, model.parameters(), model.param_gradients(), nn::Scalar{3e-3}, nn::Scalar{0.0f}, pp); - if (!opt) { CHECK(false, "优化器创建成功"); return -1.0; } + if (!opt) + { + CHECK(false, "优化器创建成功"); + return -1.0; + } nn::CrossEntropyLoss ce; ce.set_precision_profile(pp); @@ -950,22 +981,49 @@ double test_bitnet_training(nn::ComputeEngine& eng, const char* tag) for (int step = 0; step < STEPS; ++step) { auto zg = model.zero_grad(); - if (!zg) { std::printf(" [%s] zero_grad: %s\n", name, zg.error().message.c_str()); return -1.0; } + if (!zg) + { + std::printf(" [%s] zero_grad: %s\n", name, zg.error().message.c_str()); + return -1.0; + } auto logits = model.forward(xt); - if (!logits) { std::printf(" [%s] forward: %s\n", name, logits.error().message.c_str()); return -1.0; } - auto loss = ce.forward_sparse(eng, *logits, - std::span(data.labels), + if (!logits) + { + std::printf(" [%s] forward: %s\n", name, logits.error().message.c_str()); + return -1.0; + } + auto loss = ce.forward_sparse(eng, *logits, std::span(data.labels), std::span{}, V); - if (!loss) { std::printf(" [%s] loss: %s\n", name, loss.error().message.c_str()); return -1.0; } - if (!std::isfinite(*loss)) { CHECK(false, "loss 有限"); return -1.0; } + if (!loss) + { + std::printf(" [%s] loss: %s\n", name, loss.error().message.c_str()); + return -1.0; + } + if (!std::isfinite(*loss)) + { + CHECK(false, "loss 有限"); + return -1.0; + } if (step == 0) first = *loss; last = *loss; auto grad = ce.backward(); - if (!grad) { std::printf(" [%s] ce.backward 失败\n", name); return -1.0; } + if (!grad) + { + std::printf(" [%s] ce.backward 失败\n", name); + return -1.0; + } auto bw = model.backward(*grad); - if (!bw) { std::printf(" [%s] backward: %s\n", name, bw.error().message.c_str()); return -1.0; } + if (!bw) + { + std::printf(" [%s] backward: %s\n", name, bw.error().message.c_str()); + return -1.0; + } auto st = opt->step(); - if (!st) { std::printf(" [%s] step: %s\n", name, st.error().message.c_str()); return -1.0; } + if (!st) + { + std::printf(" [%s] step: %s\n", name, st.error().message.c_str()); + return -1.0; + } } std::printf(" [%s] loss: %.6f → %.6f\n", name, first, last); return last; @@ -975,7 +1033,7 @@ double test_bitnet_training(nn::ComputeEngine& eng, const char* tag) t1.param = nn::Precision::T1_58; const double t1_tied = train_one(t1, /*tie=*/true, "bitnet_t1_58_tied"); const double t1_head = train_one(t1, /*tie=*/false, "bitnet_t1_58_lmhead"); - const double base = train_one(f32p, /*tie=*/false, "f32_baseline"); + const double base = train_one(f32p, /*tie=*/false, "f32_baseline"); std::printf(" （随机基线 ln(V)=%.4f）\n", std::log(static_cast(V))); CHECK(t1_tied >= 0.0 && t1_tied < 0.30, "三值 + tied（2B4T 形状）收敛到 < 0.30"); @@ -1030,13 +1088,12 @@ int main(int argc, char* argv[]) test_gqa_head_mapping(**eng_r, "GPU"); if (!test_training(**eng_r, "GPU")) g_failures++; const double gpu_bitnet_loss = test_bitnet_training(**eng_r, "GPU"); - if (gpu_bitnet_loss < 0.0) - ++g_failures; + if (gpu_bitnet_loss < 0.0) ++g_failures; else { const double d = std::fabs(gpu_bitnet_loss - cpu_bitnet_loss); - std::printf(" [parity] 末步 loss CPU=%.6f GPU=%.6f 差=%.3e\n", - cpu_bitnet_loss, gpu_bitnet_loss, d); + std::printf(" [parity] 末步 loss CPU=%.6f GPU=%.6f 差=%.3e\n", cpu_bitnet_loss, + gpu_bitnet_loss, d); // 阈值 2e-3（原 1e-3）：合并后 GPU 的 wgrad 走 matmul K 切分 // （`#ks`/`#kr`，跨段 K 累加序 ≠ 逐 k 串联，main 侧已裁定"数学等价 // 不保证逐位一致"）→ 200 步训练的末步 loss 差被抬高到 1.051e-3。 diff --git a/src/ce_fusion_test.cpp b/src/ce_fusion_test.cpp index 118f8bdf..472ae3d8 100644 --- a/src/ce_fusion_test.cpp +++ b/src/ce_fusion_test.cpp @@ -69,10 +69,9 @@ nn::Matrix ref_denom(const nn::Matrix& logits, const nn::Matrix& col_max) } // col_softmax_sparse_forward 参考：grad 与 loss_vec -nn::Matrix ref_sparse_forward( - const nn::Matrix& logits, const std::vector& labels, - const std::vector* loss_mask, std::size_t vocab_size, - Scalar inv_num_valid, nn::Matrix& lv_out) +nn::Matrix ref_sparse_forward(const nn::Matrix& logits, const std::vector& labels, + const std::vector* loss_mask, std::size_t vocab_size, + Scalar inv_num_valid, nn::Matrix& lv_out) { const std::size_t C = logits.rows(), N = logits.cols(); nn::Matrix grad(C, N); @@ -96,8 +95,8 @@ nn::Matrix ref_sparse_forward( for (std::size_t r = 0; r < C; ++r) denom += std::exp(logits.at_unchecked(r, i) - mv); for (std::size_t r = 0; r < C; ++r) - grad.set_value_unchecked(r, i, - inv_num_valid * std::exp(logits.at_unchecked(r, i) - mv) / denom); + grad.set_value_unchecked( + r, i, inv_num_valid * std::exp(logits.at_unchecked(r, i) - mv) / denom); grad.set_value_unchecked(lbl, i, grad.at_unchecked(lbl, i) - inv_num_valid); lv.set_value_unchecked(0, i, logits.at_unchecked(lbl, i) - mv - std::log(denom)); } @@ -113,7 +112,8 @@ int run_case(nn::ComputeEngine& eng, const char* tag) std::uniform_real_distribution dist(-1.5f, 1.5f); nn::Matrix logits(C, N); - for (auto& x : logits.span()) x = dist(rng); + for (auto& x : logits.span()) + x = dist(rng); nn::Matrix col_max(1, N); for (std::size_t i = 0; i < N; ++i) { @@ -151,9 +151,19 @@ int run_case(nn::ComputeEngine& eng, const char* tag) const auto check = [&](const char* nm, nn::Result&& r, const nn::Matrix& ref) { - if (!r) { std::printf("[FAIL] %s: %s\n", nm, r.error().message.c_str()); ++fail; return; } + if (!r) + { + std::printf("[FAIL] %s: %s\n", nm, r.error().message.c_str()); + ++fail; + return; + } auto dm = eng.to_matrix(*r); - if (!dm) { std::printf("[FAIL] %s: 下载失败: %s\n", nm, dm.error().message.c_str()); ++fail; return; } + if (!dm) + { + std::printf("[FAIL] %s: 下载失败: %s\n", nm, dm.error().message.c_str()); + ++fail; + return; + } if (!check_matrix(*dm, ref, nm)) ++fail; }; @@ -165,7 +175,8 @@ int run_case(nn::ComputeEngine& eng, const char* tag) char nm[128]; std::snprintf(nm, sizeof(nm), "%s IR denom", tag); const nn::Matrix ref = ref_denom(logits, col_max); - auto r = nn::dsl::compute_reduce(eng, + auto r = nn::dsl::compute_reduce( + eng, nn::dsl::col_reduce_sum( nn::dsl::exp(nn::dsl::leaf(t_logits) - nn::dsl::col_broadcast(t_colmax))), C, N); @@ -194,24 +205,20 @@ int run_case(nn::ComputeEngine& eng, const char* tag) s += logits.at_unchecked(r, i) * logits.at_unchecked(r, i); ref.set_value_unchecked(0, i, Scalar{1} / std::sqrt(s * inv_c + eps)); } - auto r = nn::dsl::compute_reduce(eng, - nn::dsl::rsqrt(nn::dsl::col_reduce_sum( - nn::dsl::leaf(t_logits) * nn::dsl::leaf(t_logits)) - * nn::dsl::rparam(inv_c) - + nn::dsl::rparam(eps)), + auto r = nn::dsl::compute_reduce( + eng, + nn::dsl::rsqrt( + nn::dsl::col_reduce_sum(nn::dsl::leaf(t_logits) * nn::dsl::leaf(t_logits)) * + nn::dsl::rparam(inv_c) + + nn::dsl::rparam(eps)), C, N); - if (eng.device() == nn::Device::CPU) - { - check(nm, std::move(r), ref); - } + if (eng.device() == nn::Device::CPU) { check(nm, std::move(r), ref); } else { // GPU：必须显式拒绝（"静默错值"是更坏的失败模式） const bool rejected = !r; - std::printf("[%s] %s（GPU 应硬报错而非给出错值）%s%s\n", - rejected ? "PASS" : "FAIL", nm, - rejected ? " 错误: " : "", - rejected ? r.error().message.c_str() : ""); + std::printf("[%s] %s（GPU 应硬报错而非给出错值）%s%s\n", rejected ? "PASS" : "FAIL", nm, + rejected ? " 错误: " : "", rejected ? r.error().message.c_str() : ""); if (!rejected) ++fail; } } @@ -242,20 +249,22 @@ int run_case(nn::ComputeEngine& eng, const char* tag) nn::Tensor dsl_t = upload(eng, am); nn::Tensor prim_t = upload(eng, am); nn::Tensor other = upload(eng, bm); - auto r = nn::dsl::compute_into(eng, - nn::dsl::leaf(dsl_t) + nn::dsl::leaf(other), dsl_t); + auto r = nn::dsl::compute_into(eng, nn::dsl::leaf(dsl_t) + nn::dsl::leaf(other), dsl_t); char nm[128]; if (!r) { - std::printf("[FAIL] %s compute_into (A+=B): %s\n", tag, - r.error().message.c_str()); + std::printf("[FAIL] %s compute_into (A+=B): %s\n", tag, r.error().message.c_str()); ++fail; } else { auto dm = eng.to_matrix(dsl_t); std::snprintf(nm, sizeof(nm), "%s compute_into (A+=B)", tag); - if (!dm) { std::printf("[FAIL] %s 下载失败\n", nm); ++fail; } + if (!dm) + { + std::printf("[FAIL] %s 下载失败\n", nm); + ++fail; + } else if (!check_matrix(*dm, ref_add, nm)) ++fail; // 与原地原语逐字节一致 (void)eng.add_inplace(prim_t, other); @@ -270,20 +279,22 @@ int run_case(nn::ComputeEngine& eng, const char* tag) { nn::Tensor dsl_t = upload(eng, am); nn::Tensor prim_t = upload(eng, am); - auto r = nn::dsl::compute_into(eng, - nn::dsl::leaf(dsl_t) * nn::dsl::rparam(s), dsl_t); + auto r = nn::dsl::compute_into(eng, nn::dsl::leaf(dsl_t) * nn::dsl::rparam(s), dsl_t); char nm[128]; if (!r) { - std::printf("[FAIL] %s compute_into (A*=s): %s\n", tag, - r.error().message.c_str()); + std::printf("[FAIL] %s compute_into (A*=s): %s\n", tag, r.error().message.c_str()); ++fail; } else { auto dm = eng.to_matrix(dsl_t); std::snprintf(nm, sizeof(nm), "%s compute_into (A*=s)", tag); - if (!dm) { std::printf("[FAIL] %s 下载失败\n", nm); ++fail; } + if (!dm) + { + std::printf("[FAIL] %s 下载失败\n", nm); + ++fail; + } else if (!check_matrix(*dm, ref_scale, nm)) ++fail; (void)eng.scale_inplace(prim_t, s); auto pm = eng.to_matrix(prim_t); @@ -296,11 +307,10 @@ int run_case(nn::ComputeEngine& eng, const char* tag) // (3) 归约表达式必须被拒绝（compute_into 只支持逐元素） { nn::Tensor dsl_t = upload(eng, am); - auto r = nn::dsl::compute_into(eng, - nn::dsl::col_reduce_sum(nn::dsl::leaf(dsl_t)), dsl_t); + auto r = + nn::dsl::compute_into(eng, nn::dsl::col_reduce_sum(nn::dsl::leaf(dsl_t)), dsl_t); const bool rejected = !r; - std::printf("[%s] %s compute_into 拒绝归约表达式\n", - rejected ? "PASS" : "FAIL", tag); + std::printf("[%s] %s compute_into 拒绝归约表达式\n", rejected ? "PASS" : "FAIL", tag); if (!rejected) ++fail; } } @@ -326,46 +336,44 @@ int run_case(nn::ComputeEngine& eng, const char* tag) { char nm[128]; // lse = fold logsumexp（与 compute_loss.hpp 同源构造 → 同 key） - auto lse_r = eng.eval_expr( - nn::expr::make_fold_logsumexp(C, nn::expr::trans(N)), - std::span{&t_logits, 1}, N, 1); + auto lse_r = eng.eval_expr(nn::expr::make_fold_logsumexp(C, nn::expr::trans(N)), + std::span{&t_logits, 1}, N, 1); if (!lse_r) { - std::printf("[FAIL] %s lse fold: %s\n", nm_base, - lse_r.error().message.c_str()); + std::printf("[FAIL] %s lse fold: %s\n", nm_base, lse_r.error().message.c_str()); ++fail; return; } auto lse_rs = eng.reshape(*lse_r, 1, N); if (!lse_rs) { - std::printf("[FAIL] %s lse reshape: %s\n", nm_base, - lse_rs.error().message.c_str()); + std::printf("[FAIL] %s lse reshape: %s\n", nm_base, lse_rs.error().message.c_str()); ++fail; return; } // loss_sum = Σ_c (rg(logits) − lse)·mask（与 CE 同构链） - auto lv = nn::dsl::compute_reduce(eng, + auto lv = nn::dsl::compute_reduce( + eng, nn::dsl::row_reduce_sum( - (nn::dsl::row_gather(t_logits, t_lab) - - nn::dsl::col_broadcast(*lse_rs)) - * nn::dsl::col_broadcast(t_msk)), + (nn::dsl::row_gather(t_logits, t_lab) - nn::dsl::col_broadcast(*lse_rs)) * + nn::dsl::col_broadcast(t_msk)), 1, N); nn::Matrix lv_sum(1, 1); { Scalar s = 0; - for (const Scalar v : lv_ref.span()) s += v; + for (const Scalar v : lv_ref.span()) + s += v; lv_sum.set_value_unchecked(0, 0, s); } std::snprintf(nm, sizeof(nm), "%s IR loss_sum %s", tag, nm_base); check(nm, std::move(lv), lv_sum); // grad（与 CE 同构链） - auto g = nn::dsl::compute(eng, - (nn::dsl::exp(nn::dsl::leaf(t_logits) - nn::dsl::col_broadcast(*lse_rs)) - - nn::dsl::select(nn::dsl::row() == nn::dsl::col_broadcast(t_lab), - Scalar{1}, Scalar{0})) - * nn::dsl::col_broadcast(t_msk) - * nn::dsl::rparam(inv_num_valid), + auto g = nn::dsl::compute( + eng, + (nn::dsl::exp(nn::dsl::leaf(t_logits) - nn::dsl::col_broadcast(*lse_rs)) - + nn::dsl::select(nn::dsl::row() == nn::dsl::col_broadcast(t_lab), Scalar{1}, + Scalar{0})) * + nn::dsl::col_broadcast(t_msk) * nn::dsl::rparam(inv_num_valid), C, N); std::snprintf(nm, sizeof(nm), "%s IR grad %s", tag, nm_base); check(nm, std::move(g), g_ref); @@ -374,13 +382,12 @@ int run_case(nn::ComputeEngine& eng, const char* tag) // 无 mask { nn::Matrix lv_ref; - const nn::Matrix g_ref = ref_sparse_forward( - logits, labels, nullptr, vocab_size, inv_num_valid, lv_ref); + const nn::Matrix g_ref = + ref_sparse_forward(logits, labels, nullptr, vocab_size, inv_num_valid, lv_ref); // Layer 行为：mask = valid 修正（越界 label → 0） nn::Matrix mask_m(1, N); for (std::size_t i = 0; i < N; ++i) - mask_m.set_value_unchecked(0, i, - labels[i] < vocab_size ? Scalar{1} : Scalar{0}); + mask_m.set_value_unchecked(0, i, labels[i] < vocab_size ? Scalar{1} : Scalar{0}); const nn::Tensor t_mask = upload(eng, mask_m); case_run(t_labels, t_mask, g_ref, lv_ref, "(no mask)"); } @@ -389,12 +396,13 @@ int run_case(nn::ComputeEngine& eng, const char* tag) nn::Matrix mask_m(1, N); for (std::size_t i = 0; i < N; ++i) mask_m.set_value_unchecked(0, i, - (loss_mask[i] >= Scalar{0.5} && labels[i] < vocab_size) - ? Scalar{1} : Scalar{0}); + (loss_mask[i] >= Scalar{0.5} && labels[i] < vocab_size) + ? Scalar{1} + : Scalar{0}); const nn::Tensor t_mask = upload(eng, mask_m); nn::Matrix lv_ref; - const nn::Matrix g_ref = ref_sparse_forward( - logits, labels, &loss_mask, vocab_size, inv_num_valid, lv_ref); + const nn::Matrix g_ref = + ref_sparse_forward(logits, labels, &loss_mask, vocab_size, inv_num_valid, lv_ref); case_run(t_labels, t_mask, g_ref, lv_ref, "(mask)"); } // 全越界（num_valid=0）：mask 修正为 0 → grad 全 0、loss_sum 0 @@ -405,8 +413,8 @@ int run_case(nn::ComputeEngine& eng, const char* tag) bad_m.set_value_unchecked(0, i, static_cast(bad[i])); const nn::Tensor t_bad = upload(eng, bad_m); nn::Matrix lv_ref; - const nn::Matrix g_ref = ref_sparse_forward( - logits, bad, nullptr, vocab_size, Scalar{0}, lv_ref); + const nn::Matrix g_ref = + ref_sparse_forward(logits, bad, nullptr, vocab_size, Scalar{0}, lv_ref); // 越界 label 修正为 0（Layer 行为）+ mask 全 0（越界无效） nn::Matrix mask_m(1, N); for (std::size_t i = 0; i < N; ++i) @@ -421,7 +429,11 @@ int run_case(nn::ComputeEngine& eng, const char* tag) char nm[128]; nn::CrossEntropyLoss loss_layer; auto loss_r = loss_layer.forward_sparse(eng, t_logits, labels, loss_mask, vocab_size); - if (!loss_r) { std::printf("[FAIL] %s: %s\n", tag, loss_r.error().message.c_str()); ++fail; } + if (!loss_r) + { + std::printf("[FAIL] %s: %s\n", tag, loss_r.error().message.c_str()); + ++fail; + } else { // 参考 loss = -(1/num_valid) Σ log_softmax[labels[i],i]（有效位置） @@ -441,19 +453,23 @@ int run_case(nn::ComputeEngine& eng, const char* tag) ref_loss /= static_cast(num_valid); const Scalar err = std::fabs(*loss_r - ref_loss); const bool ok = err < 1e-4f; - std::printf("[%s] %s CE loss got=%.6f ref=%.6f err=%.3e\n", - ok ? "PASS" : "FAIL", tag, *loss_r, ref_loss, err); + std::printf("[%s] %s CE loss got=%.6f ref=%.6f err=%.3e\n", ok ? "PASS" : "FAIL", tag, + *loss_r, ref_loss, err); if (!ok) ++fail; // grad vs 参考（softmax - one_hot）/num_valid auto grad_r = loss_layer.backward(); auto dg = eng.to_matrix(*grad_r); - if (!dg) { std::printf("[FAIL] %s CE grad: 下载失败\n", tag); ++fail; } + if (!dg) + { + std::printf("[FAIL] %s CE grad: 下载失败\n", tag); + ++fail; + } else { nn::Matrix lv_ref; - const nn::Matrix g_ref = ref_sparse_forward( - logits, labels, &loss_mask, vocab_size, inv_num_valid, lv_ref); + const nn::Matrix g_ref = ref_sparse_forward(logits, labels, &loss_mask, vocab_size, + inv_num_valid, lv_ref); std::snprintf(nm, sizeof(nm), "%s CE grad", tag); if (!check_matrix(*dg, g_ref, nm)) ++fail; } @@ -472,11 +488,14 @@ int run_case(nn::ComputeEngine& eng, const char* tag) int run_old_new_parity(nn::ComputeEngine& eng, const char* tag) { int fail = 0; - struct Case { std::size_t C, N; Scalar scale; }; + struct Case + { + std::size_t C, N; + Scalar scale; + }; const Case cases[] = { - {8, 12, 1.5f}, {129, 2, 2.0f}, {300, 17, 2.0f}, {257, 64, 2.0f}, - {1000, 7, 2.0f}, {3, 100, 2.0f}, {4096, 16, 2.0f}, - {256, 64, 40.0f}, // 大幅值应力（数值稳定路径） + {8, 12, 1.5f}, {129, 2, 2.0f}, {300, 17, 2.0f}, {257, 64, 2.0f}, {1000, 7, 2.0f}, + {3, 100, 2.0f}, {4096, 16, 2.0f}, {256, 64, 40.0f}, // 大幅值应力（数值稳定路径） }; constexpr Scalar tol = 1e-6f; std::mt19937 rng(20261006); @@ -485,7 +504,8 @@ int run_old_new_parity(nn::ComputeEngine& eng, const char* tag) const std::size_t C = cs.C, N = cs.N; nn::Matrix logits(C, N); std::uniform_real_distribution dist(-cs.scale, cs.scale); - for (auto& x : logits.span()) x = dist(rng); + for (auto& x : logits.span()) + x = dist(rng); std::vector labels(N); for (std::size_t i = 0; i < N; ++i) labels[i] = static_cast(rng() % (C + 3)); // 少量越界 @@ -499,10 +519,10 @@ int run_old_new_parity(nn::ComputeEngine& eng, const char* tag) // 旧数学参考（grad 为 exp(x−m)/denom 形态、lv 为 x−m−log(denom)） nn::Matrix lv_ref; - const nn::Matrix g_ref = - ref_sparse_forward(logits, labels, &mask, C, inv, lv_ref); + const nn::Matrix g_ref = ref_sparse_forward(logits, labels, &mask, C, inv, lv_ref); Scalar ref_loss = 0; - for (const Scalar v : lv_ref.span()) ref_loss -= v; + for (const Scalar v : lv_ref.span()) + ref_loss -= v; if (nv) ref_loss *= inv; // 新实现（P4：fold lse + exp(x−s)） @@ -542,9 +562,8 @@ int run_old_new_parity(nn::ComputeEngine& eng, const char* tag) if (e > bound) ok = false; if (e > grad_err) grad_err = e; } - std::printf("[%s] parity(%zu,%zu,x%g) loss_err=%.3e grad_err=%.3e\n", - ok ? "PASS" : "FAIL", C, N, static_cast(cs.scale), - loss_err, grad_err); + std::printf("[%s] parity(%zu,%zu,x%g) loss_err=%.3e grad_err=%.3e\n", ok ? "PASS" : "FAIL", + C, N, static_cast(cs.scale), loss_err, grad_err); if (!ok) ++fail; } return fail; @@ -612,8 +631,7 @@ int main() { nn::CrossEntropyLoss ce; const std::span mask_span(lm.span()); - const std::size_t slots = - std::max(gpu_engine.scalar_readback_slots(), 1); + const std::size_t slots = std::max(gpu_engine.scalar_readback_slots(), 1); constexpr int ITERS = 24; Scalar max_diff = 0; int mismatches = 0, spins_total = 0; @@ -621,8 +639,7 @@ int main() for (int it = 0; it < ITERS; ++it) { // 同步参考（to_matrix 下载） - auto ref = ce.forward_sparse(gpu_engine, *t_logits_r, labels, - mask_span, C); + auto ref = ce.forward_sparse(gpu_engine, *t_logits_r, labels, mask_span, C); if (!ref) { std::printf("[FAIL] async readback: forward_sparse: %s\n", @@ -634,16 +651,15 @@ int main() // 异步路径（与 text_train 完全同序） auto bb = gpu_engine.begin_batch(); std::size_t nv = 0; - auto sum_r = ce.forward_sparse_sum(gpu_engine, *t_logits_r, labels, - mask_span, C, nv); + auto sum_r = + ce.forward_sparse_sum(gpu_engine, *t_logits_r, labels, mask_span, C, nv); if (bb && sum_r) { auto f32_r = gpu_engine.cast(*sum_r, nn::Precision::F32); if (f32_r) { auto fb = gpu_engine.flush_batch(); // 提交主帧（不等待） - const std::size_t slot = - static_cast(it) % slots; + const std::size_t slot = static_cast(it) % slots; auto sr = gpu_engine.submit_scalar_readback(slot, *f32_r); if (!fb || !sr) { @@ -665,8 +681,7 @@ int main() } if (*pr) break; ++spins; - std::this_thread::sleep_for( - std::chrono::microseconds(200)); + std::this_thread::sleep_for(std::chrono::microseconds(200)); if (spins > 20000) // 4s 上限，防死等 { std::printf("[FAIL] async readback: poll 超时\n"); @@ -745,10 +760,8 @@ int main() nn::CrossEntropyLoss ce_a, ce_b; std::size_t nv_a = 0, nv_b = 0; auto bb = gpu_engine.begin_batch(); - auto sa = ce_a.forward_sparse_sum(gpu_engine, *a_logits, labels, - mask_span, C, nv_a); - auto sb = ce_b.forward_sparse_sum(gpu_engine, *b_logits, labels, - mask_span, C, nv_b, + auto sa = ce_a.forward_sparse_sum(gpu_engine, *a_logits, labels, mask_span, C, nv_a); + auto sb = ce_b.forward_sparse_sum(gpu_engine, *b_logits, labels, mask_span, C, nv_b, /*grad_reuse=*/&(*b_logits)); ok = (bool)bb && (bool)sa && (bool)sb; if (ok) @@ -771,8 +784,8 @@ int main() auto eb = gpu_engine.end_batch(); if (!eb) ok = false; std::cout << (ok ? "[PASS] " : "[FAIL] ") - << "CE 梯度原地复用（loss 一致 + grad max_diff=" - << grad_diff << "）" << std::endl; + << "CE 梯度原地复用（loss 一致 + grad max_diff=" << grad_diff << "）" + << std::endl; if (!ok) ++fail; } } @@ -781,4 +794,3 @@ int main() std::cout << (fail == 0 ? "\nALL PASS\n" : "\nFAILED\n"); return fail == 0 ? 0 : 1; } - diff --git a/src/chat_prompt_test.cpp b/src/chat_prompt_test.cpp index c48fb52d..c6e5f6d1 100644 --- a/src/chat_prompt_test.cpp +++ b/src/chat_prompt_test.cpp @@ -59,37 +59,37 @@ std::string native_path(std::string p) #endif // ── 小工具 ─────────────────────────────────────────────────────────────── -std::string dirname_of(const std::string &path) +std::string dirname_of(const std::string& path) { const auto pos = path.find_last_of("/\\"); return pos == std::string::npos ? std::string(".") : path.substr(0, pos); } -bool file_exists(const std::string &p) +bool file_exists(const std::string& p) { std::ifstream f(p); return f.good(); } -std::string read_file(const std::string &p) +std::string read_file(const std::string& p) { std::ifstream f(p, std::ios::binary); - return std::string((std::istreambuf_iterator(f)), - std::istreambuf_iterator()); + return std::string((std::istreambuf_iterator(f)), std::istreambuf_iterator()); } std::string rstrip_newline(std::string s) { - while (!s.empty() && (s.back() == '\n' || s.back() == '\r')) s.pop_back(); + while (!s.empty() && (s.back() == '\n' || s.back() == '\r')) + s.pop_back(); return s; } // 在测试二进制附近寻找同构建的 app（与 cli_precision_warning_test 同款） -std::string find_app(const std::string &name) +std::string find_app(const std::string& name) { const std::string dir = dirname_of(g_self); const std::string candidates[] = {dir + "/../" + name, dir + "/" + name, name}; - for (const auto &c : candidates) + for (const auto& c : candidates) if (file_exists(c)) { #ifdef _WIN32 @@ -99,7 +99,7 @@ std::string find_app(const std::string &name) #endif } #ifdef _WIN32 - for (const auto &c : candidates) + for (const auto& c : candidates) if (file_exists(c + ".exe")) return native_path(c + ".exe"); #endif return {}; @@ -147,7 +147,7 @@ struct PlainTokenizer : nn::BPETokenizer void clear_markers() { markers_.system = nn::Tokenizer::npos; } }; -[[nodiscard]] std::string join_ids(const std::vector &ids) +[[nodiscard]] std::string join_ids(const std::vector& ids) { std::string s; for (std::size_t i = 0; i < ids.size(); ++i) @@ -158,8 +158,8 @@ struct PlainTokenizer : nn::BPETokenizer return s; } -[[nodiscard]] bool is_prefix(const std::vector &pre, - const std::vector &full) +[[nodiscard]] bool is_prefix(const std::vector& pre, + const std::vector& full) { if (pre.size() > full.size()) return false; for (std::size_t i = 0; i < pre.size(); ++i) @@ -167,7 +167,7 @@ struct PlainTokenizer : nn::BPETokenizer return true; } -[[nodiscard]] bool contains_id(const std::vector &ids, std::size_t id) +[[nodiscard]] bool contains_id(const std::vector& ids, std::size_t id) { for (const auto v : ids) if (v == id) return true; @@ -175,8 +175,8 @@ struct PlainTokenizer : nn::BPETokenizer } // 从 CLI 输出里抽 "[Prompt tokens: a, b, c]" / "[Generated tokens: ...]" -[[nodiscard]] std::vector parse_id_list(const std::string &text, - const std::string &tag) +[[nodiscard]] std::vector parse_id_list(const std::string& text, + const std::string& tag) { const std::string key = "[" + tag + ": "; const auto pos = text.find(key); @@ -205,7 +205,7 @@ void test_dialogue_structure() auto tok_r = train_tokenizer(300); CHECK(tok_r.has_value(), "训练小 BPE 词表"); if (!tok_r) return; - nn::BPETokenizer &tok = *tok_r; + nn::BPETokenizer& tok = *tok_r; CHECK(tok.has_dialogue_markers(), "训练出的词表带对话标记"); std::vector turns{ @@ -218,9 +218,8 @@ void test_dialogue_structure() // 期望：[BOS]<|system|>SYS<|end_of_system|><|user|>hi<|end_of_user|> // <|assistant|>yo<|end_of_assistant|><|user|>bye<|end_of_user|><|assistant|> std::vector want; - const auto push_all = [&want](const std::vector &v) { - want.insert(want.end(), v.begin(), v.end()); - }; + const auto push_all = [&want](const std::vector& v) + { want.insert(want.end(), v.begin(), v.end()); }; want.push_back(tok.bos_id()); want.push_back(tok.system_marker_id()); push_all(tok.encode("SYS")); @@ -236,20 +235,20 @@ void test_dialogue_structure() want.push_back(tok.end_user_marker_id()); want.push_back(tok.assistant_marker_id()); // ★ 末轮 user → 开标记待生成 - CHECK(ids == want, ("模板逐 token 一致（实际: " + join_ids(ids) + - "；期望: " + join_ids(want) + "）").c_str()); + CHECK( + ids == want, + ("模板逐 token 一致（实际: " + join_ids(ids) + "；期望: " + join_ids(want) + "）").c_str()); CHECK(!ids.empty() && ids.back() == tok.assistant_marker_id(), "末轮是 user 时以 <|assistant|> 收尾（提示模型作答）"); // system_prompt 为空 → 不产生 system 块（BOS 与各轮次逐 token 不变） std::vector want_no_sys{tok.bos_id()}; // 去掉的正是 BOS 之后的 <|system|> + SYS + <|end_of_system|> 三段 - want_no_sys.insert(want_no_sys.end(), want.begin() + 3 + tok.encode("SYS").size(), - want.end()); + want_no_sys.insert(want_no_sys.end(), want.begin() + 3 + tok.encode("SYS").size(), want.end()); const auto no_sys = tok.chat_prompt(turns, ""); - CHECK(no_sys == want_no_sys, - ("空 system_prompt ⇒ 无 system 块（实际 " + join_ids(no_sys) + "；期望 " + - join_ids(want_no_sys) + "）").c_str()); + CHECK(no_sys == want_no_sys, ("空 system_prompt ⇒ 无 system 块（实际 " + join_ids(no_sys) + + "；期望 " + join_ids(want_no_sys) + "）") + .c_str()); // 末轮是 assistant（完整记录）→ 以 <|end_of_assistant|> 收尾 std::vector closed(turns.begin(), turns.begin() + 2); @@ -267,7 +266,7 @@ void test_context_accumulates() auto tok_r = train_tokenizer(300); CHECK(tok_r.has_value(), "训练小 BPE 词表"); if (!tok_r) return; - nn::BPETokenizer &tok = *tok_r; + nn::BPETokenizer& tok = *tok_r; // 模拟连续 5 轮的闭环（GUI / --interactive 都是这个形状）： // 每轮 append 一个 user 轮 → 生成 → append 一个 assistant 轮。 @@ -280,8 +279,8 @@ void test_context_accumulates() if (round > 1) { CHECK(is_prefix(prev_prompt, prompt), - ("第 " + std::to_string(round) + - " 轮的 prompt 必须以上一轮为前缀（上下文丢了）").c_str()); + ("第 " + std::to_string(round) + " 轮的 prompt 必须以上一轮为前缀（上下文丢了）") + .c_str()); CHECK(prompt.size() > prev_prompt.size(), ("第 " + std::to_string(round) + " 轮的 prompt 必须严格更长").c_str()); } @@ -308,11 +307,10 @@ void test_tokens_passthrough() auto tok_r = train_tokenizer(300); CHECK(tok_r.has_value(), "训练小 BPE 词表"); if (!tok_r) return; - nn::BPETokenizer &tok = *tok_r; + nn::BPETokenizer& tok = *tok_r; // 上一轮真实生成的 token（这里用一段合法的 id 序列代表） - const std::vector generated{tok.encode("yo").front(), - tok.encode(" bye").front()}; + const std::vector generated{tok.encode("yo").front(), tok.encode(" bye").front()}; std::vector turns{ {true, "hi", {}}, {false, "", generated}, // ★ tokens 字段优先于 text @@ -336,7 +334,7 @@ void test_truncation() auto tok_r = train_tokenizer(300); CHECK(tok_r.has_value(), "训练小 BPE 词表"); if (!tok_r) return; - nn::BPETokenizer &tok = *tok_r; + nn::BPETokenizer& tok = *tok_r; std::vector turns; for (int i = 1; i <= 6; ++i) @@ -349,10 +347,8 @@ void test_truncation() // 逐级收紧上限：丢弃轮数单调不减，且恒保留 BOS + system 块 + 末轮 const std::vector limits{ - std::max(1, full.size() - 1), - std::max(1, full.size() / 2), - std::max(1, full.size() / 4), - std::size_t{1}}; + std::max(1, full.size() - 1), std::max(1, full.size() / 2), + std::max(1, full.size() / 4), std::size_t{1}}; std::size_t prev_dropped = 0; for (const auto limit : limits) { @@ -364,8 +360,7 @@ void test_truncation() CHECK(!is_prefix(full, ids) || ids.size() <= full.size(), "截断只会变短"); // 恒保留：BOS + system 块 std::vector head{tok.bos_id(), tok.system_marker_id()}; - CHECK(ids.size() >= head.size() && - std::equal(head.begin(), head.end(), ids.begin()), + CHECK(ids.size() >= head.size() && std::equal(head.begin(), head.end(), ids.begin()), "BOS 与 system 块恒保留"); // 恒保留：末轮 ⇒ 仍以 <|assistant|> 收尾 CHECK(!ids.empty() && ids.back() == tok.assistant_marker_id(), @@ -382,8 +377,7 @@ void test_truncation() std::size_t dropped = 0; const auto ids = tok.chat_prompt(turns, "SYS", 1, &dropped); CHECK(dropped == 10, "6 轮（11 条轮次）时最多丢 10 条 = 全部非末轮"); - CHECK(!ids.empty() && ids.back() == tok.assistant_marker_id(), - "极限截断下末轮结构完整"); + CHECK(!ids.empty() && ids.back() == tok.assistant_marker_id(), "极限截断下末轮结构完整"); } // 只有一轮 ⇒ 无可丢 { @@ -391,8 +385,7 @@ void test_truncation() std::size_t dropped = 99; const auto ids = tok.chat_prompt(one, "SYS", 1, &dropped); CHECK(dropped == 0, "只有一轮时无可丢（末轮恒保留）"); - CHECK(!ids.empty() && ids.back() == tok.assistant_marker_id(), - "单轮极端截断下结构仍完整"); + CHECK(!ids.empty() && ids.back() == tok.assistant_marker_id(), "单轮极端截断下结构仍完整"); } } @@ -430,14 +423,14 @@ void test_plain_text_fallback() const auto t = plain.encode(turns[i].text); want.insert(want.end(), t.begin(), t.end()); } - CHECK(ids == want, ("纯文本：BOS + 轮次文本按 \\n 拼接（实际 " + join_ids(ids) + - "；期望 " + join_ids(want) + "）").c_str()); + CHECK(ids == want, ("纯文本：BOS + 轮次文本按 \\n 拼接（实际 " + join_ids(ids) + "；期望 " + + join_ids(want) + "）") + .c_str()); // 超限：截尾（保留最近 max_tokens 个 token） const auto cut = plain.chat_prompt(turns, "", 3); CHECK(cut.size() == 3, "纯文本超限截尾到 max_tokens"); - CHECK(std::equal(cut.begin(), cut.end(), ids.end() - 3), - "截尾保留的是最近的 token（尾部）"); + CHECK(std::equal(cut.begin(), cut.end(), ids.end() - 3), "截尾保留的是最近的 token（尾部）"); } // ══════════════════════════════════════════════════════════════════════════ @@ -450,7 +443,7 @@ struct CliFixture std::string vocab_path; }; -std::string shell_quote(const std::string &s) +std::string shell_quote(const std::string& s) { #ifdef _WIN32 // cmd 会剥掉"整条命令以引号开头"时的首尾引号；无空格路径直接裸写。 @@ -468,7 +461,7 @@ std::string shell_quote(const std::string &s) // 造一个可推理的 tiny GPT（词表 = 训练出的小 BPE）+ 词表文件； // 返回 false = 环境不具备（建模型失败等），调用方据此跳过。 -bool make_cli_fixture(const std::filesystem::path &dir, CliFixture &fx) +bool make_cli_fixture(const std::filesystem::path& dir, CliFixture& fx) { auto tok_r = train_tokenizer(300); CHECK(tok_r.has_value(), "训练小 BPE 词表"); @@ -487,14 +480,13 @@ bool make_cli_fixture(const std::filesystem::path &dir, CliFixture &fx) if (!model_r) return false; fx.model_path = (dir / "chat_model.nnmodel").string(); // 空嵌入词表 ⇒ CLI 走 --vocab（与 GUI 的"模型未嵌入词表"路径一致） - auto sr = nn::save_model(fx.model_path, *model_r, spec, std::string{}, - nn::ModelRole::Model); + auto sr = nn::save_model(fx.model_path, *model_r, spec, std::string{}, nn::ModelRole::Model); CHECK(sr.has_value(), "写出 .nnmodel"); return static_cast(sr); } -void test_cli_multi_turn(const std::string &exe, const std::filesystem::path &dir, - const CliFixture &fx) +void test_cli_multi_turn(const std::string& exe, const std::filesystem::path& dir, + const CliFixture& fx) { std::printf("── [6] CLI 端到端：--turn 多轮 + --reply-out ──\n"); const std::string out1 = (dir / "turn1.out").string(); @@ -506,14 +498,14 @@ void test_cli_multi_turn(const std::string &exe, const std::filesystem::path &di const std::string sys = "SYS"; const std::string base = shell_quote(exe) + " --model " + shell_quote(fx.model_path) + - " --vocab " + shell_quote(fx.vocab_path) + - " --system " + shell_quote(sys) + - " --max-tokens 6 --seed 11 --show-tokens"; + " --vocab " + shell_quote(fx.vocab_path) + " --system " + + shell_quote(sys) + " --max-tokens 6 --seed 11 --show-tokens"; // ── 单轮（无历史）── - const int rc1 = exit_code_of(std::system( - (base + " --turn user:\"hi there\" --reply-out " + shell_quote(reply1) + - " > " + shell_quote(out1) + " 2> " + shell_quote(err1)).c_str())); + const int rc1 = exit_code_of( + std::system((base + " --turn user:\"hi there\" --reply-out " + shell_quote(reply1) + " > " + + shell_quote(out1) + " 2> " + shell_quote(err1)) + .c_str())); CHECK(rc1 == 0, "单轮 --turn 正常退出"); const std::string log1 = read_file(out1); const auto prompt1 = parse_id_list(log1, "Prompt tokens"); @@ -529,23 +521,23 @@ void test_cli_multi_turn(const std::string &exe, const std::filesystem::path &di for (char& c : r1) { const auto u = static_cast(c); - if (u < 0x20 || u >= 0x7F || c == '"' || c == '&' || c == '|' || - c == '<' || c == '>') + if (u < 0x20 || u >= 0x7F || c == '"' || c == '&' || c == '|' || c == '<' || c == '>') c = ' '; } if (r1.find_first_not_of(' ') == std::string::npos) r1 = "ok"; - const int rc3 = exit_code_of(std::system( - (base + " --turn user:\"hi there\" --turn assistant:" + shell_quote(r1) + - " --turn user:\"bye now\" --reply-out " + shell_quote(reply3) + - " > " + shell_quote(out3) + " 2> " + shell_quote(err3)).c_str())); + const int rc3 = exit_code_of( + std::system((base + " --turn user:\"hi there\" --turn assistant:" + shell_quote(r1) + + " --turn user:\"bye now\" --reply-out " + shell_quote(reply3) + " > " + + shell_quote(out3) + " 2> " + shell_quote(err3)) + .c_str())); CHECK(rc3 == 0, "三轮 --turn 正常退出"); const std::string log3 = read_file(out3); const auto prompt3 = parse_id_list(log3, "Prompt tokens"); // ★ 上下文不丢（端到端）：单轮 prompt 是三轮 prompt 的前缀，且更短 - CHECK(is_prefix(prompt1, prompt3), - ("CLI 端到端：单轮 prompt 应为三轮 prompt 的前缀（单轮 " + join_ids(prompt1) + - " / 三轮 " + join_ids(prompt3) + "）").c_str()); + CHECK(is_prefix(prompt1, prompt3), ("CLI 端到端：单轮 prompt 应为三轮 prompt 的前缀（单轮 " + + join_ids(prompt1) + " / 三轮 " + join_ids(prompt3) + "）") + .c_str()); CHECK(prompt3.size() > prompt1.size(), "三轮 prompt 严格更长（历史真的进了模型输入）"); // 与库侧同一份 turns 逐 token 一致（接线一致，而不只是"更长"） @@ -555,25 +547,29 @@ void test_cli_multi_turn(const std::string &exe, const std::filesystem::path &di {true, "bye now", {}}, }; const auto lib_prompt = fx.tok.chat_prompt(turns, sys, 64); - CHECK(lib_prompt == prompt3, - ("CLI 的 prompt tokens 与库侧 chat_prompt 逐 token 一致（CLI " + - join_ids(prompt3) + " / 库 " + join_ids(lib_prompt) + "）").c_str()); + CHECK(lib_prompt == prompt3, ("CLI 的 prompt tokens 与库侧 chat_prompt 逐 token 一致（CLI " + + join_ids(prompt3) + " / 库 " + join_ids(lib_prompt) + "）") + .c_str()); // ── --reply-out 落盘 = 生成 token 的解码（截到 <|end_of_assistant|>）── const auto gen3 = parse_id_list(log3, "Generated tokens"); CHECK(!gen3.empty(), "三轮输出里有 Generated tokens 列表"); std::vector shown = gen3; for (std::size_t i = 0; i < shown.size(); ++i) - if (shown[i] == fx.tok.end_assistant_marker_id()) { shown.resize(i); break; } + if (shown[i] == fx.tok.end_assistant_marker_id()) + { + shown.resize(i); + break; + } const std::string expect_reply = fx.tok.decode(shown); const std::string got_reply = rstrip_newline(read_file(reply3)); - CHECK(got_reply == expect_reply, - ("--reply-out 内容 = 生成 token 解码：实际 '" + got_reply + "' 期望 '" + - expect_reply + "'").c_str()); + CHECK(got_reply == expect_reply, ("--reply-out 内容 = 生成 token 解码：实际 '" + got_reply + + "' 期望 '" + expect_reply + "'") + .c_str()); } -void test_cli_rejections(const std::string &exe, const std::filesystem::path &dir, - const CliFixture &fx) +void test_cli_rejections(const std::string& exe, const std::filesystem::path& dir, + const CliFixture& fx) { std::printf("── [7] CLI 负例：--turn 形状 / 互斥项不合法即报错 ──\n"); const std::string out = (dir / "neg.out").string(); @@ -582,9 +578,9 @@ void test_cli_rejections(const std::string &exe, const std::filesystem::path &di " --vocab " + shell_quote(fx.vocab_path) + " --max-tokens 2"; struct Case { - const char *args; - const char *label; - const char *needle; + const char* args; + const char* label; + const char* needle; }; const Case cases[] = { {"--turn user:\"hi\" --prompt \"x\"", "--turn 与 --prompt 互斥", "互斥"}, @@ -595,26 +591,23 @@ void test_cli_rejections(const std::string &exe, const std::filesystem::path &di {"--turn user", "缺少冒号", ":"}, {"--turn bot:x", "role 非法", ":"}, }; - for (const auto &c : cases) + for (const auto& c : cases) { - const std::string cmd = base + " " + c.args + " > " + shell_quote(out) + " 2> " + - shell_quote(err); + const std::string cmd = + base + " " + c.args + " > " + shell_quote(out) + " 2> " + shell_quote(err); const int rc = exit_code_of(std::system(cmd.c_str())); CHECK(rc != 0, (std::string("应非零退出: ") + c.label).c_str()); const std::string errtxt = read_file(err); CHECK(errtxt.find(c.needle) != std::string::npos, - (std::string("错误信息应说明原因（含 '") + c.needle + "'）: " + c.label) - .c_str()); + (std::string("错误信息应说明原因（含 '") + c.needle + "'）: " + c.label).c_str()); } } } // namespace -int main(int argc, char **argv) +int main(int argc, char** argv) { - g_self = (argc > 0 && argv[0] != nullptr) - ? std::filesystem::absolute(argv[0]).string() - : ""; + g_self = (argc > 0 && argv[0] != nullptr) ? std::filesystem::absolute(argv[0]).string() : ""; std::printf("chat_prompt_test（P0-5 尾项：多轮对话闭环）\n"); test_dialogue_structure(); @@ -628,8 +621,7 @@ int main(int argc, char **argv) { std::printf("\n[skip] 未找到 text_infer（未构建 app / NN_BUILD_APPS=OFF）：" "CLI 半边跳过\n"); - std::printf("\n%s 失败=%d（CLI 半边跳过 = 77）\n", - g_fail == 0 ? "PASS" : "FAIL", g_fail); + std::printf("\n%s 失败=%d（CLI 半边跳过 = 77）\n", g_fail == 0 ? "PASS" : "FAIL", g_fail); return g_fail == 0 ? kSkip : 1; } const auto dir = scratch_dir(); diff --git a/src/cli_precision_warning_test.cpp b/src/cli_precision_warning_test.cpp index 443f8824..699e4aae 100644 --- a/src/cli_precision_warning_test.cpp +++ b/src/cli_precision_warning_test.cpp @@ -71,8 +71,7 @@ bool file_exists(const std::string& p) std::string read_file(const std::string& p) { std::ifstream f(p, std::ios::binary); - return std::string((std::istreambuf_iterator(f)), - std::istreambuf_iterator()); + return std::string((std::istreambuf_iterator(f)), std::istreambuf_iterator()); } // 运行 cmd，stderr 落到 err_file；返回进程退出码（-1 = 无法启动子进程）。 @@ -117,20 +116,17 @@ std::string find_app(const std::string& name) void test_helper_matrix() { using P = nn::Precision; - const auto has = [](nn::PrecisionProfile p) { - return !nn::cli::dangerous_precision_warning(p).empty(); - }; + const auto has = [](nn::PrecisionProfile p) + { return !nn::cli::dangerous_precision_warning(p).empty(); }; // 静默：f32 / --f16（profile_f16）/ --t1_58（profile_t1_58） - CHECK(nn::cli::dangerous_precision_warning(nn::profile_f32()).empty(), - "全 f32 应静默"); + CHECK(nn::cli::dangerous_precision_warning(nn::profile_f32()).empty(), "全 f32 应静默"); CHECK(nn::cli::dangerous_precision_warning(nn::profile_f16()).empty(), "--f16（profile_f16）应静默"); CHECK(nn::cli::dangerous_precision_warning(nn::profile_t1_58()).empty(), "--t1_58（profile_t1_58）应静默"); // stable=f16 单独出现不触发（24 §2.3 口径） - CHECK(!has({P::F32, P::F32, P::F16, P::F32}), - "stable=f16 单独出现不触发告警"); + CHECK(!has({P::F32, P::F32, P::F16, P::F32}), "stable=f16 单独出现不触发告警"); // param=t1_58（非 F16 存储）+ compute=f16 不触发 CHECK(!has({P::T1_58, P::F16, P::F32, P::F32}), "param=t1_58 / compute=f16（optimizer=f32）应静默"); @@ -147,8 +143,7 @@ void test_helper_matrix() "非全 f16 不应出现'数值性冻结'（那是全 f16 的形态）"); // 触发：四字段全 f16（★ F2 补进的缺口组合）→ 合并一条更完整的警告 - const std::string all16 = - nn::cli::dangerous_precision_warning(nn::profile_all_f16()); + const std::string all16 = nn::cli::dangerous_precision_warning(nn::profile_all_f16()); CHECK(!all16.empty(), "四字段全 f16 必须触发告警"); CHECK(all16.find("optimizer=f16") != std::string::npos && all16.find("stable=f16") != std::string::npos, @@ -173,10 +168,8 @@ void test_cli_stderr(const std::string& exe, const std::string& label, const std::string err_opt = "nn_precwarn_opt.err"; const std::string err_all = "nn_precwarn_all.err"; - const int rc_plain = run_capture_stderr( - quote_exe(exe) + " " + dataset_arg, err_plain); - const int rc_f16 = run_capture_stderr( - quote_exe(exe) + " --f16 " + dataset_arg, err_f16); + const int rc_plain = run_capture_stderr(quote_exe(exe) + " " + dataset_arg, err_plain); + const int rc_f16 = run_capture_stderr(quote_exe(exe) + " --f16 " + dataset_arg, err_f16); const int rc_opt = run_capture_stderr( quote_exe(exe) + " --precision-stable f16 --precision-optimizer f16 " + dataset_arg, err_opt); @@ -218,18 +211,20 @@ void test_cli_stderr(const std::string& exe, const std::string& label, int main(int argc, char** argv) { - g_self = (argc > 0 && argv[0] != nullptr) - ? std::filesystem::absolute(argv[0]).string() - : ""; + g_self = (argc > 0 && argv[0] != nullptr) ? std::filesystem::absolute(argv[0]).string() : ""; std::printf("cli_precision_warning_test（F1：危险精度组合启动期告警）\n"); test_helper_matrix(); bool any_app = false; // text_train 数据集是位置参数；mnist_train 用 --dataset <目录>。 - struct AppCase { const char* name; const char* dataset_arg; }; + struct AppCase + { + const char* name; + const char* dataset_arg; + }; const AppCase apps[] = { - {"text_train", "./nn_does_not_exist_dir_xyz/missing.nndataset"}, + {"text_train", "./nn_does_not_exist_dir_xyz/missing.nndataset"}, {"mnist_train", "--dataset ./nn_does_not_exist_dir_xyz"}, }; for (const auto& app : apps) diff --git a/src/cnn_smoke_test.cpp b/src/cnn_smoke_test.cpp index f4541f29..2deed031 100644 --- a/src/cnn_smoke_test.cpp +++ b/src/cnn_smoke_test.cpp @@ -20,7 +20,8 @@ #include #include -namespace { +namespace +{ bool cnn_all_finite(const nn::Matrix& m) { @@ -41,7 +42,8 @@ int main(int argc, char* argv[]) const std::string a = argv[i]; if (a == "--gpu") use_gpu = true; else if (a == "--f16") use_f16 = true; - else if (a == "--steps" && i + 1 < argc) steps = static_cast(std::atoi(argv[++i])); + else if (a == "--steps" && i + 1 < argc) + steps = static_cast(std::atoi(argv[++i])); else if (a == "--help") { std::cout << "用法: cnn_smoke_test [--gpu] [--f16] [--steps N]\n"; @@ -61,10 +63,10 @@ int main(int argc, char* argv[]) // 输入 12×12：conv3 → 10 → pool2 → 5 → conv3 → 3 → pool2 → 1（每层卷积后都池化） nn::CnnConfig cfg; cfg.in_channels = 1; - cfg.in_size = 12; - cfg.pool = 2; - cfg.convs = {{4, 3, 1, 0}, {8, 3, 1, 0}}; // {out_channels, kernel, stride, padding} - cfg.fc_dims = {16, 10}; // flatten(8) → 16 → 10 + cfg.in_size = 12; + cfg.pool = 2; + cfg.convs = {{4, 3, 1, 0}, {8, 3, 1, 0}}; // {out_channels, kernel, stride, padding} + cfg.fc_dims = {16, 10}; // flatten(8) → 16 → 10 // P2-3：`--f16` 走 profile_f16()（Conv2D 权重/梯度按 param=f16 建， // 融合算子按 compute=f16） if (use_f16) @@ -84,11 +86,9 @@ int main(int argc, char* argv[]) return 1; } auto cfg2 = nn::cnn_config_from_spec(spec); - if (!cfg2 || - cfg2->in_channels != cfg.in_channels || cfg2->in_size != cfg.in_size || + if (!cfg2 || cfg2->in_channels != cfg.in_channels || cfg2->in_size != cfg.in_size || cfg2->pool != cfg.pool || cfg2->convs.size() != cfg.convs.size() || - cfg2->convs[1].out_channels != cfg.convs[1].out_channels || - cfg2->fc_dims != cfg.fc_dims) + cfg2->convs[1].out_channels != cfg.convs[1].out_channels || cfg2->fc_dims != cfg.fc_dims) { std::cerr << "cnn_config_from_spec 往返不一致\n"; return 1; @@ -100,14 +100,15 @@ int main(int argc, char* argv[]) for (std::size_t i = 0; i < model.num_layers(); ++i) { nn::Layer& l = model.layer_at(i); - if (dynamic_cast(&l)) ++n_conv; + if (dynamic_cast(&l)) ++n_conv; if (dynamic_cast(&l)) ++n_pool; - if (dynamic_cast(&l)) ++n_linear; + if (dynamic_cast(&l)) ++n_linear; if (dynamic_cast(&l) || dynamic_cast(&l) || - dynamic_cast(&l)) ++n_norm; + dynamic_cast(&l)) + ++n_norm; } - std::printf("layers=%zu Conv2D=%zu MaxPool2D=%zu Linear=%zu Norm=%zu\n", - model.num_layers(), n_conv, n_pool, n_linear, n_norm); + std::printf("layers=%zu Conv2D=%zu MaxPool2D=%zu Linear=%zu Norm=%zu\n", model.num_layers(), + n_conv, n_pool, n_linear, n_norm); if (n_conv != cfg.convs.size() || n_pool != cfg.convs.size() || n_linear < 2) { std::cerr << "层组成与配置不符\n"; @@ -124,12 +125,11 @@ int main(int argc, char* argv[]) // ── 2b. 归一化挂载（NormPlace）：both → 卷积后各 1 + FC 隐藏层 1 ─── { nn::CnnConfig ncfg = cfg; - ncfg.norm_type = nn::NormType::RMSNorm; + ncfg.norm_type = nn::NormType::RMSNorm; ncfg.norm_place = nn::NormPlace::Both; - auto nspec = nn::make_cnn_spec(ncfg.in_channels, ncfg.in_size, ncfg.pool, - ncfg.convs, ncfg.fc_dims, - ncfg.norm_type, ncfg.norm_place); + auto nspec = nn::make_cnn_spec(ncfg.in_channels, ncfg.in_size, ncfg.pool, ncfg.convs, + ncfg.fc_dims, ncfg.norm_type, ncfg.norm_place); auto nmodel_r = nn::build_cnn_model_from_spec(eng, nspec); NN_EXIT(nmodel_r, 1, "norm CNN 构建失败: "); nn::Model& nmodel = *nmodel_r; @@ -157,7 +157,8 @@ int main(int argc, char* argv[]) nn::Matrix nxm(ncfg.in_channels * ncfg.in_size * ncfg.in_size, nbatch); std::mt19937_64 nrng{20261003}; std::uniform_real_distribution ndist(-1.0f, 1.0f); - for (std::size_t i = 0; i < nxm.size(); ++i) nxm.span()[i] = ndist(nrng); + for (std::size_t i = 0; i < nxm.size(); ++i) + nxm.span()[i] = ndist(nrng); std::vector nlabels(nbatch); for (std::size_t i = 0; i < nbatch; ++i) nlabels[i] = static_cast(nrng() % ncfg.fc_dims.back()); @@ -167,8 +168,7 @@ int main(int argc, char* argv[]) auto nlogits = nmodel.forward(*nx); NN_EXIT(nlogits, 1, "norm forward 失败: "); nn::CrossEntropyLoss nce; - auto nloss = nce.forward_sparse(eng, *nlogits, nlabels, {}, - ncfg.fc_dims.back()); + auto nloss = nce.forward_sparse(eng, *nlogits, nlabels, {}, ncfg.fc_dims.back()); NN_EXIT(nloss, 1, "norm loss 失败\n"); auto ngrad = nce.backward(); NN_EXIT(ngrad, 1, "norm CE backward 失败\n"); @@ -178,9 +178,17 @@ int main(int argc, char* argv[]) for (auto& g : nmodel.param_gradients()) { auto gm = eng.to_matrix(g.get()); - if (!gm) { finite = false; break; } + if (!gm) + { + finite = false; + break; + } for (auto v : gm->span()) - if (!std::isfinite(v)) { finite = false; break; } + if (!std::isfinite(v)) + { + finite = false; + break; + } if (!finite) break; } NN_EXIT(finite, 1, "norm CNN 梯度含 NaN/Inf\n"); @@ -194,20 +202,21 @@ int main(int argc, char* argv[]) std::uniform_real_distribution dist(-1.0f, 1.0f); nn::Matrix xm(cfg.in_channels * cfg.in_size * cfg.in_size, batch); - for (std::size_t i = 0; i < xm.size(); ++i) xm.span()[i] = dist(rng); + for (std::size_t i = 0; i < xm.size(); ++i) + xm.span()[i] = dist(rng); std::vector labels(batch); - for (std::size_t i = 0; i < batch; ++i) labels[i] = static_cast(rng() % classes); + for (std::size_t i = 0; i < batch; ++i) + labels[i] = static_cast(rng() % classes); auto x = eng.from_matrix(xm); NN_EXIT(x, 1, "from_matrix 失败\n"); - auto optimizer = nn::create_optimizer( - "sgd", eng, model.parameters(), model.param_gradients(), nn::Scalar{0.1}); + auto optimizer = nn::create_optimizer("sgd", eng, model.parameters(), model.param_gradients(), + nn::Scalar{0.1}); NN_EXIT(optimizer, 1, "optimizer 创建失败\n"); nn::CrossEntropyLoss ce; - if (use_f16) - ce.set_precision_profile(nn::profile_f16()); + if (use_f16) ce.set_precision_profile(nn::profile_f16()); nn::Scalar first_loss = 0, last_loss = 0; for (std::size_t step = 0; step < steps; ++step) { diff --git a/src/conv2d_gradcheck.cpp b/src/conv2d_gradcheck.cpp index 118295ca..14fc7a24 100644 --- a/src/conv2d_gradcheck.cpp +++ b/src/conv2d_gradcheck.cpp @@ -20,13 +20,14 @@ #include #include -using nn::Scalar; -using nn::Matrix; -using nn::Tensor; using nn::ComputeEngine; using nn::Conv2D; +using nn::Matrix; +using nn::Scalar; +using nn::Tensor; -namespace { +namespace +{ // 独立参考：直接卷积 + im2col 索引（不依赖层内任何实现） struct ConvRef @@ -40,8 +41,8 @@ struct ConvRef { if (ih < 0 || iw < 0 || ih >= static_cast(H) || iw >= static_cast(W)) return Scalar{0}; - return x.at_unchecked(ci * H * W + static_cast(ih) * W + - static_cast(iw), bi); + return x.at_unchecked( + ci * H * W + static_cast(ih) * W + static_cast(iw), bi); } // forward 输出布局 (C_out*OH*OW, B) Matrix fwd() const @@ -56,8 +57,8 @@ struct ConvRef for (std::size_t ci = 0; ci < C_in; ++ci) for (std::size_t kh = 0; kh < K; ++kh) for (std::size_t kw = 0; kw < K; ++kw) - s += w.at_unchecked(co, ci * K * K + kh * K + kw) - * xv(ci, static_cast(oh) + static_cast(kh), + s += w.at_unchecked(co, ci * K * K + kh * K + kw) * + xv(ci, static_cast(oh) + static_cast(kh), static_cast(ow) + static_cast(kw), bi); out.set_value_unchecked(co * OH * OW + oh * OW + ow, bi, s); } @@ -130,9 +131,11 @@ int main(int argc, char* argv[]) std::mt19937_64 rng(20260918); std::uniform_real_distribution dist(-0.5f, 0.5f); - const auto fill = [&](std::size_t r, std::size_t c) { + const auto fill = [&](std::size_t r, std::size_t c) + { Matrix m(r, c); - for (std::size_t i = 0; i < m.size(); ++i) m.span()[i] = dist(rng); + for (std::size_t i = 0; i < m.size(); ++i) + m.span()[i] = dist(rng); return m; }; @@ -151,7 +154,8 @@ int main(int argc, char* argv[]) ConvRef ref{C_in, C_out, K, H, W, OH, OW, B, fill(C_in * H * W, B), wm, bm}; int failures = 0; - const auto report = [&](const char* what, const Matrix& got, const Matrix& want) { + const auto report = [&](const char* what, const Matrix& got, const Matrix& want) + { const Scalar e = max_abs_diff(got, want); const bool ok = e < tol; std::printf(" %-14s %s max_abs_diff=%.3e\n", what, ok ? "OK" : "FAIL", diff --git a/src/cross_attention_test.cpp b/src/cross_attention_test.cpp index e967425d..2cc71f1a 100644 --- a/src/cross_attention_test.cpp +++ b/src/cross_attention_test.cpp @@ -19,11 +19,11 @@ int g_fail = 0; -namespace { +namespace +{ // 逐位比对（容差 1e-4；两侧同引擎同精度） -void compare(nn::ComputeEngine& e, const nn::Tensor& got, const nn::Tensor& ref, - const char* msg) +void compare(nn::ComputeEngine& e, const nn::Tensor& got, const nn::Tensor& ref, const char* msg) { CHECK(got.rows() == ref.rows() && got.cols() == ref.cols(), msg); if (got.rows() != ref.rows() || got.cols() != ref.cols()) return; @@ -89,10 +89,10 @@ void test_mha_consistency() CHECK(gpair.has_value(), "[1] backward_pair"); if (!gmha || !gpair) return; // grad_query + grad_memory == MHA 的 grad_input（同一投影拆分的两半） - CHECK(gpair->grad_query.rows() == d_model && - gpair->grad_query.cols() == batch * seq, "[1] grad_query 形状"); - CHECK(gpair->grad_memory.rows() == d_model && - gpair->grad_memory.cols() == batch * seq, "[1] grad_memory 形状"); + CHECK(gpair->grad_query.rows() == d_model && gpair->grad_query.cols() == batch * seq, + "[1] grad_query 形状"); + CHECK(gpair->grad_memory.rows() == d_model && gpair->grad_memory.cols() == batch * seq, + "[1] grad_memory 形状"); auto gsum = download(*engine2, gpair->grad_query); auto gmem = download(*engine2, gpair->grad_memory); auto gs = gsum.span(); @@ -104,7 +104,8 @@ void test_mha_consistency() if (std::fabs(gs[i] + gms[i] - rs[i]) > 1e-4f) { std::printf("[FAIL] [1] grad_query+grad_memory == MHA grad_input: " - "idx=%zu got=%g ref=%g\n", i, gs[i] + gms[i], rs[i]); + "idx=%zu got=%g ref=%g\n", + i, gs[i] + gms[i], rs[i]); ++g_fail; break; } @@ -127,7 +128,8 @@ void test_rectangular() "[2] 输出形状 (d_model, batch*seq_q)"); const auto om = download(*engine, *out); bool finite = true; - for (auto v : om.span()) if (!std::isfinite(v)) finite = false; + for (auto v : om.span()) + if (!std::isfinite(v)) finite = false; CHECK(finite, "[2] 输出有限"); auto grad = make_tensor(*engine, d_model, batch * seq_q, 0.02f, 0.01f); @@ -136,11 +138,12 @@ void test_rectangular() if (!g) return; CHECK(g->grad_query.rows() == d_model && g->grad_query.cols() == batch * seq_q, "[2] grad_query 形状"); - CHECK(g->grad_memory.rows() == d_model && - g->grad_memory.cols() == batch * seq_kv, "[2] grad_memory 形状"); + CHECK(g->grad_memory.rows() == d_model && g->grad_memory.cols() == batch * seq_kv, + "[2] grad_memory 形状"); const auto grm = download(*engine, g->grad_memory); finite = true; - for (auto v : grm.span()) if (!std::isfinite(v)) finite = false; + for (auto v : grm.span()) + if (!std::isfinite(v)) finite = false; CHECK(finite, "[2] grad_memory 有限"); } @@ -197,8 +200,7 @@ void test_gradcheck() auto pm = download(*engine, p.get()); auto sp = pm.span(); for (std::size_t i = 0; i < sp.size(); ++i, ++k) - sp[i] = 0.05f * std::sin(0.7f * static_cast( - (k * 17) % 23 + 1)); + sp[i] = 0.05f * std::sin(0.7f * static_cast((k * 17) % 23 + 1)); auto cp = engine->copy_from(p.get(), pm); CHECK(cp.has_value(), "[5] 权重回写"); } @@ -210,20 +212,26 @@ void test_gradcheck() for (std::size_t i = 0; i < y_n; ++i) W[i] = 0.01f * static_cast(i % 7 + 1); - auto eval = [&]() -> nn::Scalar { + auto eval = [&]() -> nn::Scalar + { auto y = xattn.forward_pair(q, m); NN_ASSERT_ALWAYS(y, "gradcheck eval: forward_pair failed"); auto ym = download(*engine, *y); const auto sp = ym.span(); nn::Scalar loss = 0; - for (std::size_t i = 0; i < sp.size(); ++i) loss += sp[i] * W[i]; + for (std::size_t i = 0; i < sp.size(); ++i) + loss += sp[i] * W[i]; return loss; }; // 解析梯度：go 第 i 列 = W 第 i 列 → backward_pair 的两路梯度即 // d/dq 与 d/dm（W 由 go 的列主序承载） nn::Matrix go_m(d_model, batch * seq_q); - { auto sp = go_m.span(); for (std::size_t i = 0; i < y_n; ++i) sp[i] = W[i]; } + { + auto sp = go_m.span(); + for (std::size_t i = 0; i < y_n; ++i) + sp[i] = W[i]; + } auto go = upload(*engine, go_m); (void)eval(); // backward_pair 需要 forward_pair 的缓存 auto g = xattn.backward_pair(go); @@ -237,7 +245,8 @@ void test_gradcheck() auto ga_m = download(*engine, g->grad_memory); // in-place 扰动辅助：改 t 的第 idx 个标量（经 copy_from 写回） - auto perturb = [&](nn::Tensor& t, std::size_t idx, nn::Scalar v) { + auto perturb = [&](nn::Tensor& t, std::size_t idx, nn::Scalar v) + { auto tm = download(*engine, t); tm.span()[idx] = v; auto cp = engine->copy_from(t, tm); @@ -265,8 +274,7 @@ void test_gradcheck() const nn::Scalar ana = ga_q.span()[idx]; if (std::fabs(num - ana) > tol) { - std::printf("[FAIL] [5] grad_query idx=%zu num=%g ana=%g\n", - idx, num, ana); + std::printf("[FAIL] [5] grad_query idx=%zu num=%g ana=%g\n", idx, num, ana); ++bad_q; } } @@ -286,8 +294,7 @@ void test_gradcheck() const nn::Scalar ana = ga_m.span()[idx]; if (std::fabs(num - ana) > tol) { - std::printf("[FAIL] [5] grad_memory idx=%zu num=%g ana=%g\n", - idx, num, ana); + std::printf("[FAIL] [5] grad_memory idx=%zu num=%g ana=%g\n", idx, num, ana); ++bad_m; } } diff --git a/src/dataset_convert.cpp b/src/dataset_convert.cpp index 180dea3a..57799c8a 100644 --- a/src/dataset_convert.cpp +++ b/src/dataset_convert.cpp @@ -20,11 +20,12 @@ #include // ==================== 帮助信息 ==================== -void print_usage(const char *prog) +void print_usage(const char* prog) { nn::cli::Help help(std::cout, prog, "旧格式转换器（词表/数据集迁移）"); - help.usage("vocab -o | csv [] -o "); + help.usage( + "vocab -o | csv [] -o "); help.section("子命令"); help.item("vocab", "旧 JSON 词表 → .nnvocab（kvrec v2）"); @@ -44,7 +45,7 @@ struct Config std::string output_path; }; -Config parse_args(int argc, char *argv[]) +Config parse_args(int argc, char* argv[]) { Config cfg; for (int i = 1; i < argc; ++i) @@ -55,16 +56,12 @@ Config parse_args(int argc, char *argv[]) print_usage(argv[0]); std::exit(0); } - else if ((arg == "-o" || arg == "--output") && i + 1 < argc) - cfg.output_path = argv[++i]; + else if ((arg == "-o" || arg == "--output") && i + 1 < argc) cfg.output_path = argv[++i]; else if (!arg.starts_with("--")) { - if (arg == "vocab" || arg == "csv") - cfg.subcommand = arg; - else if (cfg.input1.empty()) - cfg.input1 = arg; - else if (cfg.input2.empty()) - cfg.input2 = arg; + if (arg == "vocab" || arg == "csv") cfg.subcommand = arg; + else if (cfg.input1.empty()) cfg.input1 = arg; + else if (cfg.input2.empty()) cfg.input2 = arg; else { std::cerr << "多余的位置参数: " << arg << "\n使用 --help 查看用法\n"; @@ -96,7 +93,7 @@ Config parse_args(int argc, char *argv[]) } // ==================== vocab：旧 JSON 词表 → .nnvocab ==================== -int convert_vocab(const Config &cfg) +int convert_vocab(const Config& cfg) { std::ifstream ifs(cfg.input1, std::ios::binary); if (!ifs) @@ -104,8 +101,7 @@ int convert_vocab(const Config &cfg) std::cerr << "无法打开输入文件: " << cfg.input1 << "\n"; return 1; } - const std::string json((std::istreambuf_iterator(ifs)), - std::istreambuf_iterator()); + const std::string json((std::istreambuf_iterator(ifs)), std::istreambuf_iterator()); if (json.empty() || json[0] != '{') { std::cerr << "输入不是 JSON 词表（应以 '{' 开头）: " << cfg.input1 << "\n"; @@ -116,20 +112,19 @@ int convert_vocab(const Config &cfg) auto tok = nn::load_tokenizer_from_string(json); if (!tok) { - std::cerr << "无法解析 JSON 词表（缺少有效 \"type\" 字段）: " - << cfg.input1 << "\n"; + std::cerr << "无法解析 JSON 词表（缺少有效 \"type\" 字段）: " << cfg.input1 << "\n"; return 1; } auto save_r = tok->save(cfg.output_path); NN_EXIT(save_r, 1, "写入 .nnvocab 失败: "); - std::cout << "词表已转换: " << cfg.input1 << " -> " << cfg.output_path - << "（" << tok->vocab_size() << " 词）" << std::endl; + std::cout << "词表已转换: " << cfg.input1 << " -> " << cfg.output_path << "（" + << tok->vocab_size() << " 词）" << std::endl; return 0; } // ==================== csv：预留（阶段四） ==================== -int convert_csv(const Config &) +int convert_csv(const Config&) { std::cerr << "dataset_convert csv 尚未实现（tabular 形态为阶段四预留，" "见 docs/internal/19-unified-dataset.md §4.4）\n"; @@ -137,10 +132,9 @@ int convert_csv(const Config &) } // ==================== 主函数 ==================== -int main(int argc, char *argv[]) +int main(int argc, char* argv[]) { Config cfg = parse_args(argc, argv); - if (cfg.subcommand == "vocab") - return convert_vocab(cfg); + if (cfg.subcommand == "vocab") return convert_vocab(cfg); return convert_csv(cfg); } diff --git a/src/dataset_gen.cpp b/src/dataset_gen.cpp index d70f93f9..2134245c 100644 --- a/src/dataset_gen.cpp +++ b/src/dataset_gen.cpp @@ -23,7 +23,7 @@ #include // ==================== 帮助信息 ==================== -void print_usage(const char *prog) +void print_usage(const char* prog) { nn::cli::Help help(std::cout, prog, "文本数据集生成器（文本 + 词表 → .nndataset）"); @@ -33,10 +33,15 @@ void print_usage(const char *prog) help.item("", "训练文本文件路径 (UTF-8，每行 = 一篇 doc)"); help.section("选项"); - help.opt("--vocab ", "词表文件路径 (.nnvocab，必需)\n词表从配套嵌入数据集，故必须先有词表 (tokenizer_train 先行)"); + help.opt("--vocab ", + "词表文件路径 (.nnvocab，必需)\n词表从配套嵌入数据集，故必须先有词表 (tokenizer_train " + "先行)"); help.opt("-o, --output ", "输出 .nndataset 路径 (必需)"); help.opt("--test ", "测试集文本 (可选第二输入 → test 子集，结构与 train 相同)"); - help.opt("--loss-scope ", "loss 计算范围: all / assistant (默认: all)\nall: 全部有效 token 参与 loss (预训练，不写 mask 段)\nassistant: 仅 <|assistant|>…<|end_of_assistant|> 段参与 (对话 SFT；\n逐 doc 写入 mask 段，语义在生成期决定，训练端不再选)"); + help.opt("--loss-scope ", + "loss 计算范围: all / assistant (默认: all)\nall: 全部有效 token 参与 loss " + "(预训练，不写 mask 段)\nassistant: 仅 <|assistant|>…<|end_of_assistant|> 段参与 " + "(对话 SFT；\n逐 doc 写入 mask 段，语义在生成期决定，训练端不再选)"); help.opt("--license ", "许可文本文件 (可选；全文原样写入配套 license 字段)"); help.opt("--source ", "来源名称 (默认: 训练文本文件路径)"); help.opt("--help, -h", "显示此帮助信息"); @@ -54,7 +59,7 @@ struct Config std::string source_name; }; -Config parse_args(int argc, char *argv[]) +Config parse_args(int argc, char* argv[]) { Config cfg; for (int i = 1; i < argc; ++i) @@ -65,28 +70,21 @@ Config parse_args(int argc, char *argv[]) print_usage(argv[0]); std::exit(0); } - else if (arg == "--vocab" && i + 1 < argc) - cfg.vocab_path = argv[++i]; - else if ((arg == "-o" || arg == "--output") && i + 1 < argc) - cfg.output_path = argv[++i]; - else if (arg == "--test" && i + 1 < argc) - cfg.test_path = argv[++i]; + else if (arg == "--vocab" && i + 1 < argc) cfg.vocab_path = argv[++i]; + else if ((arg == "-o" || arg == "--output") && i + 1 < argc) cfg.output_path = argv[++i]; + else if (arg == "--test" && i + 1 < argc) cfg.test_path = argv[++i]; else if (arg == "--loss-scope" && i + 1 < argc) { cfg.loss_scope = argv[++i]; if (cfg.loss_scope != "all" && cfg.loss_scope != "assistant") { - std::cerr << "未知 --loss-scope: " << cfg.loss_scope - << "，可选: all, assistant\n"; + std::cerr << "未知 --loss-scope: " << cfg.loss_scope << "，可选: all, assistant\n"; std::exit(1); } } - else if (arg == "--license" && i + 1 < argc) - cfg.license_path = argv[++i]; - else if (arg == "--source" && i + 1 < argc) - cfg.source_name = argv[++i]; - else if (!arg.starts_with("--")) - cfg.text_path = arg; + else if (arg == "--license" && i + 1 < argc) cfg.license_path = argv[++i]; + else if (arg == "--source" && i + 1 < argc) cfg.source_name = argv[++i]; + else if (!arg.starts_with("--")) cfg.text_path = arg; else { std::cerr << "未知参数: " << arg << "\n使用 --help 查看用法\n"; @@ -119,30 +117,29 @@ struct SubsetBuild std::size_t asst_tokens = 0; }; -nn::Result build_subset(const nn::Tokenizer &tokenizer, - const std::vector &docs, - bool need_mask, std::string_view name) +nn::Result build_subset(const nn::Tokenizer& tokenizer, + const std::vector& docs, bool need_mask, + std::string_view name) { SubsetBuild out; out.docs = nn::encode_docs_parallel(tokenizer, docs, /*show_progress=*/true, std::string("Tokenize ") + std::string(name)); - if (!need_mask) - return out; - for (const auto &toks : out.docs) + if (!need_mask) return out; + for (const auto& toks : out.docs) { - out.masks.push_back(nn::mark_assistant_span( - toks, tokenizer.assistant_marker_id(), tokenizer.end_assistant_marker_id())); + out.masks.push_back(nn::mark_assistant_span(toks, tokenizer.assistant_marker_id(), + tokenizer.end_assistant_marker_id())); for (unsigned char m : out.masks.back()) out.asst_tokens += m; } if (out.asst_tokens == 0) NN_FAIL(std::string(name) + "中未找到 <|assistant|> 段（--loss-scope assistant " - "要求语料含对话标记）"); + "要求语料含对话标记）"); return out; } // ==================== 主函数 ==================== -int main(int argc, char *argv[]) +int main(int argc, char* argv[]) { Config cfg = parse_args(argc, argv); const bool need_mask = (cfg.loss_scope == "assistant"); @@ -160,7 +157,7 @@ int main(int argc, char *argv[]) // ── 读文本（trim/过滤规则与原 text_train read_file_lines 逐字一致）── auto train_lines_r = nn::read_text_docs(cfg.text_path); NN_EXIT(train_lines_r, 1, "读取训练文本失败: "); - const std::vector &train_lines = *train_lines_r; + const std::vector& train_lines = *train_lines_r; if (train_lines.empty()) { std::cerr << "训练文本为空: " << cfg.text_path << "\n"; @@ -186,10 +183,12 @@ int main(int argc, char *argv[]) } std::size_t train_tokens = 0; - for (const auto &d : train_r->docs) train_tokens += d.size(); + for (const auto& d : train_r->docs) + train_tokens += d.size(); std::size_t test_tokens = 0; if (test_build) - for (const auto &d : test_build->docs) test_tokens += d.size(); + for (const auto& d : test_build->docs) + test_tokens += d.size(); // ── 配套：源摘要（内容 sha256，不是文件大小）/ license / 生成参数 ── auto text_bytes_r = nn::load_text_file(cfg.text_path); @@ -225,8 +224,7 @@ int main(int argc, char *argv[]) std::cerr << "无法写输出文件: " << cfg.output_path << "\n"; return 1; } - const std::string label = nn::write_file_label(nn::NNDATASET_MAGIC, - nn::NNDATASET_FILE_VERSION); + const std::string label = nn::write_file_label(nn::NNDATASET_MAGIC, nn::NNDATASET_FILE_VERSION); ofs.write(label.data(), static_cast(label.size())); auto wr = nn::write_nndataset(ofs, build); NN_EXIT(wr, 1, "写 .nndataset 失败: "); @@ -235,14 +233,11 @@ int main(int argc, char *argv[]) std::cout << "数据集已生成: " << cfg.output_path << "\n" << " kind: text loss_scope: " << cfg.loss_scope << "\n" << " train: " << train_lines.size() << " docs / " << train_tokens << " tokens"; - if (need_mask) - std::cout << "（assistant 目标 " << train_r->asst_tokens << " tokens）"; + if (need_mask) std::cout << "（assistant 目标 " << train_r->asst_tokens << " tokens）"; std::cout << "\n"; if (test_build) - std::cout << " test: " << test_lines.size() << " docs / " << test_tokens - << " tokens\n"; - std::cout << " 词表: " << tokenizer->vocab_size() << " 词（来自 " - << cfg.vocab_path << "）\n" + std::cout << " test: " << test_lines.size() << " docs / " << test_tokens << " tokens\n"; + std::cout << " 词表: " << tokenizer->vocab_size() << " 词（来自 " << cfg.vocab_path << "）\n" << " source.sha256: " << build.source_sha256.substr(0, 16) << "...\n"; return 0; } diff --git a/src/dataset_samples_test.cpp b/src/dataset_samples_test.cpp index e141561d..81920106 100644 --- a/src/dataset_samples_test.cpp +++ b/src/dataset_samples_test.cpp @@ -130,8 +130,7 @@ void test_windowing_and_samples() auto vr = set.make_batch(idx, batch); auto units_r = set.valid_target_count(idx); CHECK(vr && *vr == 2, "make_batch 返回本批参与 loss 的位置数"); - CHECK(units_r && *units_r == 2, - "valid_target_count 与 batch 的 loss 有效单位一致"); + CHECK(units_r && *units_r == 2, "valid_target_count 与 batch 的 loss 有效单位一致"); CHECK(batch.batch == 2 && batch.seq_len == 4, "batch 形状"); CHECK(batch.valid_tokens == 2, "valid_tokens"); // input: (seq_len, batch) 列 = 样本 diff --git a/src/doc_attn_e2e_test.cpp b/src/doc_attn_e2e_test.cpp index 4aff92d3..23ab9427 100644 --- a/src/doc_attn_e2e_test.cpp +++ b/src/doc_attn_e2e_test.cpp @@ -27,23 +27,22 @@ #include #include -using nn::Scalar; -using nn::Matrix; -using nn::Tensor; +using nn::ActivationType; using nn::ComputeEngine; using nn::GPTModel; -using nn::PosEncodingType; -using nn::ActivationType; +using nn::Matrix; using nn::NormType; +using nn::PosEncodingType; +using nn::Scalar; +using nn::Tensor; -namespace { +namespace +{ // 运行一次 forward（batch=2），返回 logits（vocab × seq*batch，batch-major // 列区间：样本 b = [b*seq, b*seq+seq)）。空 doc 向量 = 纯因果。 -Matrix run_forward(ComputeEngine& eng, GPTModel& model, - const std::vector& tokens0, - const std::vector& tokens1, - const std::vector& doc0, +Matrix run_forward(ComputeEngine& eng, GPTModel& model, const std::vector& tokens0, + const std::vector& tokens1, const std::vector& doc0, const std::vector& doc1) { const std::size_t seq = tokens0.size(); @@ -57,23 +56,34 @@ Matrix run_forward(ComputeEngine& eng, GPTModel& model, } auto x = eng.from_matrix(x_m); - if (!x) { std::cerr << " from_matrix 失败\n"; std::abort(); } + if (!x) + { + std::cerr << " from_matrix 失败\n"; + std::abort(); + } // set_doc_ids 收 batch-major (b*seq+t) 全量；两样本皆空 = 关闭 std::vector ids(doc0); ids.insert(ids.end(), doc1.begin(), doc1.end()); model.set_doc_ids(ids); auto y = model.forward(*x); - if (!y) { std::cerr << " forward 失败: " << y.error().message << "\n"; std::abort(); } + if (!y) + { + std::cerr << " forward 失败: " << y.error().message << "\n"; + std::abort(); + } auto m = eng.to_matrix(*y); - if (!m) { std::cerr << " to_matrix 失败\n"; std::abort(); } + if (!m) + { + std::cerr << " to_matrix 失败\n"; + std::abort(); + } return *m; } // 比较两矩阵的指定列区间，返回最大绝对差 -Scalar max_col_diff(const Matrix& a, const Matrix& b, - std::size_t start_col, std::size_t end_col) +Scalar max_col_diff(const Matrix& a, const Matrix& b, std::size_t start_col, std::size_t end_col) { Scalar worst{0}; for (std::size_t r = 0; r < a.rows(); ++r) @@ -109,8 +119,8 @@ int main(int argc, char* argv[]) ComputeEngine& eng = *engine; const std::size_t vocab = 64, d_model = 16, seq = 8, heads = 2, d_ff = 32, layers = 2; - GPTModel model(vocab, d_model, seq, heads, d_ff, layers, - PosEncodingType::Learned, ActivationType::GeLU, NormType::LayerNorm); + GPTModel model(vocab, d_model, seq, heads, d_ff, layers, PosEncodingType::Learned, + ActivationType::GeLU, NormType::LayerNorm); NN_EXIT(model.init(eng), 1, "GPTModel init 失败: "); // batch=2（铁律 5：注意力测试必须覆盖 batch>1）： @@ -121,9 +131,9 @@ int main(int argc, char* argv[]) const std::vector doc_ids0{1, 1, 1, 1, 2, 2, 2, 2}; const std::vector doc_ids1{3, 3, 3, 3, 4, 4, 4, 4}; // 变体1 / 变体2：仅样本0 的 doc A token 不同，doc B 与样本1 相同 - const std::vector tok_v1{ 5, 7, 9, 11, 20, 21, 22, 23}; + const std::vector tok_v1{5, 7, 9, 11, 20, 21, 22, 23}; const std::vector tok_v2{50, 60, 70, 80, 20, 21, 22, 23}; - const std::vector tok_s{ 6, 8, 10, 12, 30, 31, 32, 33}; + const std::vector tok_s{6, 8, 10, 12, 30, 31, 32, 33}; // logits 列区间（batch-major (vocab, seq*batch)）：样本0=0..7，样本1=8..15 std::cout << "========================================\n"; @@ -135,15 +145,13 @@ int main(int argc, char* argv[]) auto l2_doc = run_forward(eng, model, tok_v2, tok_s, doc_ids0, doc_ids1); const Scalar diff_docB_doc = max_col_diff(l1_doc, l2_doc, 4, 7); - std::cout << " [文档感知] doc A 变化后，doc B 位置 logits 最大差 = " - << diff_docB_doc << "\n"; + std::cout << " [文档感知] doc A 变化后，doc B 位置 logits 最大差 = " << diff_docB_doc << "\n"; const bool invariant_ok = (diff_docB_doc <= Scalar{1e-4}); // ── 1b) 跨样本隔离：样本1 输入两次完全相同 → 列 8..15 必须逐位不变 // （样本0 变了还能传进来 = batch 串扰，batch=1 测不出） const Scalar diff_iso_doc = max_col_diff(l1_doc, l2_doc, 8, 15); - std::cout << " [隔离] 样本0 doc A 变化后，样本1 logits 最大差 = " - << diff_iso_doc << "\n"; + std::cout << " [隔离] 样本0 doc A 变化后，样本1 logits 最大差 = " << diff_iso_doc << "\n"; const bool isolation_ok = (diff_iso_doc <= Scalar{1e-4}); // ── 2) 负对照：纯因果下，doc B 输出应随 doc A 变化（灵敏度） ── @@ -151,20 +159,19 @@ int main(int argc, char* argv[]) auto l2_causal = run_forward(eng, model, tok_v2, tok_s, {}, {}); const Scalar diff_docB_causal = max_col_diff(l1_causal, l2_causal, 4, 7); - std::cout << " [纯因果] doc A 变化后，doc B 位置 logits 最大差 = " - << diff_docB_causal << "\n"; + std::cout << " [纯因果] doc A 变化后，doc B 位置 logits 最大差 = " << diff_docB_causal + << "\n"; const bool sensitive_ok = (diff_docB_causal > Scalar{1e-2}); // ── 2b) 纯因果下的跨样本隔离（同样应成立） ── const Scalar diff_iso_causal = max_col_diff(l1_causal, l2_causal, 8, 15); - std::cout << " [隔离·因果] 样本0 变化后，样本1 logits 最大差 = " - << diff_iso_causal << "\n"; + std::cout << " [隔离·因果] 样本0 变化后，样本1 logits 最大差 = " << diff_iso_causal << "\n"; const bool isolation_causal_ok = (diff_iso_causal <= Scalar{1e-4}); // ── 3) 顺带确认：文档感知与纯因果在 doc B 上确实不同（掩码确实生效） ── const Scalar diff_doc_vs_causal = max_col_diff(l1_doc, l1_causal, 4, 7); - std::cout << " [对照] 文档感知 vs 纯因果（同输入）doc B 最大差 = " - << diff_doc_vs_causal << "\n"; + std::cout << " [对照] 文档感知 vs 纯因果（同输入）doc B 最大差 = " << diff_doc_vs_causal + << "\n"; bool pass = true; if (!invariant_ok) diff --git a/src/dsl_shape_infer_test.cpp b/src/dsl_shape_infer_test.cpp index 348e130e..7bbdd52a 100644 --- a/src/dsl_shape_infer_test.cpp +++ b/src/dsl_shape_infer_test.cpp @@ -33,8 +33,7 @@ namespace { // ── 小工具 ──────────────────────────────────────────────────────────────── -template -nn::Tensor must(R&& r, const char* what) +template nn::Tensor must(R&& r, const char* what) { if (!r) { @@ -47,12 +46,11 @@ nn::Tensor must(R&& r, const char* what) } // 推导路径 (2/3 参) 与显式路径 (4/5 参) 的形状必须一致 -void check_same_shape(nn::CpuEngine&, const nn::Tensor& a, const nn::Tensor& b, - const char* what) +void check_same_shape(nn::CpuEngine&, const nn::Tensor& a, const nn::Tensor& b, const char* what) { const bool ok = (a.rows() == b.rows() && a.cols() == b.cols()); - std::printf("[%s] %s: inferred %zux%zu vs explicit %zux%zu\n", - ok ? " OK " : "FAIL", what, a.rows(), a.cols(), b.rows(), b.cols()); + std::printf("[%s] %s: inferred %zux%zu vs explicit %zux%zu\n", ok ? " OK " : "FAIL", what, + a.rows(), a.cols(), b.rows(), b.cols()); if (!ok) ++g_fail; } @@ -62,8 +60,8 @@ void check_same_values(nn::CpuEngine& eng, const nn::Tensor& a, const nn::Tensor { if (a.rows() != b.rows() || a.cols() != b.cols()) { - std::printf("[FAIL] %s: 形状不同（%zux%zu vs %zux%zu）\n", what, - a.rows(), a.cols(), b.rows(), b.cols()); + std::printf("[FAIL] %s: 形状不同（%zux%zu vs %zux%zu）\n", what, a.rows(), a.cols(), + b.rows(), b.cols()); ++g_fail; return; } @@ -73,8 +71,7 @@ void check_same_values(nn::CpuEngine& eng, const nn::Tensor& a, const nn::Tensor if (ma.span()[i] != mb.span()[i]) { std::printf("[FAIL] %s: [%zu] %.9g != %.9g\n", what, i, - static_cast(ma.span()[i]), - static_cast(mb.span()[i])); + static_cast(ma.span()[i]), static_cast(mb.span()[i])); ++g_fail; return; } @@ -82,13 +79,12 @@ void check_same_values(nn::CpuEngine& eng, const nn::Tensor& a, const nn::Tensor } // 期望失败：返回错误且消息含"无法推导输出形状"（提示改用显式形式） -template -void check_infer_error(R&& r, const char* what) +template void check_infer_error(R&& r, const char* what) { const bool failed = !r; - const bool explained = failed - && r.error().message.find("无法推导输出形状") != std::string::npos - && r.error().message.find("rows/cols") != std::string::npos; + const bool explained = failed && + r.error().message.find("无法推导输出形状") != std::string::npos && + r.error().message.find("rows/cols") != std::string::npos; std::printf("[%s] %s%s\n", explained ? " OK " : "FAIL", what, failed ? (" -> " + r.error().message).c_str() : " （未报错）"); if (!explained) ++g_fail; @@ -100,55 +96,50 @@ void test_infer_matches_explicit() auto eng_owner = nn::CpuEngine::create(); nn::CpuEngine& eng = *eng_owner; - auto t = [&](std::size_t r, std::size_t c, float base) { - return make_tensor(eng, r, c, base, 0.01f); - }; + auto t = [&](std::size_t r, std::size_t c, float base) + { return make_tensor(eng, r, c, base, 0.01f); }; // ① 一元 { nn::Tensor x = t(4, 3, 0.1f); auto e = nn::dsl::exp(nn::dsl::leaf(x)); check_same_shape(eng, must(nn::dsl::compute(eng, e), "unary inferred"), - must(nn::dsl::compute(eng, e, x.rows(), x.cols()), - "unary explicit"), "一元 exp(leaf)"); + must(nn::dsl::compute(eng, e, x.rows(), x.cols()), "unary explicit"), + "一元 exp(leaf)"); check_same_values(eng, must(nn::dsl::compute(eng, e), "unary inferred"), - must(nn::dsl::compute(eng, e, x.rows(), x.cols()), - "unary explicit"), "一元 exp(leaf)"); + must(nn::dsl::compute(eng, e, x.rows(), x.cols()), "unary explicit"), + "一元 exp(leaf)"); } // ② 二元（同形状） { nn::Tensor a = t(4, 3, 0.1f), b = t(4, 3, 0.2f); auto e = nn::dsl::leaf(a) * nn::dsl::leaf(b) + nn::dsl::rparam(Scalar{0.5}); check_same_shape(eng, must(nn::dsl::compute(eng, e), "binary inferred"), - must(nn::dsl::compute(eng, e, 4, 3), "binary explicit"), - "二元 a*b+k"); + must(nn::dsl::compute(eng, e, 4, 3), "binary explicit"), "二元 a*b+k"); } // ③ 列广播 (1,C) + 全网格 { nn::Tensor a = t(4, 3, 0.1f), v = t(1, 3, 0.7f); auto e = nn::dsl::leaf(a) * nn::dsl::col_broadcast(v); check_same_shape(eng, must(nn::dsl::compute(eng, e), "colbcast inferred"), - must(nn::dsl::compute(eng, e, 4, 3), "colbcast explicit"), - "列广播 (1,C)"); + must(nn::dsl::compute(eng, e, 4, 3), "colbcast explicit"), "列广播 (1,C)"); } // ④ 行广播 (F,1) + 全网格 { nn::Tensor a = t(4, 3, 0.1f), v = t(4, 1, 0.7f); auto e = nn::dsl::leaf(a) + nn::dsl::row_broadcast(v); check_same_shape(eng, must(nn::dsl::compute(eng, e), "rowbcast inferred"), - must(nn::dsl::compute(eng, e, 4, 3), "rowbcast explicit"), - "行广播 (F,1)"); + must(nn::dsl::compute(eng, e, 4, 3), "rowbcast explicit"), "行广播 (F,1)"); check_same_values(eng, must(nn::dsl::compute(eng, e), "rowbcast inferred"), - must(nn::dsl::compute(eng, e, 4, 3), "rowbcast explicit"), + must(nn::dsl::compute(eng, e, 4, 3), "rowbcast explicit"), "行广播 (F,1)"); } // ⑤ 条件选择 { nn::Tensor a = t(4, 3, -0.3f); - auto e = nn::dsl::select(nn::dsl::leaf(a) > Scalar{0}, nn::dsl::leaf(a), - Scalar{0}); + auto e = nn::dsl::select(nn::dsl::leaf(a) > Scalar{0}, nn::dsl::leaf(a), Scalar{0}); check_same_shape(eng, must(nn::dsl::compute(eng, e), "select inferred"), - must(nn::dsl::compute(eng, e, 4, 3), "select explicit"), + must(nn::dsl::compute(eng, e, 4, 3), "select explicit"), "select(c>0,a,0)"); } // ⑥ matmul + 行广播 bias（Linear::forward 的形态）──────── transA/transB 四组合 @@ -156,61 +147,59 @@ void test_infer_matches_explicit() nn::Tensor W = t(5, 7, 0.01f), x = t(7, 3, 0.1f), bias = t(1, 3, 0.5f); auto e = nn::dsl::matmul(W, x, false, false) + nn::dsl::row_broadcast(bias); check_same_shape(eng, must(nn::dsl::compute(eng, e), "mm+bias inferred"), - must(nn::dsl::compute(eng, e, 5, 3), "mm+bias explicit"), + must(nn::dsl::compute(eng, e, 5, 3), "mm+bias explicit"), "matmul + row_broadcast(bias)"); check_same_values(eng, must(nn::dsl::compute(eng, e), "mm+bias inferred"), - must(nn::dsl::compute(eng, e, 5, 3), "mm+bias explicit"), + must(nn::dsl::compute(eng, e, 5, 3), "mm+bias explicit"), "matmul + row_broadcast(bias)"); } { // transA：A 存 (K,M) nn::Tensor A = t(7, 5, 0.01f), x = t(7, 3, 0.1f); auto e = nn::dsl::matmul(A, x, true, false); check_same_shape(eng, must(nn::dsl::compute(eng, e), "mm-transA inferred"), - must(nn::dsl::compute(eng, e, 5, 3), "mm-transA explicit"), + must(nn::dsl::compute(eng, e, 5, 3), "mm-transA explicit"), "matmul transA"); } { // transB：B 存 (N,K) nn::Tensor W = t(5, 7, 0.01f), B = t(3, 7, 0.1f); auto e = nn::dsl::matmul(W, B, false, true); check_same_shape(eng, must(nn::dsl::compute(eng, e), "mm-transB inferred"), - must(nn::dsl::compute(eng, e, 5, 3), "mm-transB explicit"), + must(nn::dsl::compute(eng, e, 5, 3), "mm-transB explicit"), "matmul transB"); } // ⑦ batch>1 的 matmul（注意力 Q·Kᵀ 形态：A/B 都按 batch 垂直切分） { const std::size_t BH = 3, seq = 5, dk = 4; nn::Tensor Q = t(BH * dk, seq, 0.01f), K = t(BH * dk, seq, 0.02f); - auto e = nn::dsl::matmul_kv(Q, K, true, false, - static_cast(BH), 1, 1); + auto e = nn::dsl::matmul_kv(Q, K, true, false, static_cast(BH), 1, 1); check_same_shape(eng, must(nn::dsl::compute(eng, e), "bmm inferred"), - must(nn::dsl::compute(eng, e, BH * seq, seq), - "bmm explicit"), "batch=BH 的 matmul (Q·Kᵀ)"); + must(nn::dsl::compute(eng, e, BH * seq, seq), "bmm explicit"), + "batch=BH 的 matmul (Q·Kᵀ)"); check_same_values(eng, must(nn::dsl::compute(eng, e), "bmm inferred"), - must(nn::dsl::compute(eng, e, BH * seq, seq), - "bmm explicit"), "batch=BH 的 matmul (Q·Kᵀ)"); + must(nn::dsl::compute(eng, e, BH * seq, seq), "bmm explicit"), + "batch=BH 的 matmul (Q·Kᵀ)"); } // ⑧ GQA：KV 侧只存 batch/ratio 份行块（b_kv_ratio=2） { const std::size_t BH = 4, seq = 5, dk = 4, ratio = 2; nn::Tensor Q = t(BH * dk, seq, 0.01f), K = t((BH / ratio) * dk, seq, 0.02f); - auto e = nn::dsl::matmul_kv(Q, K, true, false, - static_cast(BH), 1, + auto e = nn::dsl::matmul_kv(Q, K, true, false, static_cast(BH), 1, static_cast(ratio)); check_same_shape(eng, must(nn::dsl::compute(eng, e), "gqa inferred"), - must(nn::dsl::compute(eng, e, BH * seq, seq), - "gqa explicit"), "GQA matmul (ratio=2)"); + must(nn::dsl::compute(eng, e, BH * seq, seq), "gqa explicit"), + "GQA matmul (ratio=2)"); } // ⑨ 分组归约视图根（MaxPool 形态：输出 (rows/R, cols)） { const std::size_t C = 2, R = 4, PB = 3; nn::Tensor col = t(C * R, PB, 0.1f); - auto e = nn::dsl::grouped_reduce_max(col, static_cast(R)) - + nn::dsl::rparam(Scalar{0}); + auto e = nn::dsl::grouped_reduce_max(col, static_cast(R)) + + nn::dsl::rparam(Scalar{0}); check_same_shape(eng, must(nn::dsl::compute(eng, e), "gred inferred"), - must(nn::dsl::compute(eng, e, C, PB), "gred explicit"), + must(nn::dsl::compute(eng, e, C, PB), "gred explicit"), "grouped_reduce_max + 0"); check_same_values(eng, must(nn::dsl::compute(eng, e), "gred inferred"), - must(nn::dsl::compute(eng, e, C, PB), "gred explicit"), + must(nn::dsl::compute(eng, e, C, PB), "gred explicit"), "grouped_reduce_max + 0"); } // ⑩ 归约视图嵌套在逐元素链里（softmax 形态）：输出仍是整网格 @@ -219,11 +208,11 @@ void test_infer_matches_explicit() // 归约视图本身即"按行广播向量"，直接参与算术（无需 row_broadcast 包裹） auto e = nn::dsl::exp(nn::dsl::leaf(x) - nn::dsl::row_reduce_max(x)); check_same_shape(eng, must(nn::dsl::compute(eng, e), "reduce-nest inferred"), - must(nn::dsl::compute(eng, e, 4, 3), - "reduce-nest explicit"), "归约视图嵌套（softmax 形态）"); + must(nn::dsl::compute(eng, e, 4, 3), "reduce-nest explicit"), + "归约视图嵌套（softmax 形态）"); check_same_values(eng, must(nn::dsl::compute(eng, e), "reduce-nest inferred"), - must(nn::dsl::compute(eng, e, 4, 3), - "reduce-nest explicit"), "归约视图嵌套（softmax 形态）"); + must(nn::dsl::compute(eng, e, 4, 3), "reduce-nest explicit"), + "归约视图嵌套（softmax 形态）"); } } @@ -238,10 +227,9 @@ void test_explicit_precision() // F16：3 参重载必须把 P 透传（输出存储精度 f16） nn::Tensor inf = must(nn::dsl::compute(eng, e, nn::Precision::F16), "f16 inferred"); - nn::Tensor exp_ = must(nn::dsl::compute(eng, e, 4, 3, nn::Precision::F16), - "f16 explicit"); - const bool ok = (inf.precision() == nn::Precision::F16) - && (exp_.precision() == nn::Precision::F16); + nn::Tensor exp_ = must(nn::dsl::compute(eng, e, 4, 3, nn::Precision::F16), "f16 explicit"); + const bool ok = + (inf.precision() == nn::Precision::F16) && (exp_.precision() == nn::Precision::F16); std::printf("[%s] 3 参重载透传 P（f16 输出）\n", ok ? " OK " : "FAIL"); if (!ok) ++g_fail; check_same_shape(eng, inf, exp_, "P=F16 的 3 参重载"); @@ -258,8 +246,7 @@ void test_compute_reduce_infer() nn::Tensor inf = must(nn::dsl::compute_reduce(eng, e), "reduce inferred"); nn::Tensor exp_ = must(nn::dsl::compute_reduce(eng, e, 4, 3), "reduce explicit"); const bool ok = (inf.rows() == 4 && inf.cols() == 1); - std::printf("[%s] compute_reduce 无形状重载 → 归约向量 (rows,1)=(4,1)\n", - ok ? " OK " : "FAIL"); + std::printf("[%s] compute_reduce 无形状重载 → 归约向量 (rows,1)=(4,1)\n", ok ? " OK " : "FAIL"); if (!ok) ++g_fail; check_same_shape(eng, inf, exp_, "compute_reduce 无形状重载"); check_same_values(eng, inf, exp_, "compute_reduce 无形状重载"); @@ -278,13 +265,11 @@ void test_negative_cases() "形状冲突 (4,3)+(2,3) → 报错"); } // 无任何可确定网格的操作数（标量 + 网格索引叶） - check_infer_error( - nn::dsl::compute(eng, nn::dsl::row() + nn::dsl::rparam(Scalar{0})), - "无网格贡献者（row()+rparam）→ 报错"); + check_infer_error(nn::dsl::compute(eng, nn::dsl::row() + nn::dsl::rparam(Scalar{0})), + "无网格贡献者（row()+rparam）→ 报错"); // compute_reduce 同样不猜 - check_infer_error( - nn::dsl::compute_reduce(eng, nn::dsl::rparam(Scalar{1})), - "compute_reduce 无网格贡献者 → 报错"); + check_infer_error(nn::dsl::compute_reduce(eng, nn::dsl::rparam(Scalar{1})), + "compute_reduce 无网格贡献者 → 报错"); } } // namespace diff --git a/src/embedding_test.cpp b/src/embedding_test.cpp index 6dffa3b9..ce5ed864 100644 --- a/src/embedding_test.cpp +++ b/src/embedding_test.cpp @@ -19,11 +19,11 @@ int g_fail = 0; -namespace { +namespace +{ // 取表的某一行（宿主侧对拍基准） -std::vector table_row(nn::ComputeEngine& eng, nn::Layer& emb, - std::size_t id) +std::vector table_row(nn::ComputeEngine& eng, nn::Layer& emb, std::size_t id) { auto& t = emb.parameters().at(0).get(); auto m = download(eng, t); @@ -42,7 +42,11 @@ void test_forward_gather() // ids (2,3)：batch>1 覆盖（flat 遍历 = 列主序索引） const nn::Scalar ids_v[6] = {0, 3, 9, 3, 7, 0}; nn::Matrix ids_m(2, 3); - { auto s = ids_m.span(); for (int i = 0; i < 6; ++i) s[i] = ids_v[i]; } + { + auto s = ids_m.span(); + for (int i = 0; i < 6; ++i) + s[i] = ids_v[i]; + } auto ids = upload(*engine, ids_m); auto out = emb.forward(ids); @@ -71,13 +75,21 @@ void test_backward_scatter() // ids 含重复：id 2 出现两次 → 梯度应为两次之和 const nn::Scalar ids_v[4] = {2, 0, 2, 5}; nn::Matrix ids_m(4, 1); - { auto s = ids_m.span(); for (int i = 0; i < 4; ++i) s[i] = ids_v[i]; } + { + auto s = ids_m.span(); + for (int i = 0; i < 4; ++i) + s[i] = ids_v[i]; + } auto ids = upload(*engine, ids_m); CHECK(emb.forward(ids).has_value(), "[3] forward"); // grad_output (d_model=3, N=4)：第 i 列 = i+1 nn::Matrix g(3, 4); - { auto s = g.span(); for (int i = 0; i < 12; ++i) s[i] = static_cast(i + 1); } + { + auto s = g.span(); + for (int i = 0; i < 12; ++i) + s[i] = static_cast(i + 1); + } auto grad = upload(*engine, g); auto gi = emb.backward(grad); CHECK(gi.has_value(), "[3] backward 成功"); @@ -86,7 +98,8 @@ void test_backward_scatter() // 索引梯度 = 零 auto gim = download(*engine, *gi); bool all_zero = true; - for (auto v : gim.span()) if (v != 0.0f) all_zero = false; + for (auto v : gim.span()) + if (v != 0.0f) all_zero = false; CHECK(all_zero, "[4] 索引无梯度（零张量）"); // 表梯度：grad (3,4) 行主序 → 列0=(1,5,9) 列2=(3,7,11)； @@ -96,8 +109,7 @@ void test_backward_scatter() const auto sp = gtm.span(); const float expect[3] = {4.f, 12.f, 20.f}; for (std::size_t c = 0; c < 3; ++c) - CHECK(std::fabs(sp[2 * 3 + c] - expect[c]) < 1e-5f, - "[3] 重复 id 梯度累加"); + CHECK(std::fabs(sp[2 * 3 + c] - expect[c]) < 1e-5f, "[3] 重复 id 梯度累加"); // id 0 的行 = 列1 = (2,6,10) const float expect0[3] = {2.f, 6.f, 10.f}; for (std::size_t c = 0; c < 3; ++c) @@ -119,12 +131,17 @@ void test_padding() CHECK(emb.init(*engine).has_value(), "init"); // [6] padding 行初值恒零 const auto row1 = table_row(*engine, emb, 1); - for (auto v : row1) CHECK(v == 0.0f, "[6] padding 行初值为零"); + for (auto v : row1) + CHECK(v == 0.0f, "[6] padding 行初值为零"); // forward 含 padding id（输出列应为零向量） const nn::Scalar ids_v[3] = {1, 0, 1}; nn::Matrix ids_m(3, 1); - { auto s = ids_m.span(); for (int i = 0; i < 3; ++i) s[i] = ids_v[i]; } + { + auto s = ids_m.span(); + for (int i = 0; i < 3; ++i) + s[i] = ids_v[i]; + } auto ids = upload(*engine, ids_m); auto out = emb.forward(ids); CHECK(out.has_value(), "[6] forward 含 padding id"); @@ -139,7 +156,11 @@ void test_padding() // backward 后 padding 行梯度仍为零 nn::Matrix g(3, 3); - { auto s = g.span(); for (int i = 0; i < 9; ++i) s[i] = 1.0f; } + { + auto s = g.span(); + for (int i = 0; i < 9; ++i) + s[i] = 1.0f; + } auto grad = upload(*engine, g); CHECK(emb.backward(grad).has_value(), "[3] backward（含 padding id）"); auto& gt = emb.param_gradients().at(0).get(); @@ -184,13 +205,16 @@ void test_param_protocol() CHECK(!emb.grad_buffers_allocated(), "[6] init 后仍未分配梯度（惰性）"); CHECK(emb.param_gradients().size() == 1, "[6] param_gradients 恰 1 条"); CHECK(emb.grad_buffers_allocated(), "[6] 取梯度后已分配"); - CHECK(emb.parameters().at(0).get().rows() == 4 && - emb.parameters().at(0).get().cols() == 2, + CHECK(emb.parameters().at(0).get().rows() == 4 && emb.parameters().at(0).get().cols() == 2, "[6] 表形状 (vocab, d_model)"); // backward 前无索引 → fail-fast nn::Matrix g(2, 1); - { auto s = g.span(); s[0] = 1.f; s[1] = 1.f; } + { + auto s = g.span(); + s[0] = 1.f; + s[1] = 1.f; + } auto grad = upload(*engine, g); auto r = emb.backward(grad); CHECK(!r.has_value(), "[4] 无缓存索引时 backward fail-fast"); diff --git a/src/expr_dsl_test.cpp b/src/expr_dsl_test.cpp index a268be23..8738dab0 100644 --- a/src/expr_dsl_test.cpp +++ b/src/expr_dsl_test.cpp @@ -36,8 +36,7 @@ int g_fail = 0; // 手写 RoPE forward 参考（与 RotateHalfRef/RowModRef 语义一致） float rope_ref(nn::ComputeEngine& eng, const nn::Tensor& q, const nn::Tensor& cos, - const nn::Tensor& sin, std::size_t r, std::size_t c, std::size_t dk, - bool backward) + const nn::Tensor& sin, std::size_t r, std::size_t c, std::size_t dk, bool backward) { const std::size_t cols = q.cols(); const auto qm = download(eng, q); @@ -56,13 +55,12 @@ float rope_ref(nn::ComputeEngine& eng, const nn::Tensor& q, const nn::Tensor& co // 内联 RoPE 表达式（与 Layer 里写的同一段，唯一事实来源） template -auto rope_expr(const nn::Tensor& q, const nn::Tensor& cos, const nn::Tensor& sin, - std::uint32_t dk) +auto rope_expr(const nn::Tensor& q, const nn::Tensor& cos, const nn::Tensor& sin, std::uint32_t dk) { auto term1 = leaf(q) * row_mod(cos, dk); auto term2 = rotate_half(q, dk) * row_mod(sin, dk); if constexpr (Backward) return term1 - term2; - else return term1 + term2; + else return term1 + term2; } void test_rope() @@ -86,7 +84,11 @@ void test_rope() { const float want = rope_ref(eng, q, cos, sin, r, c, DK, false); const float got = os[r * COLS + c]; - if (std::fabs(want - got) > 1e-5f) { ok = false; break; } + if (std::fabs(want - got) > 1e-5f) + { + ok = false; + break; + } } CHECK(ok, "RoPE forward CPU 求值与参考一致"); @@ -103,7 +105,8 @@ void test_rope() // 规范 key：同表达式折叠两次 key 相同（AOT 匹配前提） const std::string k1 = nn::expr_spec_key(spec); - const std::string k2 = nn::expr_spec_key(to_expr_spec(rope_expr(q, cos, sin, DK)).first); + const std::string k2 = + nn::expr_spec_key(to_expr_spec(rope_expr(q, cos, sin, DK)).first); CHECK(k1 == k2, "同表达式两次折叠 key 相同"); } @@ -118,13 +121,17 @@ void test_rope() for (std::size_t c = 0; c < COLS; ++c) { const float want = rope_ref(eng, q, cos, sin, r, c, DK, true); - if (std::fabs(want - os[r * COLS + c]) > 1e-5f) { ok = false; break; } + if (std::fabs(want - os[r * COLS + c]) > 1e-5f) + { + ok = false; + break; + } } CHECK(ok, "RoPE backward CPU 求值与参考一致"); auto [spec, _] = to_expr_spec(expr); CHECK(spec.instrs.back().op == static_cast(nn::ExprOp::Sub), "backward 末条=Sub"); CHECK(nn::expr_spec_key(spec) != - nn::expr_spec_key(to_expr_spec(rope_expr(q, cos, sin, DK)).first), + nn::expr_spec_key(to_expr_spec(rope_expr(q, cos, sin, DK)).first), "forward/backward key 可区分"); // 形状无关融合：RowMod/RotateHalf 的 param（周期/块大小 = d_k）是 @@ -132,7 +139,7 @@ void test_rope() // shader，任意 d_k 全融合；glsl_gen 把 param 作为 push constant 读取） constexpr std::uint32_t DK2 = 64; CHECK(nn::expr_spec_key(spec) == - nn::expr_spec_key(to_expr_spec(rope_expr(q, cos, sin, DK2)).first), + nn::expr_spec_key(to_expr_spec(rope_expr(q, cos, sin, DK2)).first), "不同 d_k 的 key 相同（形状无关融合）"); } } @@ -157,7 +164,11 @@ void test_elementwise() const auto bs = bm.span(); bool ok = true; for (std::size_t i = 0; i < R * C; ++i) - if (std::fabs((as[i] * 2.0f + bs[i]) - os[i]) > 1e-5f) { ok = false; break; } + if (std::fabs((as[i] * 2.0f + bs[i]) - os[i]) > 1e-5f) + { + ok = false; + break; + } CHECK(ok, "a*2+b 标量广播正确"); auto [spec, inputs] = to_expr_spec(expr); CHECK(spec.instrs.size() == 2 && spec.num_regs == 2, "a*2+b: 2 指令 2 寄存器"); @@ -174,7 +185,11 @@ void test_elementwise() const auto as = am.span(); bool ok = true; for (std::size_t i = 0; i < R * C; ++i) - if (std::fabs((as[i] > 0.0f ? as[i] : 0.0f) - os[i]) > 1e-5f) { ok = false; break; } + if (std::fabs((as[i] > 0.0f ? as[i] : 0.0f) - os[i]) > 1e-5f) + { + ok = false; + break; + } CHECK(ok, "relu(a) 正确"); auto [spec, _] = to_expr_spec(expr); CHECK(spec.instrs.size() == 1 && spec.num_regs == 1, "relu: 单条指令"); @@ -191,7 +206,11 @@ void test_elementwise() const auto as = am.span(); bool ok = true; for (std::size_t i = 0; i < R * C; ++i) - if (std::fabs((as[i] > 0.0f ? as[i] : 0.0f) - os[i]) > 1e-5f) { ok = false; break; } + if (std::fabs((as[i] > 0.0f ? as[i] : 0.0f) - os[i]) > 1e-5f) + { + ok = false; + break; + } CHECK(ok, "select(a>0,a,0) 正确"); auto [spec, _] = to_expr_spec(expr); CHECK(spec.instrs.size() == 2, "select: 比较+Select 共 2 指令"); @@ -211,7 +230,11 @@ void test_elementwise() const auto bs = bm.span(); bool ok = true; for (std::size_t i = 0; i < R * C; ++i) - if (std::fabs((std::tanh(as[i]) + std::exp(bs[i])) - os[i]) > 1e-4f) { ok = false; break; } + if (std::fabs((std::tanh(as[i]) + std::exp(bs[i])) - os[i]) > 1e-4f) + { + ok = false; + break; + } CHECK(ok, "tanh(a)+exp(b) 正确"); auto [spec, _] = to_expr_spec(expr); CHECK(spec.instrs.size() == 3, "tanh+exp+: 3 指令"); @@ -231,8 +254,7 @@ void test_swiglu() // grad_gate = g*up*s*(1 + gate*(1-s))（内联，与 Layer 相同） { const Scalar one{1}; - auto expr = leaf(g) * leaf(up) - * (leaf(s) * (one + leaf(gate) * (one - leaf(s)))); + auto expr = leaf(g) * leaf(up) * (leaf(s) * (one + leaf(gate) * (one - leaf(s)))); nn::Tensor out = eval_cpu(expr, R, C); const auto om = download(eng, out); const auto os = om.span(); @@ -248,13 +270,18 @@ void test_swiglu() for (std::size_t i = 0; i < R * C; ++i) { const float want = gs[i] * us[i] * ss[i] * (1.0f + gts[i] * (1.0f - ss[i])); - if (std::fabs(want - os[i]) > 1e-4f) { ok = false; break; } + if (std::fabs(want - os[i]) > 1e-4f) + { + ok = false; + break; + } } CHECK(ok, "swiglu grad_gate CPU 正确"); auto [spec, _] = to_expr_spec(expr); CHECK(nn::expr_spec_key(spec) == - nn::expr_spec_key(to_expr_spec( - leaf(g) * leaf(up) * (leaf(s) * (one + leaf(gate) * (one - leaf(s))))).first), + nn::expr_spec_key(to_expr_spec(leaf(g) * leaf(up) * + (leaf(s) * (one + leaf(gate) * (one - leaf(s))))) + .first), "swiglu grad_gate 同表达式 key 稳定"); } // grad_up = g*gate*s @@ -271,13 +298,18 @@ void test_swiglu() const auto gts = gtm.span(); bool ok = true; for (std::size_t i = 0; i < R * C; ++i) - if (std::fabs((gs[i] * gts[i] * ss[i]) - os[i]) > 1e-4f) { ok = false; break; } + if (std::fabs((gs[i] * gts[i] * ss[i]) - os[i]) > 1e-4f) + { + ok = false; + break; + } CHECK(ok, "swiglu grad_up CPU 正确"); auto [spec, _] = to_expr_spec(expr); CHECK(nn::expr_spec_key(spec) != - nn::expr_spec_key(to_expr_spec( - leaf(g) * leaf(up) - * (leaf(s) * (Scalar{1} + leaf(gate) * (Scalar{1} - leaf(s))))).first), + nn::expr_spec_key( + to_expr_spec(leaf(g) * leaf(up) * + (leaf(s) * (Scalar{1} + leaf(gate) * (Scalar{1} - leaf(s))))) + .first), "grad_gate 与 grad_up key 不同"); } } @@ -290,9 +322,7 @@ int main() test_elementwise(); test_swiglu(); - if (g_fail == 0) - std::printf("ALL PASS\n"); - else - std::printf("%d FAILED\n", g_fail); + if (g_fail == 0) std::printf("ALL PASS\n"); + else std::printf("%d FAILED\n", g_fail); return g_fail == 0 ? 0 : 1; } \ No newline at end of file diff --git a/src/expr_fold_test.cpp b/src/expr_fold_test.cpp index 0f05e4df..f51ddc03 100644 --- a/src/expr_fold_test.cpp +++ b/src/expr_fold_test.cpp @@ -33,8 +33,8 @@ #include #include #include -#include // 注意力 fold 构造（定义于 Layer：表达式文本只写在 Layer） -#include "test_common.hpp" // M1：upload/download（库外经引擎进出张量） +#include // 注意力 fold 构造（定义于 Layer：表达式文本只写在 Layer） +#include "test_common.hpp" // M1：upload/download（库外经引擎进出张量） using nn::Scalar; @@ -50,8 +50,9 @@ int test_expr_fold() // IEEE fmax(err, NaN) = err 会静默吞掉 NaN diff——若无此守卫，被测端 // 输出 NaN 时对拍仍可能 PASS，故此处直接把非有限 diff 记为 inf。 const auto err_of = [](Scalar err, Scalar diff) -> Scalar - { return std::isfinite(diff) ? std::fmax(err, diff) - : std::numeric_limits::infinity(); }; + { + return std::isfinite(diff) ? std::fmax(err, diff) : std::numeric_limits::infinity(); + }; auto engine_owner = nn::CpuEngine::create(); nn::CpuEngine& engine = *engine_owner; @@ -59,14 +60,15 @@ int test_expr_fold() std::uniform_real_distribution dist(-2.0f, 2.0f); constexpr Scalar kTol = Scalar{1e-4}; - for (const std::size_t K : {std::size_t{1}, std::size_t{7}, std::size_t{32}, - std::size_t{33}, std::size_t{100}, std::size_t{1000}}) + for (const std::size_t K : {std::size_t{1}, std::size_t{7}, std::size_t{32}, std::size_t{33}, + std::size_t{100}, std::size_t{1000}}) { const std::size_t rows = 5; nn::Matrix mx(rows, K); { auto sp = mx.span(); - for (auto& v : sp) v = dist(rng); + for (auto& v : sp) + v = dist(rng); } nn::Tensor x = upload(engine, mx); const std::array ins{x}; @@ -77,7 +79,10 @@ int test_expr_fold() check(nn::validate_expr_spec(spec, 1).has_value(), "rowmax validate K=" + std::to_string(K)); auto out = engine.eval_expr(spec, ins, rows, 1); - if (!out) { check(false, "rowmax eval K=" + std::to_string(K) + ": " + out.error().message); } + if (!out) + { + check(false, "rowmax eval K=" + std::to_string(K) + ": " + out.error().message); + } else { const auto om = download(engine, *out); @@ -92,8 +97,8 @@ int test_expr_fold() ref = std::fmax(ref, xs[r * K + k]); err = err_of(err, std::fabs(os[r] - ref)); } - check(err <= kTol, "rowmax vs ref K=" + std::to_string(K) + - " err=" + std::to_string(err)); + check(err <= kTol, + "rowmax vs ref K=" + std::to_string(K) + " err=" + std::to_string(err)); } } @@ -103,7 +108,10 @@ int test_expr_fold() check(nn::validate_expr_spec(spec, 1).has_value(), "rowsum validate K=" + std::to_string(K)); auto out = engine.eval_expr(spec, ins, rows, 1); - if (!out) { check(false, "rowsum eval K=" + std::to_string(K) + ": " + out.error().message); } + if (!out) + { + check(false, "rowsum eval K=" + std::to_string(K) + ": " + out.error().message); + } else { const auto om = download(engine, *out); @@ -114,11 +122,13 @@ int test_expr_fold() for (std::size_t r = 0; r < rows; ++r) { Scalar ref = 0; - for (std::size_t k = 0; k < K; ++k) ref += xs[r * K + k]; - err = err_of(err, std::fabs(os[r] - ref) / std::fmax(Scalar{1}, std::fabs(ref))); + for (std::size_t k = 0; k < K; ++k) + ref += xs[r * K + k]; + err = + err_of(err, std::fabs(os[r] - ref) / std::fmax(Scalar{1}, std::fabs(ref))); } - check(err <= kTol, "rowsum vs ref K=" + std::to_string(K) + - " err=" + std::to_string(err)); + check(err <= kTol, + "rowsum vs ref K=" + std::to_string(K) + " err=" + std::to_string(err)); } } @@ -128,7 +138,11 @@ int test_expr_fold() check(nn::validate_expr_spec(spec, 1).has_value(), "softmax_denom validate K=" + std::to_string(K)); auto out = engine.eval_expr(spec, ins, rows, 1); - if (!out) { check(false, "softmax_denom eval K=" + std::to_string(K) + ": " + out.error().message); } + if (!out) + { + check(false, + "softmax_denom eval K=" + std::to_string(K) + ": " + out.error().message); + } else { const auto om = download(engine, *out); @@ -144,10 +158,11 @@ int test_expr_fold() Scalar ref = 0; for (std::size_t k = 0; k < K; ++k) ref += std::exp(xs[r * K + k] - m); - err = err_of(err, std::fabs(os[r] - ref) / std::fmax(Scalar{1}, std::fabs(ref))); + err = + err_of(err, std::fabs(os[r] - ref) / std::fmax(Scalar{1}, std::fabs(ref))); } check(err <= kTol, "softmax_denom vs ref K=" + std::to_string(K) + - " err=" + std::to_string(err)); + " err=" + std::to_string(err)); } } } @@ -164,29 +179,32 @@ int test_expr_fold() nn::ExprSpec s = nn::expr::make_fold_rowsum(16); s.fold->body = {}; nn::ExprInstr b; - b.op = static_cast(nn::ExprOp::Add); + b.op = static_cast(nn::ExprOp::Add); b.dst = 0; - b.a = nn::expr::input(0); - b.b = nn::expr::cst(0); - s.fold->body = { b }; + b.a = nn::expr::input(0); + b.b = nn::expr::cst(0); + s.fold->body = {b}; rejects(s, "state written from element source"); } // (b) fold + 顶层 instrs 非空 { nn::ExprSpec s = nn::expr::make_fold_rowsum(16); nn::ExprInstr t; - t.op = static_cast(nn::ExprOp::Add); + t.op = static_cast(nn::ExprOp::Add); t.dst = 1; - t.a = nn::expr::reg(0); - t.b = nn::expr::cst(0); - s.instrs = { t }; + t.a = nn::expr::reg(0); + t.b = nn::expr::cst(0); + s.instrs = {t}; rejects(s, "fold + non-empty top-level instrs"); } // (c) fold + matmul 段互斥 { nn::ExprSpec s = nn::expr::make_fold_rowsum(16); nn::MatmulSpec mm; - mm.a_input = 0; mm.b_input = 0; mm.k = 16; mm.batch = 1; + mm.a_input = 0; + mm.b_input = 0; + mm.k = 16; + mm.batch = 1; s.matmul = mm; rejects(s, "fold + matmul segment"); } @@ -195,11 +213,11 @@ int test_expr_fold() nn::ExprSpec s = nn::expr::make_fold_rowsum(16); s.fold->finalize = {}; nn::ExprInstr t; - t.op = static_cast(nn::ExprOp::Add); + t.op = static_cast(nn::ExprOp::Add); t.dst = 0; - t.a = nn::expr::reg(0); - t.b = nn::expr::input(0); - s.fold->finalize = { t }; + t.a = nn::expr::reg(0); + t.b = nn::expr::input(0); + s.fold->finalize = {t}; rejects(s, "finalize reading inputs"); } } @@ -214,7 +232,8 @@ int test_expr_fold() nn::Tensor Q = upload(engine, mQ); nn::Tensor K = upload(engine, mK); nn::Tensor Vt = upload(engine, mVt); - nn::ExprSpec sp = nn::expr::make_fold_attn_o(1, 1, 1, nn::expr::AttnMaskKind::Plain, /*score_bias=*/false); + nn::ExprSpec sp = nn::expr::make_fold_attn_o(1, 1, 1, nn::expr::AttnMaskKind::Plain, + /*score_bias=*/false); const std::vector ii{Q, K, Vt}; auto o = engine.eval_expr(sp, ii, 1, 1); Scalar got = Scalar{-999}; @@ -224,8 +243,8 @@ int test_expr_fold() got = om.span()[0]; } check(o.has_value() && std::fabs(got - Scalar{5}) < Scalar{1e-5}, - "attn fold mini(1x1) O=" + std::to_string(got) + " expect 5" - + (o ? "" : (" err=" + o.error().message))); + "attn fold mini(1x1) O=" + std::to_string(got) + " expect 5" + + (o ? "" : (" err=" + o.error().message))); // 分辨"多 j"vs"多行"： // A 单行多 j：bh=1 seq=4 dk=1；B 多行单 j：bh=2 seq=1 dk=1 for (int ab = 0; ab < 2; ++ab) @@ -244,7 +263,8 @@ int test_expr_fold() nn::Tensor Q2 = upload(engine, mQ2); nn::Tensor K2 = upload(engine, mK2); nn::Tensor V2 = upload(engine, mV2); - nn::ExprSpec sp2 = nn::expr::make_fold_attn_o(seq2, 1, bh2, nn::expr::AttnMaskKind::Plain, /*score_bias=*/false); + nn::ExprSpec sp2 = nn::expr::make_fold_attn_o( + seq2, 1, bh2, nn::expr::AttnMaskKind::Plain, /*score_bias=*/false); const std::vector i2{Q2, K2, V2}; auto o2 = engine.eval_expr(sp2, i2, bh2 * seq2, 1); std::string vals = "?"; @@ -257,8 +277,8 @@ int test_expr_fold() vals += " " + std::to_string(sp2o[i]); } // 常数 Q=2,K=3 → 每行 softmax=1 → O = V = 5（所有行） - check(o2.has_value(), std::string("mini-") + (ab ? "B" : "A") - + " eval" + (o2 ? "" : (": " + o2.error().message))); + check(o2.has_value(), std::string("mini-") + (ab ? "B" : "A") + " eval" + + (o2 ? "" : (": " + o2.error().message))); if (o2) { bool okv = true; @@ -266,27 +286,29 @@ int test_expr_fold() const auto sp2o = om2.span(); for (std::size_t i = 0; i < bh2 * seq2; ++i) if (std::fabs(sp2o[i] - Scalar{5}) > Scalar{1e-5}) okv = false; - check(okv, std::string("mini-") + (ab ? "B(bh2seq1)" : "A(seq4)") - + " O=" + vals + " expect 5×" + std::to_string(bh2 * seq2)); + check(okv, std::string("mini-") + (ab ? "B(bh2seq1)" : "A(seq4)") + " O=" + vals + + " expect 5×" + std::to_string(bh2 * seq2)); } } // C: 非均匀 Vt、Q=K=0（s 全 0 → p=1/4 均匀）→ 手算期望： // d0 = 0.25·(10+20+30+40) = 25；d1 = 0.25·(1+2+3+4) = 2.5 { - nn::Matrix mQ3(2, 4), mK3(2, 4), - mV3(4, 2); // Q/K: (bh·dk, seq)=(2,4) - for (auto& v : mQ3.span()) v = 0; - for (auto& v : mK3.span()) v = 0; + nn::Matrix mQ3(2, 4), mK3(2, 4), mV3(4, 2); // Q/K: (bh·dk, seq)=(2,4) + for (auto& v : mQ3.span()) + v = 0; + for (auto& v : mK3.span()) + v = 0; const Scalar vv[8] = {10, 1, 20, 2, 30, 3, 40, 4}; - for (int i = 0; i < 8; ++i) mV3.span()[i] = vv[i]; + for (int i = 0; i < 8; ++i) + mV3.span()[i] = vv[i]; nn::Tensor Q3 = upload(engine, mQ3); nn::Tensor K3 = upload(engine, mK3); nn::Tensor V3 = upload(engine, mV3); - nn::ExprSpec sp3 = nn::expr::make_fold_attn_o(4, 2, 1, nn::expr::AttnMaskKind::Plain, /*score_bias=*/false); + nn::ExprSpec sp3 = nn::expr::make_fold_attn_o(4, 2, 1, nn::expr::AttnMaskKind::Plain, + /*score_bias=*/false); const std::vector i3{Q3, K3, V3}; auto o3 = engine.eval_expr(sp3, i3, 4, 2); - check(o3.has_value(), "mini-C eval" - + (o3 ? "" : (": " + o3.error().message))); + check(o3.has_value(), "mini-C eval" + (o3 ? "" : (": " + o3.error().message))); if (o3) { const auto om3 = download(engine, *o3); @@ -307,20 +329,25 @@ int test_expr_fold() } } - struct Mk { nn::expr::AttnMaskKind mask; bool score_bias; const char* name; }; + struct Mk + { + nn::expr::AttnMaskKind mask; + bool score_bias; + const char* name; + }; for (const Mk mk : {Mk{nn::expr::AttnMaskKind::Plain, false, "plain"}, - Mk{nn::expr::AttnMaskKind::Causal, false, "causal"}, - Mk{nn::expr::AttnMaskKind::Causal, true, "alibi"}, - Mk{nn::expr::AttnMaskKind::CausalDoc, false, "doc"}, - Mk{nn::expr::AttnMaskKind::CausalDoc, true, "alibidoc"}}) + Mk{nn::expr::AttnMaskKind::Causal, false, "causal"}, + Mk{nn::expr::AttnMaskKind::Causal, true, "alibi"}, + Mk{nn::expr::AttnMaskKind::CausalDoc, false, "doc"}, + Mk{nn::expr::AttnMaskKind::CausalDoc, true, "alibidoc"}}) { const char* mname = mk.name; - struct Sh { std::uint32_t bh, seq, dk; }; + struct Sh + { std::uint32_t bh, seq, dk; }; // {2,133,4}：seq > EXPR_FOLD_BLOCK(128) → 跨块流式（m/l 进位 + vecacc // rescale）与 causal 整块跳过分支（k0>qt 空块）首次被执行； // 其余形状覆盖单块与奇数行（clamp + row_ok） - for (const Sh sh : {Sh{2, 9, 4}, Sh{1, 33, 8}, Sh{3, 5, 2}, - Sh{2, 133, 4}}) + for (const Sh sh : {Sh{2, 9, 4}, Sh{1, 33, 8}, Sh{3, 5, 2}, Sh{2, 133, 4}}) { const std::size_t rows_out = static_cast(sh.bh) * sh.seq; nn::Matrix mQ(static_cast(sh.bh) * sh.dk, sh.seq); @@ -333,8 +360,7 @@ int test_expr_fold() // (1, batch*seq) 会越界（AGENTS 教训 #2） const std::uint32_t sseq = sh.seq; nn::Matrix mdoc_col(rows_out, 1); - nn::Matrix mdoc_ids_t( - 1, static_cast(sh.bh) * sh.seq); + nn::Matrix mdoc_ids_t(1, static_cast(sh.bh) * sh.seq); // doc 分段边界：seq > EXPR_FOLD_BLOCK(128) 时放在 **128 之后** // （如 133 → 129），使查询位置 i ≥ 边界的行其首个 fold 块 // (j<128) 被文档掩码**全部**屏蔽——触发首块全屏蔽的分支 @@ -342,21 +368,22 @@ int test_expr_fold() // 边界若落在 128 之内（如 seq/2），首块永远留有有效项， // 该分支测不到。seq ≤ 128 时用 seq/2 二分段。 const std::uint32_t doc_boundary = - sseq > nn::EXPR_FOLD_BLOCK ? nn::EXPR_FOLD_BLOCK + 1u - : sseq / 2u; + sseq > nn::EXPR_FOLD_BLOCK ? nn::EXPR_FOLD_BLOCK + 1u : sseq / 2u; const auto doc_of = [doc_boundary](std::uint32_t pos) -> Scalar { return pos < doc_boundary ? Scalar{1} : Scalar{2}; }; for (std::size_t r = 0; r < rows_out; ++r) - mdoc_col.span()[r] = - doc_of(static_cast(r % sh.seq)); + mdoc_col.span()[r] = doc_of(static_cast(r % sh.seq)); for (std::uint32_t blk = 0; blk < sh.bh; ++blk) for (std::uint32_t j = 0; j < sh.seq; ++j) - mdoc_ids_t.span() - [static_cast(blk) * sh.seq + j] = doc_of(j); - for (auto& v : mQ.span()) v = dist(rng); - for (auto& v : mK.span()) v = dist(rng); - for (auto& v : mVt.span()) v = dist(rng); - for (auto& v : mslopes.span()) v = dist(rng) * Scalar{0.1}; + mdoc_ids_t.span()[static_cast(blk) * sh.seq + j] = doc_of(j); + for (auto& v : mQ.span()) + v = dist(rng); + for (auto& v : mK.span()) + v = dist(rng); + for (auto& v : mVt.span()) + v = dist(rng); + for (auto& v : mslopes.span()) + v = dist(rng) * Scalar{0.1}; // 宿主矩阵全部写完后经引擎上传（出生绑定同一 engine） nn::Tensor Q = upload(engine, mQ); nn::Tensor K = upload(engine, mK); @@ -368,20 +395,23 @@ int test_expr_fold() const bool docm = (mk.mask == nn::expr::AttnMaskKind::CausalDoc); std::vector ins{Q, K, Vt}; if (alibi) ins.push_back(slopes); - if (docm) { ins.push_back(doc_col); ins.push_back(doc_ids_t); } + if (docm) + { + ins.push_back(doc_col); + ins.push_back(doc_ids_t); + } - nn::ExprSpec spec = nn::expr::make_fold_attn_o(sh.seq, sh.dk, sh.bh, mk.mask, mk.score_bias); + nn::ExprSpec spec = + nn::expr::make_fold_attn_o(sh.seq, sh.dk, sh.bh, mk.mask, mk.score_bias); if (auto v = nn::validate_expr_spec(spec, ins.size()); !v) { - check(false, std::string("attn fold ") + mname + " validate: " - + v.error().message); + check(false, std::string("attn fold ") + mname + " validate: " + v.error().message); continue; } auto out = engine.eval_expr(spec, ins, rows_out, sh.dk); if (!out) { - check(false, std::string("attn fold ") + mname + " eval: " - + out.error().message); + check(false, std::string("attn fold ") + mname + " eval: " + out.error().message); continue; } @@ -409,8 +439,8 @@ int test_expr_fold() { Scalar acc = 0; for (std::uint32_t d = 0; d < sh.dk; ++d) - acc += qs[(b * sh.dk + d) * sh.seq + i] * - ks[(b * sh.dk + d) * sh.seq + j]; + acc += + qs[(b * sh.dk + d) * sh.seq + i] * ks[(b * sh.dk + d) * sh.seq + j]; if (mk.mask != nn::expr::AttnMaskKind::Plain && j > i) acc = -std::numeric_limits::infinity(); // 文档块对角：跨文档 -inf（fold body 内 causal→doc→alibi @@ -418,8 +448,8 @@ int test_expr_fold() if (docm && ((i < doc_boundary) != (j < doc_boundary))) acc = -std::numeric_limits::infinity(); if (alibi) - acc += sl[b] * static_cast( - static_cast(j) - static_cast(i)); + acc += sl[b] * + static_cast(static_cast(j) - static_cast(i)); srow[j] = acc; m = std::fmax(m, acc); } @@ -435,18 +465,16 @@ int test_expr_fold() for (std::uint32_t j = 0; j < sh.seq; ++j) acc += prow[j] * vs[(b * sh.seq + j) * sh.dk + d]; const Scalar ref = acc / l; - const std::size_t oi = - static_cast(b) * sh.seq + i; + const std::size_t oi = static_cast(b) * sh.seq + i; ref_out[oi * sh.dk + d] = ref; err = err_of(err, std::fabs(os[oi * sh.dk + d] - ref) / - std::fmax(Scalar{1}, std::fabs(ref))); + std::fmax(Scalar{1}, std::fabs(ref))); } } } - check(err <= kTol, - std::string("attn fold ") + mname + " bh=" + std::to_string(sh.bh) - + " seq=" + std::to_string(sh.seq) + " dk=" + std::to_string(sh.dk) - + " err=" + std::to_string(err)); + check(err <= kTol, std::string("attn fold ") + mname + " bh=" + std::to_string(sh.bh) + + " seq=" + std::to_string(sh.seq) + + " dk=" + std::to_string(sh.dk) + " err=" + std::to_string(err)); if (err > kTol) { std::cout << " out[0..5]:"; @@ -466,15 +494,15 @@ int test_expr_fold() // 注册表自身的两条不变量：① 同结构重复 add 只留一条；② tri_skip 进 key。 { nn::fused::ExprRegistry reg; - const nn::ExprSpec sk = nn::expr::make_fold_attn_o( - 64, 4, 2, nn::expr::AttnMaskKind::Causal, /*score_bias=*/false); // tri_skip=true - const nn::ExprSpec pl = nn::expr::make_fold_attn_o( - 64, 4, 2, nn::expr::AttnMaskKind::Plain, /*score_bias=*/false); // tri_skip=false + const nn::ExprSpec sk = nn::expr::make_fold_attn_o(64, 4, 2, nn::expr::AttnMaskKind::Causal, + /*score_bias=*/false); // tri_skip=true + const nn::ExprSpec pl = nn::expr::make_fold_attn_o(64, 4, 2, nn::expr::AttnMaskKind::Plain, + /*score_bias=*/false); // tri_skip=false reg.add(sk); reg.add(pl); reg.add(nn::expr::make_fold_rowsum(64)); const std::size_t n1 = reg.specs.size(); - reg.add(sk); // 重复 add → 按 key 去重 + reg.add(sk); // 重复 add → 按 key 去重 check(n1 == 3 && reg.specs.size() == 3, "registry add 去重 n=" + std::to_string(reg.specs.size())); // tri_skip 进 key 必须**只翻该字段**来验：Causal vs Plain 的 body 本就 @@ -485,12 +513,12 @@ int test_expr_fold() nn::ExprSpec b = sk; b.fold->tri_skip = !a.fold->tri_skip; check(nn::expr_spec_key(nn::canonicalize_expr_spec(a)) != - nn::expr_spec_key(nn::canonicalize_expr_spec(b)), + nn::expr_spec_key(nn::canonicalize_expr_spec(b)), "fold tri_skip 进 key：同结构仅翻 tri_skip 位后 key 必须不同"); } // 顺带锁死"不同掩码 = 不同结构"（body 差异进 key） check(nn::expr_spec_key(nn::canonicalize_expr_spec(sk)) != - nn::expr_spec_key(nn::canonicalize_expr_spec(pl)), + nn::expr_spec_key(nn::canonicalize_expr_spec(pl)), "Causal vs Plain 掩码 body 不同 → key 不同"); bool saw_skip = false; for (const auto& s : reg.specs) @@ -507,31 +535,29 @@ int test_expr_fold() // ② 10 个组合（掩码清单 × {无偏置, 位置偏置} −"Plain+偏置"，× {MHA,GQA}） // 两两不同（登记少一条 = GPU 硬报错）。 { - const auto key_of = [](nn::expr::AttnMaskKind m, bool sb, - std::uint32_t kv_ratio, std::uint32_t seq, - std::uint32_t dk, std::uint32_t bh) + const auto key_of = [](nn::expr::AttnMaskKind m, bool sb, std::uint32_t kv_ratio, + std::uint32_t seq, std::uint32_t dk, std::uint32_t bh) { return nn::expr_spec_key(nn::canonicalize_expr_spec( nn::expr::make_fold_attn_o(seq, dk, bh, m, sb, kv_ratio))); }; - const auto combo_name = [](nn::expr::AttnMaskKind m, bool sb, - std::uint32_t kv_ratio) + const auto combo_name = [](nn::expr::AttnMaskKind m, bool sb, std::uint32_t kv_ratio) { - std::string n = m == nn::expr::AttnMaskKind::Plain ? "plain" - : m == nn::expr::AttnMaskKind::Causal ? "causal" - : "doc"; + std::string n = m == nn::expr::AttnMaskKind::Plain ? "plain" + : m == nn::expr::AttnMaskKind::Causal ? "causal" + : "doc"; if (sb) n += "+alibi"; n += kv_ratio > 1 ? " gqa" : " mha"; return n; }; std::vector keys; - for (const nn::expr::AttnMaskKind m : {nn::expr::AttnMaskKind::Plain, - nn::expr::AttnMaskKind::Causal, - nn::expr::AttnMaskKind::CausalDoc}) + for (const nn::expr::AttnMaskKind m : + {nn::expr::AttnMaskKind::Plain, nn::expr::AttnMaskKind::Causal, + nn::expr::AttnMaskKind::CausalDoc}) for (const bool sb : {false, true}) { if (m == nn::expr::AttnMaskKind::Plain && sb) - continue; // 本仓不构造该组合（与自登记表同规则） + continue; // 本仓不构造该组合（与自登记表同规则） for (const std::uint32_t kv_ratio : {1u, 2u}) { const std::string k = key_of(m, sb, kv_ratio, 64, 2, 2); @@ -545,9 +571,8 @@ int test_expr_fold() } std::sort(keys.begin(), keys.end()); keys.erase(std::unique(keys.begin(), keys.end()), keys.end()); - check(keys.size() == 10, - "attn fold 10 个组合（掩码 × 偏置 × {MHA,GQA}）互不相同 n=" + - std::to_string(keys.size())); + check(keys.size() == 10, "attn fold 10 个组合（掩码 × 偏置 × {MHA,GQA}）互不相同 n=" + + std::to_string(keys.size())); } return fail; diff --git a/src/expr_fuse_shapes.cpp b/src/expr_fuse_shapes.cpp index d3eb655c..337c30ce 100644 --- a/src/expr_fuse_shapes.cpp +++ b/src/expr_fuse_shapes.cpp @@ -17,15 +17,16 @@ #ifdef NN_EXPR_SCAN namespace { -const int g_scan_self_registered = [] { +const int g_scan_self_registered = [] +{ auto eng_owner = nn::CpuEngine::create(); nn::CpuEngine& eng = *eng_owner; constexpr std::size_t R = 5, C = 7, M = 3, K = 4, N = 2; - auto x = eng.create_tensor(R, C); - auto g = eng.create_tensor(R, 1); - auto w = eng.create_tensor(R, C); - auto a0 = eng.create_tensor(M, K); - auto b0 = eng.create_tensor(K, N); + auto x = eng.create_tensor(R, C); + auto g = eng.create_tensor(R, 1); + auto w = eng.create_tensor(R, C); + auto a0 = eng.create_tensor(M, K); + auto b0 = eng.create_tensor(K, N); auto bias = eng.create_tensor(M, N); // 录制段 1：3 节点链（用例 1） diff --git a/src/expr_fuse_shapes.hpp b/src/expr_fuse_shapes.hpp index e5753b2a..b3a6c5c9 100644 --- a/src/expr_fuse_shapes.hpp +++ b/src/expr_fuse_shapes.hpp @@ -18,48 +18,38 @@ namespace nn::test_fuse // 链 1（3 节点）：t1 = x*2；t2 = t1+3；t3 = t2*row_broadcast(g) template -[[nodiscard]] auto chain3_a(Eng& eng, const Tensor& x, - std::size_t R, std::size_t C) +[[nodiscard]] auto chain3_a(Eng& eng, const Tensor& x, std::size_t R, std::size_t C) { return dsl::compute(eng, dsl::leaf(x) * Scalar{2}, R, C); } template -[[nodiscard]] auto chain3_b(Eng& eng, const Tensor& t, - std::size_t R, std::size_t C) +[[nodiscard]] auto chain3_b(Eng& eng, const Tensor& t, std::size_t R, std::size_t C) { return dsl::compute(eng, dsl::leaf(t) + Scalar{3}, R, C); } template -[[nodiscard]] auto chain3_c(Eng& eng, const Tensor& t, const Tensor& g, - std::size_t R, std::size_t C) +[[nodiscard]] auto chain3_c(Eng& eng, const Tensor& t, const Tensor& g, std::size_t R, + std::size_t C) { return dsl::compute(eng, dsl::leaf(t) * dsl::row_broadcast(g), R, C); } // 链 2（独立分支 B）：u = w+1（分支 A 复用 chain3_a） template -[[nodiscard]] auto indep_b(Eng& eng, const Tensor& w, - std::size_t R, std::size_t C) +[[nodiscard]] auto indep_b(Eng& eng, const Tensor& w, std::size_t R, std::size_t C) { return dsl::compute(eng, dsl::leaf(w) + Scalar{1}, R, C); } // 链 3（S5 matmul）：t1 = matmul(a0,b0)+bias；t2 = max(t1, 0)（relu） template -[[nodiscard]] auto mm_a(Eng& eng, const Tensor& a0, const Tensor& b0, - const Tensor& bias, std::size_t M, std::size_t N) -{ - return dsl::compute(eng, - dsl::matmul(a0, b0, false, false, 1) + dsl::leaf(bias), M, N); -} +[[nodiscard]] auto mm_a(Eng& eng, const Tensor& a0, const Tensor& b0, const Tensor& bias, + std::size_t M, std::size_t N) +{ return dsl::compute(eng, dsl::matmul(a0, b0, false, false, 1) + dsl::leaf(bias), M, N); } template -[[nodiscard]] auto mm_b(Eng& eng, const Tensor& t, - std::size_t M, std::size_t N) +[[nodiscard]] auto mm_b(Eng& eng, const Tensor& t, std::size_t M, std::size_t N) { return dsl::compute(eng, dsl::max(dsl::leaf(t), Scalar{0}), M, N); } // 单表达式（未录制回归 / vec4 路径）：out = x*row_broadcast(g)*2 template -[[nodiscard]] auto one_shot_rb2(Eng& eng, const Tensor& x, const Tensor& g, - std::size_t R, std::size_t C) -{ - return dsl::compute(eng, - dsl::leaf(x) * dsl::row_broadcast(g) * Scalar{2}, R, C); -} +[[nodiscard]] auto one_shot_rb2(Eng& eng, const Tensor& x, const Tensor& g, std::size_t R, + std::size_t C) +{ return dsl::compute(eng, dsl::leaf(x) * dsl::row_broadcast(g) * Scalar{2}, R, C); } // ── 链 4（Adam 金刚石，P2 多输出）──────────────────────────────────────── // K1: m = 0.9·m0 + 0.1·g （状态写穿 → m 张量） @@ -68,37 +58,30 @@ template // K4: p += d（compute_into 目标传递 → p 张量） // 目标：fuse 成**单 kernel 4 输出**（tail=p，extras=[m, v, d]）。 template -[[nodiscard]] auto adam_a(Eng& eng, const Tensor& m0, const Tensor& g, - std::size_t R, std::size_t C) +[[nodiscard]] auto adam_a(Eng& eng, const Tensor& m0, const Tensor& g, std::size_t R, std::size_t C) { // 超参走 RParam（同真实 adam_update_ 的 lr/β/ε 形态）——拼接器的 // RParam 槽偏移曾漏过（GPU 锚位偏移真凶），本链即其回归锁 - return dsl::compute(eng, - dsl::leaf(m0) * dsl::rparam(0.9f) + dsl::leaf(g) * dsl::rparam(0.1f), R, C); + return dsl::compute(eng, dsl::leaf(m0) * dsl::rparam(0.9f) + dsl::leaf(g) * dsl::rparam(0.1f), + R, C); } template -[[nodiscard]] auto adam_b(Eng& eng, const Tensor& v0, const Tensor& g, - std::size_t R, std::size_t C) +[[nodiscard]] auto adam_b(Eng& eng, const Tensor& v0, const Tensor& g, std::size_t R, std::size_t C) { - return dsl::compute(eng, - dsl::leaf(v0) * dsl::rparam(0.99f) - + dsl::leaf(g) * dsl::leaf(g) * dsl::rparam(0.01f), + return dsl::compute( + eng, dsl::leaf(v0) * dsl::rparam(0.99f) + dsl::leaf(g) * dsl::leaf(g) * dsl::rparam(0.01f), R, C); } template -[[nodiscard]] auto adam_c(Eng& eng, const Tensor& m, const Tensor& v, - std::size_t R, std::size_t C) +[[nodiscard]] auto adam_c(Eng& eng, const Tensor& m, const Tensor& v, std::size_t R, std::size_t C) { - return dsl::compute(eng, - dsl::leaf(m) * dsl::rparam(2.0f) - dsl::leaf(v) * dsl::rparam(3.0f), R, C); + return dsl::compute(eng, dsl::leaf(m) * dsl::rparam(2.0f) - dsl::leaf(v) * dsl::rparam(3.0f), R, + C); } -template -[[nodiscard]] auto adam_d_into(Eng& eng, Tensor& p, const Tensor& d) -{ - return dsl::compute_into(eng, dsl::leaf(p) + dsl::leaf(d), p); -} +template [[nodiscard]] auto adam_d_into(Eng& eng, Tensor& p, const Tensor& d) +{ return dsl::compute_into(eng, dsl::leaf(p) + dsl::leaf(d), p); } } // namespace nn::test_fuse diff --git a/src/expr_fuse_test.cpp b/src/expr_fuse_test.cpp index ce3cbbe1..fb64a1bd 100644 --- a/src/expr_fuse_test.cpp +++ b/src/expr_fuse_test.cpp @@ -51,12 +51,14 @@ int g_fail = 0; // Tensor 不再保存引擎指针；测试里的 matrix_of 按张量设备选择对应引擎。 nn::CpuEngine* g_cpu_engine = nullptr; nn::GpuEngine* g_gpu_engine = nullptr; -#define CHECK(cond, msg) \ - do { \ - if (!(cond)) { \ - std::printf("[FAIL] %s\n", msg); \ - ++g_fail; \ - } \ +#define CHECK(cond, msg) \ + do \ + { \ + if (!(cond)) \ + { \ + std::printf("[FAIL] %s\n", msg); \ + ++g_fail; \ + } \ } while (0) nn::Tensor make_cpu_tensor(nn::CpuEngine& cpu, std::size_t rows, std::size_t cols, @@ -64,7 +66,8 @@ nn::Tensor make_cpu_tensor(nn::CpuEngine& cpu, std::size_t rows, std::size_t col { std::uniform_real_distribution dist(-1.0f, 1.0f); nn::Matrix m(rows, cols); - for (auto& v : m.span()) v = dist(rng); + for (auto& v : m.span()) + v = dist(rng); auto t = cpu.from_matrix(m); if (!t) { @@ -95,9 +98,8 @@ nn::Matrix matrix_of(const nn::Tensor& t) ++g_fail; return nn::Matrix(0, 0); } - nn::ComputeEngine& engine = t.is_gpu() - ? static_cast(*g_gpu_engine) - : static_cast(*g_cpu_engine); + nn::ComputeEngine& engine = t.is_gpu() ? static_cast(*g_gpu_engine) + : static_cast(*g_cpu_engine); auto m = engine.to_matrix(t); if (!m) { @@ -137,7 +139,11 @@ int run_chain3(nn::CpuEngine& cpu, nn::GpuEngine& gpu) auto c1 = nn::test_fuse::chain3_a(cpu, x, R, C); auto c2 = nn::test_fuse::chain3_b(cpu, *c1, R, C); auto c3 = nn::test_fuse::chain3_c(cpu, *c2, g, R, C); - if (!c1 || !c2 || !c3) { std::printf("[FAIL] CPU 参考求值失败\n"); return 1; } + if (!c1 || !c2 || !c3) + { + std::printf("[FAIL] CPU 参考求值失败\n"); + return 1; + } // GPU：begin_expr 录制 → 三表达式 → end_expr 融合执行 CHECK(static_cast(gpu.begin_expr()), "begin_expr 成功"); @@ -156,8 +162,8 @@ int run_chain3(nn::CpuEngine& cpu, nn::GpuEngine& gpu) const Scalar err = max_abs_diff(matrix_of(*c3), matrix_of(*g3)); const bool ok = err < 1e-4f; - std::printf("[%s] 3 节点链融合（begin_expr/end_expr） err=%.2e\n", - ok ? "PASS" : "FAIL", static_cast(err)); + std::printf("[%s] 3 节点链融合（begin_expr/end_expr） err=%.2e\n", ok ? "PASS" : "FAIL", + static_cast(err)); if (!ok) return 1; // P2-12：重复录制（同结构 → 图级缓存命中），新输入仍正确 @@ -175,8 +181,8 @@ int run_chain3(nn::CpuEngine& cpu, nn::GpuEngine& gpu) if (!d3 || !h3) return 1; const Scalar err2 = max_abs_diff(matrix_of(*d3), matrix_of(*h3)); const bool ok2 = err2 < 1e-4f; - std::printf("[%s] 重复录制（P2-12 图级缓存命中） err=%.2e\n", - ok2 ? "PASS" : "FAIL", static_cast(err2)); + std::printf("[%s] 重复录制（P2-12 图级缓存命中） err=%.2e\n", ok2 ? "PASS" : "FAIL", + static_cast(err2)); return (ok && ok2) ? 0 : 1; } @@ -192,7 +198,11 @@ int run_independent(nn::CpuEngine& cpu, nn::GpuEngine& gpu) auto c1 = nn::test_fuse::chain3_a(cpu, x, R, C); auto c2 = nn::test_fuse::indep_b(cpu, w, R, C); - if (!c1 || !c2) { std::printf("[FAIL] CPU 参考求值失败\n"); return 1; } + if (!c1 || !c2) + { + std::printf("[FAIL] CPU 参考求值失败\n"); + return 1; + } CHECK(static_cast(gpu.begin_expr()), "begin_expr 成功"); auto g1 = nn::test_fuse::chain3_a(gpu, xg, R, C); @@ -203,9 +213,8 @@ int run_independent(nn::CpuEngine& cpu, nn::GpuEngine& gpu) const Scalar e1 = max_abs_diff(matrix_of(*c1), matrix_of(*g1)); const Scalar e2 = max_abs_diff(matrix_of(*c2), matrix_of(*g2)); const bool ok = e1 < 1e-4f && e2 < 1e-4f; - std::printf("[%s] 独立分支各自物化（2 kernel） err=%.2e/%.2e\n", - ok ? "PASS" : "FAIL", static_cast(e1), - static_cast(e2)); + std::printf("[%s] 独立分支各自物化（2 kernel） err=%.2e/%.2e\n", ok ? "PASS" : "FAIL", + static_cast(e1), static_cast(e2)); return ok ? 0 : 1; } @@ -223,7 +232,11 @@ int run_matmul_chain(nn::CpuEngine& cpu, nn::GpuEngine& gpu) auto c1 = nn::test_fuse::mm_a(cpu, a0, b0, bias, M, N); auto c2 = nn::test_fuse::mm_b(cpu, *c1, M, N); - if (!c1 || !c2) { std::printf("[FAIL] CPU 参考求值失败\n"); return 1; } + if (!c1 || !c2) + { + std::printf("[FAIL] CPU 参考求值失败\n"); + return 1; + } CHECK(static_cast(gpu.begin_expr()), "begin_expr 成功"); auto g1 = nn::test_fuse::mm_a(gpu, a0g, b0g, biasg, M, N); @@ -233,8 +246,8 @@ int run_matmul_chain(nn::CpuEngine& cpu, nn::GpuEngine& gpu) const Scalar err = max_abs_diff(matrix_of(*c2), matrix_of(*g2)); const bool ok = err < 1e-4f; - std::printf("[%s] S5 matmul 链融合（matmul+bias → relu） err=%.2e\n", - ok ? "PASS" : "FAIL", static_cast(err)); + std::printf("[%s] S5 matmul 链融合（matmul+bias → relu） err=%.2e\n", ok ? "PASS" : "FAIL", + static_cast(err)); return ok ? 0 : 1; } @@ -243,8 +256,7 @@ int run_non_recording(nn::CpuEngine& cpu, nn::GpuEngine& gpu) { const std::size_t R = 4; int fail = 0; - for (const std::size_t C : {std::size_t{6}, std::size_t{4}, - std::size_t{8}, std::size_t{12}}) + for (const std::size_t C : {std::size_t{6}, std::size_t{4}, std::size_t{8}, std::size_t{12}}) { std::mt19937 rng(42 + static_cast(C)); const nn::Tensor x = make_cpu_tensor(cpu, R, C, rng); @@ -261,9 +273,8 @@ int run_non_recording(nn::CpuEngine& cpu, nn::GpuEngine& gpu) } const Scalar err = max_abs_diff(matrix_of(*cr), matrix_of(*gr)); const bool ok = err < 1e-4f; - std::printf("[%s] 未录制路径 cols=%zu%s err=%.2e\n", ok ? "PASS" : "FAIL", - C, (C % 4 == 0) ? "（vec4）" : "（标量）", - static_cast(err)); + std::printf("[%s] 未录制路径 cols=%zu%s err=%.2e\n", ok ? "PASS" : "FAIL", C, + (C % 4 == 0) ? "（vec4）" : "（标量）", static_cast(err)); if (!ok) ++fail; } return fail == 0 ? 0 : 1; @@ -278,11 +289,11 @@ int run_adam_diamond(nn::CpuEngine& cpu, nn::GpuEngine& gpu) std::mt19937 rng(55); const nn::Tensor m0 = make_cpu_tensor(cpu, R, C, rng); const nn::Tensor v0 = make_cpu_tensor(cpu, R, C, rng); - const nn::Tensor g = make_cpu_tensor(cpu, R, C, rng); + const nn::Tensor g = make_cpu_tensor(cpu, R, C, rng); nn::Tensor p0 = make_cpu_tensor(cpu, R, C, rng); const nn::Tensor m0g = to_gpu(gpu, m0); const nn::Tensor v0g = to_gpu(gpu, v0); - const nn::Tensor gg = to_gpu(gpu, g); + const nn::Tensor gg = to_gpu(gpu, g); nn::Tensor pg = to_gpu(gpu, p0); nn::Tensor pu = to_gpu(gpu, p0); // GPU 未融合参考的目标（与融合同源） @@ -329,8 +340,8 @@ int run_adam_diamond(nn::CpuEngine& cpu, nn::GpuEngine& gpu) auto g2 = nn::test_fuse::adam_b(gpu, v0g, gg, R, C); auto g3 = nn::test_fuse::adam_c(gpu, *g1, *g2, R, C); auto gd = nn::test_fuse::adam_d_into(gpu, pg, *g3); - CHECK(static_cast(g1) && static_cast(g2) && static_cast(g3) - && static_cast(gd), + CHECK(static_cast(g1) && static_cast(g2) && static_cast(g3) && + static_cast(gd), "录制期金刚石求值返回占位"); CHECK(static_cast(gpu.end_expr()), "end_expr 融合执行成功"); if (!g1 || !g2 || !g3) return 1; @@ -341,8 +352,8 @@ int run_adam_diamond(nn::CpuEngine& cpu, nn::GpuEngine& gpu) const Scalar ep = max_abs_diff(matrix_of(p0), matrix_of(pg)); // p 目标传递 const bool ok = e1 < 1e-4f && e2 < 1e-4f && e3 < 1e-4f && ep < 1e-4f; std::printf("[%s] Adam 金刚石单 kernel 4 输出（m/v/d 写穿 + p 目标）" - " err=%.2e/%.2e/%.2e/%.2e\n", ok ? "PASS" : "FAIL", - static_cast(e1), static_cast(e2), + " err=%.2e/%.2e/%.2e/%.2e\n", + ok ? "PASS" : "FAIL", static_cast(e1), static_cast(e2), static_cast(e3), static_cast(ep)); if (!ok) return 1; @@ -353,9 +364,8 @@ int run_adam_diamond(nn::CpuEngine& cpu, nn::GpuEngine& gpu) const Scalar bp = max_abs_diff(matrix_of(pu), matrix_of(pg)); const bool okb = b1 == 0.0f && b2 == 0.0f && b3 == 0.0f && bp == 0.0f; std::printf("[%s] 融合 vs 未融合 GPU 逐位一致 diff=%.2e/%.2e/%.2e/%.2e\n", - okb ? "PASS" : "FAIL", static_cast(b1), - static_cast(b2), static_cast(b3), - static_cast(bp)); + okb ? "PASS" : "FAIL", static_cast(b1), static_cast(b2), + static_cast(b3), static_cast(bp)); return okb ? 0 : 1; } @@ -377,8 +387,8 @@ int run_contract_guard(nn::CpuEngine& cpu, nn::GpuEngine& gpu) if (!ref) return 1; const Scalar err = max_abs_diff(matrix_of(*ref), matrix_of(dst_r)); const bool ok2 = err < 1e-4f; - std::printf("[%s] 目标传递 dst 写穿 = x*2 err=%.2e\n", - ok2 ? "PASS" : "FAIL", static_cast(err)); + std::printf("[%s] 目标传递 dst 写穿 = x*2 err=%.2e\n", ok2 ? "PASS" : "FAIL", + static_cast(err)); return ok2 ? 0 : 1; } @@ -388,7 +398,8 @@ void set_env_flag(const char* name, bool on) #ifdef _MSC_VER _putenv_s(name, on ? "1" : ""); #else - if (on) ::setenv(name, "1", 1); else ::unsetenv(name); + if (on) ::setenv(name, "1", 1); + else ::unsetenv(name); #endif } @@ -396,8 +407,7 @@ void set_env_flag(const char* name, bool on) bool byte_equal(const nn::Matrix& a, const nn::Matrix& b) { const auto as = a.span(), bs = b.span(); - return as.size() == bs.size() && - std::memcmp(as.data(), bs.data(), as.size_bytes()) == 0; + return as.size() == bs.size() && std::memcmp(as.data(), bs.data(), as.size_bytes()) == 0; } // ── 用例 8（P3 跨链批量派发）：批量 vs 逐个逐字节一致 ───────────────────── @@ -411,7 +421,8 @@ bool byte_equal(const nn::Matrix& a, const nn::Matrix& b) // CPU 逐节点参考按既有容差断言；并断言批量派发恰 2 组（否则 A/B 自比恒真）。 int run_batch_dispatch(nn::CpuEngine& cpu, nn::GpuEngine& gpu) { - struct Shape { std::size_t R, C; }; + struct Shape + { std::size_t R, C; }; const Shape sh[3] = {{3, 5}, {4, 7}, {2, 11}}; const std::size_t CR = 3, CC = 5; // P1 链形状 std::mt19937 rng(20261006); @@ -492,8 +503,7 @@ int run_batch_dispatch(nn::CpuEngine& cpu, nn::GpuEngine& gpu) std::vector> pd(3); CHECK(static_cast(gpu.begin_expr()), "批量用例 begin_expr 成功"); for (int i = 0; i < 3; ++i) - s1[i] = nn::test_fuse::one_shot_rb2(gpu, xg[i], gsg[i], - sh[i].R, sh[i].C); + s1[i] = nn::test_fuse::one_shot_rb2(gpu, xg[i], gsg[i], sh[i].R, sh[i].C); for (int i = 0; i < 3; ++i) { ma[i] = nn::test_fuse::adam_a(gpu, m0g[i], ggg[i], sh[i].R, sh[i].C); @@ -531,30 +541,26 @@ int run_batch_dispatch(nn::CpuEngine& cpu, nn::GpuEngine& gpu) }; const std::uint64_t b0 = backend.fused_batch_dispatches(); - if (!run_gpu(false, outs[0])) - return 1; + if (!run_gpu(false, outs[0])) return 1; const std::uint64_t b1 = backend.fused_batch_dispatches(); - if (!run_gpu(true, outs[1])) - return 1; + if (!run_gpu(true, outs[1])) return 1; // ① 批量 vs 逐个：逐字节相等（memcmp/哈希口径） bool same = outs[0].size() == outs[1].size(); for (std::size_t i = 0; same && i < outs[0].size(); ++i) same = byte_equal(outs[0][i], outs[1][i]); - std::printf("[%s] 批量 vs 逐个逐字节一致（%zu 张量，memcmp）\n", - same ? "PASS" : "FAIL", outs[0].size()); - if (!same) - return 1; + std::printf("[%s] 批量 vs 逐个逐字节一致（%zu 张量，memcmp）\n", same ? "PASS" : "FAIL", + outs[0].size()); + if (!same) return 1; // ② CPU 对照（容差口径同既有用例） Scalar worst = 0.0f; for (std::size_t i = 0; i < cpu_out.size() && i < outs[0].size(); ++i) worst = std::max(worst, max_abs_diff(cpu_out[i], outs[0][i])); const bool cpu_ok = worst < 1e-4f; - std::printf("[%s] CPU 逐节点对照 max_err=%.2e\n", - cpu_ok ? "PASS" : "FAIL", static_cast(worst)); - if (!cpu_ok) - return 1; + std::printf("[%s] CPU 逐节点对照 max_err=%.2e\n", cpu_ok ? "PASS" : "FAIL", + static_cast(worst)); + if (!cpu_ok) return 1; // ③ 批量路径确实走到（否则 ① 是静默降级后的自比恒真）：两个批量组 // = **2 次派发**（3× 同签名单节点 + 3× 金刚石各自成组）——断言恰为 2 @@ -563,11 +569,9 @@ int run_batch_dispatch(nn::CpuEngine& cpu, nn::GpuEngine& gpu) // 次派发或①字节比较失配。设备无 BDA 时 `#b` 不注册、合法逐 kernel // 降级 → 跳过该断言（A/B 一致性仍成立）。 const std::uint64_t batched = b1 - b0; - std::printf("[info] 批量派发次数 = %llu%s\n", - static_cast(batched), + std::printf("[info] 批量派发次数 = %llu%s\n", static_cast(batched), backend.device().bda_available() ? "" : "（无 BDA：预期 0）"); - if (backend.device().bda_available()) - CHECK(batched == 2, "批量派发恰 2 组（同组混形状：3+3）"); + if (backend.device().bda_available()) CHECK(batched == 2, "批量派发恰 2 组（同组混形状：3+3）"); return 0; } diff --git a/src/expr_fused_key_test.cpp b/src/expr_fused_key_test.cpp index d170819a..1d651504 100644 --- a/src/expr_fused_key_test.cpp +++ b/src/expr_fused_key_test.cpp @@ -29,8 +29,7 @@ int g_fail = 0; [[nodiscard]] std::string fk_key(const nn::ExprSpec& s) { return nn::expr_spec_key(nn::canonicalize_expr_spec(s)); } -void fk_eq(const char* name, const std::string& a, const std::string& b, - const char* why) +void fk_eq(const char* name, const std::string& a, const std::string& b, const char* why) { if (a == b) { @@ -38,8 +37,8 @@ void fk_eq(const char* name, const std::string& a, const std::string& b, return; } ++g_fail; - std::printf("[FAIL] %s：key 应相同但不同（%s）\n a=%s\n b=%s\n", - name, why, a.c_str(), b.c_str()); + std::printf("[FAIL] %s：key 应相同但不同（%s）\n a=%s\n b=%s\n", name, why, + a.c_str(), b.c_str()); } } // namespace @@ -48,46 +47,40 @@ int main() auto eng_owner = nn::CpuEngine::create(); nn::CpuEngine& eng = *eng_owner; // 真实实例的操作数（形状/精度都真实） - const nn::Tensor x = eng.create_tensor(8, 4, nn::Precision::F32); - const nn::Tensor w = eng.create_tensor(5, 8, nn::Precision::F32); - const nn::Tensor b = eng.create_tensor(5, 1, nn::Precision::F32); - const nn::Tensor g = eng.create_tensor(5, 4, nn::Precision::F32); + const nn::Tensor x = eng.create_tensor(8, 4, nn::Precision::F32); + const nn::Tensor w = eng.create_tensor(5, 8, nn::Precision::F32); + const nn::Tensor b = eng.create_tensor(5, 1, nn::Precision::F32); + const nn::Tensor g = eng.create_tensor(5, 4, nn::Precision::F32); const nn::Tensor S{}; // 符号张量（默认构造：未绑定、无存储） using namespace nn; // ── ① 符号实例 key ≡ 真实实例 key（FusedAnchor 的全部依据）────────────── - struct Pair { const char* name; std::string real; std::string sym; }; + struct Pair + { + const char* name; + std::string real; + std::string sym; + }; const Pair pairs[] = { - {"relu_fwd", fk_key(dsl::to_expr_spec( - dsl::max(dsl::leaf(x), Scalar{0})).first), - fk_key(dsl::to_expr_spec( - dsl::max(dsl::leaf(S), Scalar{0})).first)}, + {"relu_fwd", fk_key(dsl::to_expr_spec(dsl::max(dsl::leaf(x), Scalar{0})).first), + fk_key(dsl::to_expr_spec(dsl::max(dsl::leaf(S), Scalar{0})).first)}, {"linear_fwd(matmul+row_broadcast)", - fk_key(dsl::to_expr_spec( - dsl::matmul(w, x, false, false) + dsl::row_broadcast(b)).first), - fk_key(dsl::to_expr_spec( - dsl::matmul(S, S, false, false) + dsl::row_broadcast(S)).first)}, - {"matmul_transA", - fk_key(dsl::to_expr_spec(dsl::matmul(w, g, true, false)).first), - fk_key(dsl::to_expr_spec(dsl::matmul(S, S, true, false)).first)}, - {"matmul_transB", - fk_key(dsl::to_expr_spec(dsl::matmul(g, x, false, true)).first), - fk_key(dsl::to_expr_spec(dsl::matmul(S, S, false, true)).first)}, - {"row_reduce_sum", - fk_key(dsl::to_expr_spec(dsl::row_reduce_sum(dsl::leaf(g))).first), - fk_key(dsl::to_expr_spec(dsl::row_reduce_sum(dsl::leaf(S))).first)}, + fk_key(dsl::to_expr_spec(dsl::matmul(w, x, false, false) + dsl::row_broadcast(b)).first), + fk_key(dsl::to_expr_spec(dsl::matmul(S, S, false, false) + dsl::row_broadcast(S)).first)}, + {"matmul_transA", fk_key(dsl::to_expr_spec(dsl::matmul(w, g, true, false)).first), + fk_key(dsl::to_expr_spec(dsl::matmul(S, S, true, false)).first)}, + {"matmul_transB", fk_key(dsl::to_expr_spec(dsl::matmul(g, x, false, true)).first), + fk_key(dsl::to_expr_spec(dsl::matmul(S, S, false, true)).first)}, + {"row_reduce_sum", fk_key(dsl::to_expr_spec(dsl::row_reduce_sum(dsl::leaf(g))).first), + fk_key(dsl::to_expr_spec(dsl::row_reduce_sum(dsl::leaf(S))).first)}, {"grouped_reduce_max", - fk_key(dsl::to_expr_spec( - dsl::grouped_reduce_max(g, 4u) + dsl::rparam(Scalar{0})).first), - fk_key(dsl::to_expr_spec( - dsl::grouped_reduce_max(S, 1u) + dsl::rparam(Scalar{0})).first)}, - {"rotate_half", - fk_key(dsl::to_expr_spec(dsl::rotate_half(x, 16u)).first), - fk_key(dsl::to_expr_spec(dsl::rotate_half(S, 1u)).first)}, - {"row_access", - fk_key(dsl::to_expr_spec(dsl::row_access(x, 8u, 8u)).first), - fk_key(dsl::to_expr_spec(dsl::row_access(S, 1u, 1u)).first)}, + fk_key(dsl::to_expr_spec(dsl::grouped_reduce_max(g, 4u) + dsl::rparam(Scalar{0})).first), + fk_key(dsl::to_expr_spec(dsl::grouped_reduce_max(S, 1u) + dsl::rparam(Scalar{0})).first)}, + {"rotate_half", fk_key(dsl::to_expr_spec(dsl::rotate_half(x, 16u)).first), + fk_key(dsl::to_expr_spec(dsl::rotate_half(S, 1u)).first)}, + {"row_access", fk_key(dsl::to_expr_spec(dsl::row_access(x, 8u, 8u)).first), + fk_key(dsl::to_expr_spec(dsl::row_access(S, 1u, 1u)).first)}, }; for (const Pair& p : pairs) fk_eq(p.name, p.real, p.sym, "符号实例 key ≡ 真实实例 key"); @@ -102,19 +95,18 @@ int main() // ── ② 常量池**值**不进 key（只喂个数）──────────────────────────────── fk_eq("const_value_not_in_key", fk_key(dsl::to_expr_spec(dsl::max(dsl::leaf(x), Scalar{1e-6f})).first), - fk_key(dsl::to_expr_spec(dsl::max(dsl::leaf(x), Scalar{1e-5f})).first), - "常量池只喂个数"); + fk_key(dsl::to_expr_spec(dsl::max(dsl::leaf(x), Scalar{1e-5f})).first), "常量池只喂个数"); // ── ③ matmul 转置不进 key（运行期 operand layout）───────────────────── { - const std::string k00 = fk_key(dsl::to_expr_spec( - dsl::matmul(w, x, false, false) + dsl::row_broadcast(b)).first); - const std::string k01 = fk_key(dsl::to_expr_spec( - dsl::matmul(w, x, false, true) + dsl::row_broadcast(b)).first); - const std::string k10 = fk_key(dsl::to_expr_spec( - dsl::matmul(w, x, true, false) + dsl::row_broadcast(b)).first); - const std::string k11 = fk_key(dsl::to_expr_spec( - dsl::matmul(w, x, true, true) + dsl::row_broadcast(b)).first); + const std::string k00 = fk_key( + dsl::to_expr_spec(dsl::matmul(w, x, false, false) + dsl::row_broadcast(b)).first); + const std::string k01 = + fk_key(dsl::to_expr_spec(dsl::matmul(w, x, false, true) + dsl::row_broadcast(b)).first); + const std::string k10 = + fk_key(dsl::to_expr_spec(dsl::matmul(w, x, true, false) + dsl::row_broadcast(b)).first); + const std::string k11 = + fk_key(dsl::to_expr_spec(dsl::matmul(w, x, true, true) + dsl::row_broadcast(b)).first); fk_eq("matmul_trans_not_in_key", k00, k01, "4 种转置组合共享一个 shader"); fk_eq("matmul_trans_not_in_key", k00, k10, "4 种转置组合共享一个 shader"); fk_eq("matmul_trans_not_in_key", k00, k11, "4 种转置组合共享一个 shader"); @@ -124,19 +116,21 @@ int main() // ── ④ 分组归约 R 不进 key（运行期视图参数）──────────────────────────── fk_eq("grouped_reduce_R_not_in_key", - fk_key(dsl::to_expr_spec( - dsl::grouped_reduce_max(x, 4u) + dsl::rparam(Scalar{0})).first), - fk_key(dsl::to_expr_spec( - dsl::grouped_reduce_max(x, 9u) + dsl::rparam(Scalar{0})).first), + fk_key(dsl::to_expr_spec(dsl::grouped_reduce_max(x, 4u) + dsl::rparam(Scalar{0})).first), + fk_key(dsl::to_expr_spec(dsl::grouped_reduce_max(x, 9u) + dsl::rparam(Scalar{0})).first), "R 经 vp 槽运行期传入"); // sum / max 是**两种视图**（结构）→ 必须仍然区分 { - const std::string ks = fk_key(dsl::to_expr_spec( - dsl::grouped_reduce_sum(x, 4u) + dsl::rparam(Scalar{0})).first); - const std::string km = fk_key(dsl::to_expr_spec( - dsl::grouped_reduce_max(x, 4u) + dsl::rparam(Scalar{0})).first); + const std::string ks = fk_key( + dsl::to_expr_spec(dsl::grouped_reduce_sum(x, 4u) + dsl::rparam(Scalar{0})).first); + const std::string km = fk_key( + dsl::to_expr_spec(dsl::grouped_reduce_max(x, 4u) + dsl::rparam(Scalar{0})).first); if (ks != km) std::printf("[PASS] sum/max 是不同视图（key 不同）\n"); - else { ++g_fail; std::printf("[FAIL] grouped_reduce sum/max key 应不同\n"); } + else + { + ++g_fail; + std::printf("[FAIL] grouped_reduce sum/max key 应不同\n"); + } } std::printf("\nexpr_fused_key_test: %d failure(s)\n", g_fail); diff --git a/src/expr_graph_test.cpp b/src/expr_graph_test.cpp index c40bad3e..595ac606 100644 --- a/src/expr_graph_test.cpp +++ b/src/expr_graph_test.cpp @@ -41,16 +41,18 @@ namespace { int g_fail = 0; -#define CHECK(cond, msg) \ - do { \ - if (!(cond)) { \ - std::printf("[FAIL] %s\n", msg); \ - ++g_fail; \ - } \ +#define CHECK(cond, msg) \ + do \ + { \ + if (!(cond)) \ + { \ + std::printf("[FAIL] %s\n", msg); \ + ++g_fail; \ + } \ } while (0) -nn::Tensor make_tensor(nn::CpuEngine& eng, std::size_t rows, std::size_t cols, - float base = 0.0f, float step = 0.01f) +nn::Tensor make_tensor(nn::CpuEngine& eng, std::size_t rows, std::size_t cols, float base = 0.0f, + float step = 0.01f) { nn::Tensor t = eng.create_tensor(rows, cols); std::vector buf(rows * cols); @@ -156,8 +158,8 @@ void test_chain_fuse() const int nA = g.add_node(make_unary_const(nn::ExprOp::Mul, Scalar{2}), std::vector{x}, rows, cols, false); - nn::Tensor ytag = make_placeholder(eng, rows, cols, - g.tag_of_node[static_cast(nA)]); + nn::Tensor ytag = + make_placeholder(eng, rows, cols, g.tag_of_node[static_cast(nA)]); const int nB = g.add_node(make_unary_const(nn::ExprOp::Add, Scalar{1}), std::vector{ytag}, rows, cols, false); (void)nB; @@ -173,7 +175,8 @@ void test_chain_fuse() // 融合 kernel 求值（P2 多输出：主输出 = B，extras = A 写穿物化） std::vector kinputs; - for (auto& in : kernels[0].inputs) kinputs.push_back(in.external); + for (auto& in : kernels[0].inputs) + kinputs.push_back(in.external); std::vector kouts; kouts.push_back(eng.create_tensor(rows, cols)); kouts.push_back(eng.create_tensor(rows, cols)); @@ -218,20 +221,17 @@ void test_chain3_fuse() nn::ExprGraph g; nn::Tensor x = make_tensor(eng, 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 + 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); + 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 = make_placeholder(eng, rows, cols, - g.tag_of_node[static_cast(nA)]); + nn::Tensor yt = make_placeholder(eng, rows, cols, g.tag_of_node[static_cast(nA)]); const int nB = add_node_chain(nn::ExprOp::Add, Scalar{3}, yt, rows, cols); - nn::Tensor zt = make_placeholder(eng, rows, cols, - g.tag_of_node[static_cast(nB)]); + nn::Tensor zt = make_placeholder(eng, rows, cols, g.tag_of_node[static_cast(nB)]); const int nC = add_node_chain(nn::ExprOp::Sub, Scalar{1}, zt, rows, cols); (void)nC; @@ -242,30 +242,28 @@ void test_chain3_fuse() CHECK(kernels[0].spec.extras.size() == 2, "P2 写穿：A/B 输出进 extras"); std::vector kinputs; - for (auto& in : kernels[0].inputs) kinputs.push_back(in.external); + for (auto& in : kernels[0].inputs) + kinputs.push_back(in.external); // 输出槽序：[0]=tail(C)，[1..]=成员序去尾（A, B） std::vector kouts; - for (int k = 0; k < 3; ++k) kouts.push_back(eng.create_tensor(rows, cols)); + for (int k = 0; k < 3; ++k) + kouts.push_back(eng.create_tensor(rows, cols)); auto mul_r = eng.eval_expr_multi_into(kernels[0].spec, kinputs, rows, cols, kouts); if (!mul_r || nn::dsl::env_flag("NN_DUMP_GLSL")) { std::printf(" [diag] multi_into: %s\n", mul_r.error().message.c_str()); const auto& sp = kernels[0].spec; std::printf(" [diag] key=%s instrs=%zu views=%zu num_regs=%u extras=%zu\n", - nn::expr_spec_key(sp).c_str(), sp.instrs.size(), - sp.views.size(), sp.num_regs, sp.extras.size()); + nn::expr_spec_key(sp).c_str(), sp.instrs.size(), sp.views.size(), sp.num_regs, + sp.extras.size()); for (const auto& v : sp.views) std::printf(" [diag] view kind=%u\n", static_cast(v.kind)); for (const auto& ins : sp.instrs) - std::printf(" [diag] ins op=%u dst=%u a=(%u,%u)\n", - static_cast(ins.op), static_cast(ins.dst), - static_cast(ins.a.kind), static_cast(ins.a.idx)); + std::printf(" [diag] ins op=%u dst=%u a=(%u,%u)\n", static_cast(ins.op), + static_cast(ins.dst), static_cast(ins.a.kind), + static_cast(ins.a.idx)); auto gl = nn::emitter_registry::make("glsl"); - if (gl) - { - std::printf(" [diag] ---- GLSL ----\n%s\n", - gl->generate("diag", sp).c_str()); - } + if (gl) { std::printf(" [diag] ---- GLSL ----\n%s\n", gl->generate("diag", sp).c_str()); } } CHECK(static_cast(mul_r), "融合 kernel 多输出求值成功"); @@ -299,10 +297,10 @@ void test_independent() nn::Tensor x = make_tensor(eng, rows, cols, 0.0f, 0.1f); nn::Tensor w = make_tensor(eng, 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); + (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"); @@ -312,9 +310,11 @@ void test_independent() // 两个 kernel 求值都正确 std::vector k0in; - for (auto& in : kernels[0].inputs) k0in.push_back(in.external); + 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); + 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 求值成功"); @@ -333,14 +333,12 @@ void test_multi_consumer() // 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 = make_placeholder(eng, rows, cols, - g.tag_of_node[static_cast(nA)]); - (void)g.add_node(make_unary_const(nn::ExprOp::Add, Scalar{1}), - std::vector{yt}, rows, cols, false); - nn::Tensor yt2 = make_placeholder(eng, rows, cols, - g.tag_of_node[static_cast(nA)]); - (void)g.add_node(make_unary_const(nn::ExprOp::Sub, Scalar{1}), - std::vector{yt2}, rows, cols, false); + nn::Tensor yt = make_placeholder(eng, rows, cols, g.tag_of_node[static_cast(nA)]); + (void)g.add_node(make_unary_const(nn::ExprOp::Add, Scalar{1}), std::vector{yt}, + rows, cols, false); + nn::Tensor yt2 = make_placeholder(eng, rows, cols, g.tag_of_node[static_cast(nA)]); + (void)g.add_node(make_unary_const(nn::ExprOp::Sub, Scalar{1}), std::vector{yt2}, + rows, cols, false); auto kernels = fuse_expr_graph(g); // P2 写穿物化：A 的输出落回张量（extras）→ 多消费者不再是融合边界， @@ -352,9 +350,11 @@ void test_multi_consumer() // 数值验证：三输出各自与逐节点参考一致 std::vector kinputs; - for (auto& in : kernels[0].inputs) kinputs.push_back(in.external); + for (auto& in : kernels[0].inputs) + kinputs.push_back(in.external); std::vector kouts; - for (int k = 0; k < 3; ++k) kouts.push_back(eng.create_tensor(rows, cols)); + for (int k = 0; k < 3; ++k) + kouts.push_back(eng.create_tensor(rows, cols)); auto mul_r = eng.eval_expr_multi_into(kernels[0].spec, kinputs, rows, cols, kouts); CHECK(static_cast(mul_r), "多输出求值成功"); auto a_ref = eng.eval_expr(make_unary_const(nn::ExprOp::Mul, Scalar{2}), @@ -385,10 +385,8 @@ void test_reduce_boundary() // 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 = make_placeholder(eng, rows, 1, - g.tag_of_node[static_cast(nA)]); + const int nA = g.add_node(make_row_sum(), std::vector{x}, rows, cols, true); + nn::Tensor yt = make_placeholder(eng, rows, 1, g.tag_of_node[static_cast(nA)]); nn::ExprSpec sb; sb.views.push_back(nn::expr::row_broadcast()); // (rows,1) 按行广播 nn::ExprInstr in; @@ -435,12 +433,10 @@ void test_limit_pressure() nn::Tensor x = make_tensor(eng, 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 = make_placeholder(eng, rows, cols, - g.tag_of_node[static_cast(nA)]); - (void)g.add_node(make_unary_const(nn::ExprOp::Add, Scalar{1}), - std::vector{yt}, rows, cols, false); + const int nA = g.add_node(make_abs_chain(40), std::vector{x}, rows, cols, false); + nn::Tensor yt = make_placeholder(eng, rows, cols, g.tag_of_node[static_cast(nA)]); + (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 内 → 融合"); @@ -448,10 +444,8 @@ void test_limit_pressure() // A：60 条 abs；B：5 条 abs 链 → 拼接后 65 > 64 → 不融合 nn::ExprGraph g2; nn::Tensor x2 = make_tensor(eng, 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 = make_placeholder(eng, rows, cols, - g.tag_of_node[static_cast(nC)]); + const int nC = g2.add_node(make_abs_chain(60), std::vector{x2}, rows, cols, false); + nn::Tensor yt2 = make_placeholder(eng, rows, cols, g.tag_of_node[static_cast(nC)]); // B：5 条 abs 链（消费 A 输出） nn::ExprSpec sb = make_abs_chain(5); (void)g2.add_node(sb, std::vector{yt2}, rows, cols, false); @@ -473,10 +467,10 @@ void test_determinism() { const int nA = g.add_node(make_unary_const(nn::ExprOp::Mul, Scalar{2}), std::vector{x}, rows, cols, false); - nn::Tensor yt = make_placeholder(eng, rows, cols, - g.tag_of_node[static_cast(nA)]); - (void)g.add_node(make_unary_const(nn::ExprOp::Add, Scalar{1}), - std::vector{yt}, rows, cols, false); + nn::Tensor yt = + make_placeholder(eng, rows, cols, g.tag_of_node[static_cast(nA)]); + (void)g.add_node(make_unary_const(nn::ExprOp::Add, Scalar{1}), std::vector{yt}, + rows, cols, false); }; nn::ExprGraph g1, g2; @@ -505,8 +499,7 @@ void test_dep_recognition() const int nA = g.add_node(make_unary_const(nn::ExprOp::Mul, Scalar{2}), std::vector{x}, rows, cols, false); // 占位带 tag → 依赖识别 - nn::Tensor yt = make_placeholder(eng, rows, cols, - g.tag_of_node[static_cast(nA)]); + nn::Tensor yt = make_placeholder(eng, rows, cols, g.tag_of_node[static_cast(nA)]); 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, @@ -514,14 +507,12 @@ void test_dep_recognition() // 普通输入 → 外部 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 应为外部输入"); + CHECK(g.nodes[static_cast(nC)].dep_of_input[0] == -1, "C 的输入 0 应为外部输入"); // 未知 tag → 外部 nn::Tensor fake = make_placeholder(eng, rows, cols, 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 视为外部输入"); + CHECK(g.nodes[static_cast(nD)].dep_of_input[0] == -1, "未知 tag 视为外部输入"); std::printf("[OK] test_dep_recognition\n"); } @@ -538,11 +529,10 @@ void test_input_order() nn::Tensor x3 = make_tensor(eng, 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); + 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 = make_placeholder(eng, rows, cols, - g.tag_of_node[static_cast(nA)]); + nn::Tensor yt = make_placeholder(eng, rows, cols, g.tag_of_node[static_cast(nA)]); nn::ExprSpec sb; sb.views.push_back(nn::expr::linear()); sb.views.push_back(nn::expr::linear()); @@ -570,7 +560,8 @@ void test_input_order() // 数值验证：z = x1*x2 + x3（P2 多输出：主输出 = B，extras = A 写穿） std::vector kinputs; - for (auto& in : kernels[0].inputs) kinputs.push_back(in.external); + for (auto& in : kernels[0].inputs) + kinputs.push_back(in.external); std::vector kouts; kouts.push_back(eng.create_tensor(rows, cols)); kouts.push_back(eng.create_tensor(rows, cols)); @@ -613,14 +604,16 @@ void test_matmul_chain_fuse() // A：matmul(A0,B0) + bias nn::ExprSpec sa; - sa.views = {nn::expr::linear(), nn::expr::linear(), nn::expr::linear()}; + sa.views = {nn::expr::linear(), nn::expr::linear(), nn::expr::linear()}; sa.num_regs = 1; - sa.matmul = nn::MatmulSpec{0, 1, 0, 0, static_cast(K), 1}; - sa.instrs.push_back({static_cast(nn::ExprOp::Add), 0, - nn::expr::matmul_op(), nn::expr::input(2), {}}); + sa.matmul = nn::MatmulSpec{0, 1, 0, 0, static_cast(K), 1}; + sa.instrs.push_back({static_cast(nn::ExprOp::Add), + 0, + nn::expr::matmul_op(), + nn::expr::input(2), + {}}); const int nA = g.add_node(sa, std::vector{A0, B0, bias}, M, N, false); - nn::Tensor yt = make_placeholder(eng, M, N, - g.tag_of_node[static_cast(nA)]); + nn::Tensor yt = make_placeholder(eng, M, N, g.tag_of_node[static_cast(nA)]); // B：relu（max(y, 0)）消费 A const int nB = g.add_node(make_unary_const(nn::ExprOp::Max, Scalar{0}), std::vector{yt}, M, N, false); @@ -638,7 +631,8 @@ void test_matmul_chain_fuse() // 融合 kernel 求值 vs 逐节点参考 std::vector kinputs; - for (auto& in : kernels[0].inputs) kinputs.push_back(in.external); + for (auto& in : kernels[0].inputs) + kinputs.push_back(in.external); auto fused = eng.eval_expr(kernels[0].spec, kinputs, M, N); CHECK(static_cast(fused), "融合 kernel 求值成功"); if (!fused) return; @@ -668,18 +662,20 @@ void test_matmul_pure_chain_fuse() // A：纯 matmul（空指令链） nn::ExprSpec sa; - sa.views = {nn::expr::linear(), nn::expr::linear()}; + sa.views = {nn::expr::linear(), nn::expr::linear()}; sa.num_regs = 0; - sa.matmul = nn::MatmulSpec{0, 1, 0, 0, static_cast(K), 1}; + sa.matmul = nn::MatmulSpec{0, 1, 0, 0, static_cast(K), 1}; const int nA = g.add_node(sa, std::vector{A0, B0}, M, N, false); - nn::Tensor yt = make_placeholder(eng, M, N, - g.tag_of_node[static_cast(nA)]); + nn::Tensor yt = make_placeholder(eng, M, N, g.tag_of_node[static_cast(nA)]); // B：Add(y, bias) nn::ExprSpec sb; - sb.views = {nn::expr::linear(), nn::expr::linear()}; + sb.views = {nn::expr::linear(), nn::expr::linear()}; sb.num_regs = 1; - sb.instrs.push_back({static_cast(nn::ExprOp::Add), 0, - nn::expr::input(0), nn::expr::input(1), {}}); + sb.instrs.push_back({static_cast(nn::ExprOp::Add), + 0, + nn::expr::input(0), + nn::expr::input(1), + {}}); const int nB = g.add_node(sb, std::vector{yt, bias}, M, N, false); (void)nB; @@ -695,7 +691,8 @@ void test_matmul_pure_chain_fuse() if (!kernels[0].spec.matmul) return; std::vector kinputs; - for (auto& in : kernels[0].inputs) kinputs.push_back(in.external); + for (auto& in : kernels[0].inputs) + kinputs.push_back(in.external); auto fused = eng.eval_expr(kernels[0].spec, kinputs, M, N); CHECK(static_cast(fused), "融合 kernel 求值成功"); if (!fused) return; @@ -719,7 +716,7 @@ void test_emitter() 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 cpu = nn::emitter_registry::make("cpu"); auto none = nn::emitter_registry::make("nonexistent"); CHECK(static_cast(glsl), "GlslEmitter 已注册可创建"); // CpuEmitter 产物从不参与编译（gen_fused 硬编码 glsl），死代码已清除（docs/17 §1.3） diff --git a/src/expr_matmul_test.cpp b/src/expr_matmul_test.cpp index 0ea82797..e8cf7b32 100644 --- a/src/expr_matmul_test.cpp +++ b/src/expr_matmul_test.cpp @@ -31,7 +31,7 @@ // 测试写在全局作用域（非 namespace nn），避免 ADL 把匹配拉进 nn 命名空间。 using namespace nn::dsl; -namespace dsl = nn::dsl; // 别名：全局作用域可用 dsl::leaf 等限定调用 +namespace dsl = nn::dsl; // 别名：全局作用域可用 dsl::leaf 等限定调用 using nn::Scalar; namespace @@ -40,7 +40,6 @@ namespace int g_fail = 0; - // ── 参考：逐元素加 bias（CPU span 循环，**独立于**被测的 eval_expr/DSL 路径）── nn::Tensor add_bias(nn::ComputeEngine& e, const nn::Tensor& a, const nn::Tensor& bias) { @@ -52,13 +51,14 @@ nn::Tensor add_bias(nn::ComputeEngine& e, const nn::Tensor& a, const nn::Tensor& const auto sa = ma.span(); const auto sb = mb.span(); auto so = out.span(); - for (std::size_t i = 0; i < so.size(); ++i) so[i] = sa[i] + sb[i]; + for (std::size_t i = 0; i < so.size(); ++i) + so[i] = sa[i] + sb[i]; return upload(e, out); } // ── 参考实现：独立 matmul + bias + relu ────────────────────────────────── -nn::Tensor ref_matmul_bias_relu(nn::ComputeEngine& e, const nn::Tensor& A, - const nn::Tensor& B, const nn::Tensor& bias) +nn::Tensor ref_matmul_bias_relu(nn::ComputeEngine& e, const nn::Tensor& A, const nn::Tensor& B, + const nn::Tensor& bias) { auto m = e.matmul(A, B, false, false); nn::Tensor t = add_bias(e, *m, bias); @@ -80,14 +80,18 @@ void test_pure_matmul() // 直接构造 spec（结构 = DSL 折叠产物） nn::ExprSpec s; - s.views = {nn::expr::linear(), nn::expr::linear()}; + s.views = {nn::expr::linear(), nn::expr::linear()}; s.num_regs = 0; - s.matmul = nn::MatmulSpec{/*a_input=*/0, /*b_input=*/1, 0, 0, static_cast(K)}; + s.matmul = nn::MatmulSpec{/*a_input=*/0, /*b_input=*/1, 0, 0, static_cast(K)}; auto v = nn::validate_expr_spec(s, 2); CHECK(static_cast(v), "纯 matmul spec 校验应通过（允许空指令表）"); auto r = eng.eval_expr(s, std::vector{A, B}, M, N); - if (!r) { CHECK(false, "纯 matmul eval_expr 失败"); return; } + if (!r) + { + CHECK(false, "纯 matmul eval_expr 失败"); + return; + } auto ref = eng.matmul(A, B, false, false); check_close(eng, *r, *ref, "纯 matmul：eval_expr == 独立 matmul", g_fail); } @@ -103,23 +107,34 @@ void test_matmul_bias() const nn::Tensor bias = make_tensor(eng, M, N, 0.2f, 0.01f); nn::ExprSpec s; - s.views = {nn::expr::linear(), nn::expr::linear(), nn::expr::linear()}; + s.views = {nn::expr::linear(), nn::expr::linear(), nn::expr::linear()}; s.num_regs = 1; - s.matmul = nn::MatmulSpec{0, 1, 0, 0, static_cast(K)}; + s.matmul = nn::MatmulSpec{0, 1, 0, 0, static_cast(K)}; // Add r0 = matmul + bias（Input 2） - s.instrs.push_back({static_cast(nn::ExprOp::Add), 0, - nn::expr::matmul_op(), nn::expr::input(2), {}}); + s.instrs.push_back({static_cast(nn::ExprOp::Add), + 0, + nn::expr::matmul_op(), + nn::expr::input(2), + {}}); auto v = nn::validate_expr_spec(s, 3); CHECK(static_cast(v), "matmul+bias spec 校验应通过"); auto r = eng.eval_expr(s, std::vector{A, B, bias}, M, N); - if (!r) { CHECK(false, "matmul+bias eval_expr 失败"); return; } + if (!r) + { + CHECK(false, "matmul+bias eval_expr 失败"); + return; + } auto ref = add_bias(eng, *eng.matmul(A, B, false, false), bias); check_close(eng, *r, ref, "matmul+bias：eval_expr == 参考", g_fail); // DSL 路径：dsl::matmul(A,B) + bias auto rd = nn::dsl::compute(eng, dsl::matmul(A, B) + dsl::leaf(bias), M, N); - if (!rd) { CHECK(false, "DSL matmul+bias compute 失败"); return; } + if (!rd) + { + CHECK(false, "DSL matmul+bias compute 失败"); + return; + } check_close(eng, *rd, ref, "DSL matmul+bias：dsl::compute == 参考", g_fail); } @@ -134,16 +149,19 @@ void test_matmul_bias_relu() const nn::Tensor bias = make_tensor(eng, M, N, -0.4f, 0.03f); // DSL：max(matmul(A,B) + bias, 0) —— canonicalize 后与 scan 端一致 - auto r = nn::dsl::compute(eng, - dsl::max(dsl::matmul(A, B) + dsl::leaf(bias), Scalar{0}), M, N); - if (!r) { CHECK(false, "DSL matmul+bias+relu compute 失败"); return; } + auto r = nn::dsl::compute(eng, dsl::max(dsl::matmul(A, B) + dsl::leaf(bias), Scalar{0}), M, N); + if (!r) + { + CHECK(false, "DSL matmul+bias+relu compute 失败"); + return; + } auto ref = ref_matmul_bias_relu(eng, A, B, bias); check_close(eng, *r, ref, "matmul+bias+relu：dsl::compute == 参考", g_fail); // 折叠出的 spec 应含 matmul 段 - auto [spec, inputs] = nn::dsl::to_expr_spec( - dsl::max(dsl::matmul(A, B) + dsl::leaf(bias), Scalar{0})); + auto [spec, inputs] = + nn::dsl::to_expr_spec(dsl::max(dsl::matmul(A, B) + dsl::leaf(bias), Scalar{0})); CHECK(nn::expr_spec_has_matmul(spec), "折叠 spec 应含 matmul 段"); CHECK(inputs.size() == 3, "折叠 spec 应有 3 个输入（A/B/bias）"); } @@ -156,30 +174,39 @@ void test_transpose_variants() nn::CpuEngine& eng = *eng_owner; // A0: (M,K)；At: A^T 存储 (K,M)；B0: (K,N)；Bt: B^T 存储 (N,K) const nn::Tensor A0 = make_tensor(eng, M, K, 0.2f, 0.04f); - const nn::Tensor At = make_tensor(eng, K, M, 0.2f, 0.04f); // A^T + const nn::Tensor At = make_tensor(eng, K, M, 0.2f, 0.04f); // A^T const nn::Tensor B0 = make_tensor(eng, K, N, -0.1f, 0.03f); - const nn::Tensor Bt = make_tensor(eng, N, K, -0.1f, 0.03f); // B^T + const nn::Tensor Bt = make_tensor(eng, N, K, -0.1f, 0.03f); // B^T const nn::Tensor bias = make_tensor(eng, M, N, 0.1f, 0.01f); - const auto run = [&](bool trA, bool trB) { + const auto run = [&](bool trA, bool trB) + { nn::ExprSpec s; - s.views = {nn::expr::linear(), nn::expr::linear(), nn::expr::linear()}; + s.views = {nn::expr::linear(), nn::expr::linear(), nn::expr::linear()}; s.num_regs = 1; - s.matmul = nn::MatmulSpec{0, 1, - static_cast(trA ? 1u : 0u), - static_cast(trB ? 1u : 0u), - static_cast(K)}; - s.instrs.push_back({static_cast(nn::ExprOp::Add), 0, - nn::expr::matmul_op(), nn::expr::input(2), {}}); + s.matmul = + nn::MatmulSpec{0, 1, static_cast(trA ? 1u : 0u), + static_cast(trB ? 1u : 0u), static_cast(K)}; + s.instrs.push_back({static_cast(nn::ExprOp::Add), + 0, + nn::expr::matmul_op(), + nn::expr::input(2), + {}}); const nn::Tensor& A = trA ? At : A0; const nn::Tensor& B = trB ? Bt : B0; auto r = eng.eval_expr(s, std::vector{A, B, bias}, M, N); - if (!r) { CHECK(false, "transpose variant eval_expr 失败"); return std::string{}; } + if (!r) + { + CHECK(false, "transpose variant eval_expr 失败"); + return std::string{}; + } auto ref = add_bias(eng, *eng.matmul(A, B, trA, trB), bias); - check_close(eng, *r, ref, trA && trB ? "matmul(A^T,B^T)+bias == 参考" - : (trA ? "matmul(A^T,B)+bias == 参考" - : (trB ? "matmul(A,B^T)+bias == 参考" - : "matmul(A,B)+bias == 参考")), g_fail); + check_close(eng, *r, ref, + trA && trB + ? "matmul(A^T,B^T)+bias == 参考" + : (trA ? "matmul(A^T,B)+bias == 参考" + : (trB ? "matmul(A,B^T)+bias == 参考" : "matmul(A,B)+bias == 参考")), + g_fail); return nn::expr_spec_key(s); }; @@ -196,32 +223,37 @@ void test_transpose_variants() // 运行期描述子：4 种组合必须编码成 mm_trans 的 4 个不同位模式 // （GPU 端靠它选加载路径；编码错 = 静默读错布局） - const auto mk = [](bool a, bool b) { + const auto mk = [](bool a, bool b) + { nn::ExprSpec s; - s.views = {nn::expr::linear(), nn::expr::linear()}; + s.views = {nn::expr::linear(), nn::expr::linear()}; s.matmul = nn::MatmulSpec{0, 1, static_cast(a ? 1u : 0u), static_cast(b ? 1u : 0u), 1}; return nn::expr_spec_runtime_matmul_trans(s); }; - CHECK(mk(false, false) == 0u && mk(true, false) == 1u && - mk(false, true) == 2u && mk(true, true) == 3u, + CHECK(mk(false, false) == 0u && mk(true, false) == 1u && mk(false, true) == 2u && + mk(true, true) == 3u, "expr_spec_runtime_matmul_trans 位编码（bit0=transA, bit1=transB）"); } // ── 5) key 与形状无关：k 不进 key，同结构不同 K 共享融合 shader ──────── void test_key_shape_invariance() { - auto make_spec = [](std::uint32_t k) { + auto make_spec = [](std::uint32_t k) + { nn::ExprSpec s; - s.views = {nn::expr::linear(), nn::expr::linear(), nn::expr::linear()}; + s.views = {nn::expr::linear(), nn::expr::linear(), nn::expr::linear()}; s.num_regs = 1; - s.matmul = nn::MatmulSpec{0, 1, 0, 0, k}; - s.instrs.push_back({static_cast(nn::ExprOp::Add), 0, - nn::expr::matmul_op(), nn::expr::input(2), {}}); + s.matmul = nn::MatmulSpec{0, 1, 0, 0, k}; + s.instrs.push_back({static_cast(nn::ExprOp::Add), + 0, + nn::expr::matmul_op(), + nn::expr::input(2), + {}}); return s; }; // canonicalize 后 key 一致（k 不参与哈希） - const std::string k8 = nn::expr_spec_key(nn::canonicalize_expr_spec(make_spec(8))); + const std::string k8 = nn::expr_spec_key(nn::canonicalize_expr_spec(make_spec(8))); const std::string k32 = nn::expr_spec_key(nn::canonicalize_expr_spec(make_spec(32))); const std::string k127 = nn::expr_spec_key(nn::canonicalize_expr_spec(make_spec(127))); CHECK(k8 == k32 && k32 == k127, @@ -229,10 +261,13 @@ void test_key_shape_invariance() // 无 matmul 段的同类表达式 key 必须不同（matmul 段结构进 key） nn::ExprSpec s_nomm; - s_nomm.views = {nn::expr::linear(), nn::expr::linear(), nn::expr::linear()}; + s_nomm.views = {nn::expr::linear(), nn::expr::linear(), nn::expr::linear()}; s_nomm.num_regs = 1; - s_nomm.instrs.push_back({static_cast(nn::ExprOp::Add), 0, - nn::expr::input(0), nn::expr::input(2), {}}); + s_nomm.instrs.push_back({static_cast(nn::ExprOp::Add), + 0, + nn::expr::input(0), + nn::expr::input(2), + {}}); const std::string knomm = nn::expr_spec_key(nn::canonicalize_expr_spec(s_nomm)); CHECK(knomm != k8, "有/无 matmul 段的 spec key 必须不同"); @@ -245,7 +280,11 @@ void test_key_shape_invariance() const nn::Tensor bias = make_tensor(eng, M, N, 0.0f, 0.05f); auto r = eng.eval_expr(make_spec(static_cast(K)), std::vector{A, B, bias}, M, N); - if (!r) { CHECK(false, "K=8 求值失败"); return; } + if (!r) + { + CHECK(false, "K=8 求值失败"); + return; + } auto ref = add_bias(eng, *eng.matmul(A, B, false, false), bias); check_close(eng, *r, ref, "K=8 求值 == 参考（k 只影响形状，不影响结构）", g_fail); } @@ -261,18 +300,22 @@ void test_matmul_reduce_instr() // row_sum(matmul(A,B))：归约指令 RowSum 的源是 Matmul 操作数 nn::ExprSpec s; - s.views = {nn::expr::linear(), nn::expr::linear()}; + s.views = {nn::expr::linear(), nn::expr::linear()}; s.num_regs = 1; - s.matmul = nn::MatmulSpec{0, 1, 0, 0, static_cast(K)}; - s.instrs.push_back({static_cast(nn::ExprOp::RowSum), 0, - nn::expr::matmul_op(), {}, {}}); + s.matmul = nn::MatmulSpec{0, 1, 0, 0, static_cast(K)}; + s.instrs.push_back( + {static_cast(nn::ExprOp::RowSum), 0, nn::expr::matmul_op(), {}, {}}); auto v = nn::validate_expr_spec(s, 2); CHECK(static_cast(v), "row_sum(matmul) spec 校验应通过"); CHECK(nn::expr_spec_reduce_axis(s) == 0, "row_sum(matmul) 归约轴应为行(0)"); // eval_expr_reduce：输出 (rows, 1) auto r = eng.eval_expr_reduce(s, std::vector{A, B}, M, N); - if (!r) { CHECK(false, "row_sum(matmul) eval_expr_reduce 失败"); return; } + if (!r) + { + CHECK(false, "row_sum(matmul) eval_expr_reduce 失败"); + return; + } CHECK(r->rows() == M && r->cols() == 1, "row_sum(matmul) 输出形状应为 (M,1)"); auto mm = eng.matmul(A, B, false, false); auto ref = eng.row_reduce_sum(*mm); @@ -280,7 +323,11 @@ void test_matmul_reduce_instr() // 广播输出：eval_expr（非 vector_out）应广播归约结果 auto rb = eng.eval_expr(s, std::vector{A, B}, M, N); - if (!rb) { CHECK(false, "row_sum(matmul) eval_expr 失败"); return; } + if (!rb) + { + CHECK(false, "row_sum(matmul) eval_expr 失败"); + return; + } nn::Matrix refb_m(M, N); { const nn::Matrix rrm = download(eng, *ref); @@ -300,26 +347,30 @@ void test_matmul_reduce_instr() void test_col_matmul_reduce_instr() { const std::size_t K = 4, N = 5, batch = 2, M = 3; - const std::size_t rows = batch * M; // 6：覆盖 rows != M 的 batch 分解 + const std::size_t rows = batch * M; // 6：覆盖 rows != M 的 batch 分解 auto eng_owner = nn::CpuEngine::create(); nn::CpuEngine& eng = *eng_owner; - const nn::Tensor A = make_tensor(eng, rows, K, 0.2f, 0.03f); // transA=0 → (rows, K) + const nn::Tensor A = make_tensor(eng, rows, K, 0.2f, 0.03f); // transA=0 → (rows, K) const nn::Tensor B = make_tensor(eng, batch * K, N, -0.2f, 0.04f); // transB=0 → (batch*K, N) nn::ExprSpec s; - s.views = {nn::expr::linear(), nn::expr::linear()}; + s.views = {nn::expr::linear(), nn::expr::linear()}; s.num_regs = 1; - s.matmul = nn::MatmulSpec{0, 1, 0, 0, static_cast(K), - static_cast(batch)}; - s.instrs.push_back({static_cast(nn::ExprOp::ColMax), 0, - nn::expr::matmul_op(), {}, {}}); + s.matmul = nn::MatmulSpec{ + 0, 1, 0, 0, static_cast(K), static_cast(batch)}; + s.instrs.push_back( + {static_cast(nn::ExprOp::ColMax), 0, nn::expr::matmul_op(), {}, {}}); auto v = nn::validate_expr_spec(s, 2); CHECK(static_cast(v), "col_max(matmul) spec 校验应通过"); CHECK(nn::expr_spec_reduce_axis(s) == 1, "col_max(matmul) 归约轴应为列(1)"); // 归约向量输出：(1, N) auto r = eng.eval_expr_reduce(s, std::vector{A, B}, rows, N); - if (!r) { CHECK(false, "col_max(matmul) eval_expr_reduce 失败"); return; } + if (!r) + { + CHECK(false, "col_max(matmul) eval_expr_reduce 失败"); + return; + } CHECK(r->rows() == 1 && r->cols() == N, "col_max(matmul) 输出形状应为 (1,N)"); // 独立标量参考（不走引擎解释器，避免同源）： @@ -337,11 +388,15 @@ void test_col_matmul_reduce_instr() bool first = true; for (std::size_t i = 0; i < rows; ++i) { - const std::size_t b = i / M; // 全局行 → batch（rows = batch*M） + const std::size_t b = i / M; // 全局行 → batch（rows = batch*M） float acc = 0.0f; for (std::size_t kk = 0; kk < K; ++kk) acc += as[i * K + kk] * bs[(b * K + kk) * N + j]; - if (first || acc > mx) { mx = acc; first = false; } + if (first || acc > mx) + { + mx = acc; + first = false; + } } rs[j] = mx; } @@ -351,7 +406,11 @@ void test_col_matmul_reduce_instr() // 广播输出：eval_expr（非 vector_out）按列广播归约结果 auto rb = eng.eval_expr(s, std::vector{A, B}, rows, N); - if (!rb) { CHECK(false, "col_max(matmul) eval_expr 失败"); return; } + if (!rb) + { + CHECK(false, "col_max(matmul) eval_expr 失败"); + return; + } nn::Matrix refb_m(rows, N); { const nn::Matrix rm = download(eng, ref); @@ -369,19 +428,22 @@ void test_col_matmul_reduce_instr() void test_validation() { nn::ExprSpec s; - s.views = {nn::expr::linear(), nn::expr::linear()}; + s.views = {nn::expr::linear(), nn::expr::linear()}; s.num_regs = 1; // 无 matmul 段却使用 Matmul 操作数 - s.instrs.push_back({static_cast(nn::ExprOp::Add), 0, - nn::expr::matmul_op(), nn::expr::input(1), {}}); + s.instrs.push_back({static_cast(nn::ExprOp::Add), + 0, + nn::expr::matmul_op(), + nn::expr::input(1), + {}}); auto v = nn::validate_expr_spec(s, 2); CHECK(!v, "无 matmul 段的 Matmul 操作数必须被校验拒绝"); // matmul A/B 输入越界 nn::ExprSpec s2; - s2.views = {nn::expr::linear()}; + s2.views = {nn::expr::linear()}; s2.num_regs = 0; - s2.matmul = nn::MatmulSpec{5, 1, 0, 0, 4}; + s2.matmul = nn::MatmulSpec{5, 1, 0, 0, 4}; auto v2 = nn::validate_expr_spec(s2, 1); CHECK(!v2, "matmul 输入下标越界必须被校验拒绝"); } @@ -395,7 +457,11 @@ void test_dsl_pure_matmul() const nn::Tensor A = make_tensor(eng, M, K, 0.3f, 0.02f); const nn::Tensor B = make_tensor(eng, K, N, 0.1f, 0.03f); auto r = nn::dsl::compute(eng, dsl::matmul(A, B), M, N); - if (!r) { CHECK(false, "DSL 纯 matmul compute 失败"); return; } + if (!r) + { + CHECK(false, "DSL 纯 matmul compute 失败"); + return; + } auto ref = eng.matmul(A, B, false, false); check_close(eng, *r, *ref, "DSL 纯 matmul == 独立 matmul", g_fail); } diff --git a/src/expr_opt_test.cpp b/src/expr_opt_test.cpp index 8295cd38..82842863 100644 --- a/src/expr_opt_test.cpp +++ b/src/expr_opt_test.cpp @@ -36,8 +36,7 @@ int g_fail = 0; // 辅助：直接构造 ExprSpec（控制性测试用）；num_inputs = views.size() -nn::ExprSpec build_spec(std::vector instrs, - std::vector views, +nn::ExprSpec build_spec(std::vector instrs, std::vector views, std::vector consts) { nn::ExprSpec s; @@ -57,18 +56,14 @@ void test_fold() // x*0 → 0（中间指令折叠为常量）；Add(Const0, input1) 为输出 → 保留 { auto s = build_spec( - {{static_cast(nn::ExprOp::Mul), 0, nn::expr::input(0), - nn::expr::cst(0), {}}, - {static_cast(nn::ExprOp::Add), 1, nn::expr::reg(0), - nn::expr::input(1), {}}}, + {{static_cast(nn::ExprOp::Mul), 0, nn::expr::input(0), nn::expr::cst(0), {}}, + {static_cast(nn::ExprOp::Add), 1, nn::expr::reg(0), nn::expr::input(1), {}}}, {nn::expr::linear(), nn::expr::linear()}, {Scalar{0}}); auto c = nn::canonicalize_expr_spec(s); CHECK(c.instrs.size() == 1, "fold: x*0 → 常量 0 后仅剩输出 Add"); - CHECK(static_cast(c.instrs[0].op) == nn::ExprOp::Add, - "fold: 输出指令仍为 Add"); - CHECK(c.instrs[0].a.kind == - static_cast(nn::ExprOperandKind::Const) && - c.consts[c.instrs[0].a.idx] == Scalar{0}, + CHECK(static_cast(c.instrs[0].op) == nn::ExprOp::Add, "fold: 输出指令仍为 Add"); + CHECK(c.instrs[0].a.kind == static_cast(nn::ExprOperandKind::Const) && + c.consts[c.instrs[0].a.idx] == Scalar{0}, "fold: Add 的 a 操作数被折叠成常量 0"); auto v = nn::validate_expr_spec(c, 2); CHECK(v.has_value(), "fold: canonical 通过 validate"); @@ -77,16 +72,13 @@ void test_fold() // x + 0 → x（中间指令消失） { auto s = build_spec( - {{static_cast(nn::ExprOp::Add), 0, nn::expr::input(0), - nn::expr::cst(0), {}}, + {{static_cast(nn::ExprOp::Add), 0, nn::expr::input(0), nn::expr::cst(0), {}}, {static_cast(nn::ExprOp::Neg), 1, nn::expr::reg(0), {}, {}}}, {nn::expr::linear()}, {Scalar{0}}); auto c = nn::canonicalize_expr_spec(s); CHECK(c.instrs.size() == 1, "fold: x+0 → x 后仅剩输出 Neg"); - CHECK(static_cast(c.instrs[0].op) == nn::ExprOp::Neg, - "fold: 输出指令仍为 Neg"); - CHECK(c.instrs[0].a.kind == - static_cast(nn::ExprOperandKind::Input), + CHECK(static_cast(c.instrs[0].op) == nn::ExprOp::Neg, "fold: 输出指令仍为 Neg"); + CHECK(c.instrs[0].a.kind == static_cast(nn::ExprOperandKind::Input), "fold: Neg 的操作数直接引用输入（x+0 消除）"); } @@ -95,41 +87,34 @@ void test_fold() auto s = build_spec( {{static_cast(nn::ExprOp::Neg), 0, nn::expr::input(0), {}, {}}, {static_cast(nn::ExprOp::Neg), 1, nn::expr::reg(0), {}, {}}, - {static_cast(nn::ExprOp::Add), 2, nn::expr::reg(1), - nn::expr::input(1), {}}}, + {static_cast(nn::ExprOp::Add), 2, nn::expr::reg(1), nn::expr::input(1), {}}}, {nn::expr::linear(), nn::expr::linear()}, {}); auto c = nn::canonicalize_expr_spec(s); // neg(neg(x))→x；输出 Add(x, input1) CHECK(c.instrs.size() == 1, "fold: neg(neg(x))→x 消除两条 Neg"); - CHECK(c.instrs[0].a.kind == - static_cast(nn::ExprOperandKind::Input), + CHECK(c.instrs[0].a.kind == static_cast(nn::ExprOperandKind::Input), "fold: Add 的 a 直接引用输入（neg(neg) 消除）"); } // max(x, x) → x；纯常量算术 Add(2,3) → 5 { auto s = build_spec( - {{static_cast(nn::ExprOp::Max), 0, nn::expr::input(0), - nn::expr::input(0), {}}, - {static_cast(nn::ExprOp::Add), 1, nn::expr::cst(0), - nn::expr::cst(1), {}}, - {static_cast(nn::ExprOp::Mul), 2, nn::expr::reg(0), - nn::expr::reg(1), {}}}, + {{static_cast(nn::ExprOp::Max), 0, nn::expr::input(0), nn::expr::input(0), {}}, + {static_cast(nn::ExprOp::Add), 1, nn::expr::cst(0), nn::expr::cst(1), {}}, + {static_cast(nn::ExprOp::Mul), 2, nn::expr::reg(0), nn::expr::reg(1), {}}}, {nn::expr::linear()}, {Scalar{2}, Scalar{3}}); auto c = nn::canonicalize_expr_spec(s); // max(x,x)→input0（reg0 别名）；Add(2,3)→Const(5)（reg1 别名）； // 输出 Mul(input0, Const(5))——但 Mul 是输出，只 remap 不化简 CHECK(c.instrs.size() == 1, "fold: max(x,x) 与常量折叠后仅剩输出 Mul"); - CHECK(static_cast(c.instrs[0].op) == nn::ExprOp::Mul, - "fold: 输出指令为 Mul"); - CHECK(c.instrs[0].a.kind == - static_cast(nn::ExprOperandKind::Input), + CHECK(static_cast(c.instrs[0].op) == nn::ExprOp::Mul, "fold: 输出指令为 Mul"); + CHECK(c.instrs[0].a.kind == static_cast(nn::ExprOperandKind::Input), "fold: Mul 的 a 直接引用输入（max(x,x) 消除）"); - CHECK(c.instrs[0].b.kind == - static_cast(nn::ExprOperandKind::Const), + CHECK(c.instrs[0].b.kind == static_cast(nn::ExprOperandKind::Const), "fold: Mul 的 b 为常量（Add(2,3) 折叠）"); bool found5 = false; - for (Scalar v : c.consts) if (v == Scalar{5}) found5 = true; + for (Scalar v : c.consts) + if (v == Scalar{5}) found5 = true; CHECK(found5, "fold: 常量池含折叠结果 5"); } @@ -137,14 +122,12 @@ void test_fold() { auto s = build_spec( {{static_cast(nn::ExprOp::Exp), 0, nn::expr::cst(0), {}, {}}, - {static_cast(nn::ExprOp::Add), 1, nn::expr::reg(0), - nn::expr::input(0), {}}}, + {static_cast(nn::ExprOp::Add), 1, nn::expr::reg(0), nn::expr::input(0), {}}}, {nn::expr::linear()}, {Scalar{2.0f}}); auto c = nn::canonicalize_expr_spec(s); // exp 不折叠 → 指令保留；Add 输出保留 CHECK(c.instrs.size() == 2, "fold: 超越函数 exp 不被折叠（保守）"); - CHECK(static_cast(c.instrs[0].op) == nn::ExprOp::Exp, - "fold: exp 指令保留"); + CHECK(static_cast(c.instrs[0].op) == nn::ExprOp::Exp, "fold: exp 指令保留"); } } @@ -154,15 +137,16 @@ void test_dce() // r0 = Mul 结果未被引用 → 死代码 { auto s = build_spec( - {{static_cast(nn::ExprOp::Mul), 0, nn::expr::input(0), - nn::expr::input(1), {}}, - {static_cast(nn::ExprOp::Add), 1, nn::expr::input(0), - nn::expr::input(2), {}}}, + {{static_cast(nn::ExprOp::Mul), 0, nn::expr::input(0), nn::expr::input(1), {}}, + {static_cast(nn::ExprOp::Add), + 1, + nn::expr::input(0), + nn::expr::input(2), + {}}}, {nn::expr::linear(), nn::expr::linear(), nn::expr::linear()}, {}); auto c = nn::canonicalize_expr_spec(s); CHECK(c.instrs.size() == 1, "dce: 未引用指令被消除"); - CHECK(static_cast(c.instrs[0].op) == nn::ExprOp::Add, - "dce: 保留输出 Add"); + CHECK(static_cast(c.instrs[0].op) == nn::ExprOp::Add, "dce: 保留输出 Add"); auto v = nn::validate_expr_spec(c, 3); CHECK(v.has_value(), "dce: canonical 通过 validate"); } @@ -174,18 +158,14 @@ void test_cse() // 重复子表达式 Mul(input0, input1) 出现两次 → 合并 { auto s = build_spec( - {{static_cast(nn::ExprOp::Mul), 0, nn::expr::input(0), - nn::expr::input(1), {}}, - {static_cast(nn::ExprOp::Mul), 1, nn::expr::input(0), - nn::expr::input(1), {}}, - {static_cast(nn::ExprOp::Add), 2, nn::expr::reg(0), - nn::expr::reg(1), {}}}, + {{static_cast(nn::ExprOp::Mul), 0, nn::expr::input(0), nn::expr::input(1), {}}, + {static_cast(nn::ExprOp::Mul), 1, nn::expr::input(0), nn::expr::input(1), {}}, + {static_cast(nn::ExprOp::Add), 2, nn::expr::reg(0), nn::expr::reg(1), {}}}, {nn::expr::linear(), nn::expr::linear()}, {}); auto c = nn::canonicalize_expr_spec(s); CHECK(c.instrs.size() == 2, "cse: 两条相同 Mul 合并为一条"); // Add 的两个操作数应引用同一寄存器 - CHECK(c.instrs[1].a == c.instrs[1].b, - "cse: Add 两个操作数引用同一子表达式"); + CHECK(c.instrs[1].a == c.instrs[1].b, "cse: Add 两个操作数引用同一子表达式"); auto v = nn::validate_expr_spec(c, 2); CHECK(v.has_value(), "cse: canonical 通过 validate"); } @@ -193,15 +173,14 @@ void test_cse() // CSE 与 DCE 交互：死代码的重复子表达式也应被消除 { auto s = build_spec( - {{static_cast(nn::ExprOp::Mul), 0, nn::expr::input(0), - nn::expr::input(1), {}}, - {static_cast(nn::ExprOp::Mul), 1, nn::expr::input(0), - nn::expr::input(1), {}}, - {static_cast(nn::ExprOp::Add), 2, nn::expr::input(2), - nn::expr::input(3), {}}}, - {nn::expr::linear(), nn::expr::linear(), nn::expr::linear(), - nn::expr::linear()}, - {}); + {{static_cast(nn::ExprOp::Mul), 0, nn::expr::input(0), nn::expr::input(1), {}}, + {static_cast(nn::ExprOp::Mul), 1, nn::expr::input(0), nn::expr::input(1), {}}, + {static_cast(nn::ExprOp::Add), + 2, + nn::expr::input(2), + nn::expr::input(3), + {}}}, + {nn::expr::linear(), nn::expr::linear(), nn::expr::linear(), nn::expr::linear()}, {}); auto c = nn::canonicalize_expr_spec(s); CHECK(c.instrs.size() == 1, "cse+dce: 死代码 Mul 全部消除"); } @@ -214,24 +193,18 @@ void test_regalloc() // r0 = in0*in1; r1 = in2*in3; r2 = r0+r1; r3 = in0*in4; r4 = r2+r3 (out) { auto s = build_spec( - {{static_cast(nn::ExprOp::Mul), 0, nn::expr::input(0), - nn::expr::input(1), {}}, - {static_cast(nn::ExprOp::Mul), 1, nn::expr::input(2), - nn::expr::input(3), {}}, - {static_cast(nn::ExprOp::Add), 2, nn::expr::reg(0), - nn::expr::reg(1), {}}, - {static_cast(nn::ExprOp::Mul), 3, nn::expr::input(0), - nn::expr::input(4), {}}, - {static_cast(nn::ExprOp::Add), 4, nn::expr::reg(2), - nn::expr::reg(3), {}}}, - {nn::expr::linear(), nn::expr::linear(), nn::expr::linear(), - nn::expr::linear(), nn::expr::linear()}, + {{static_cast(nn::ExprOp::Mul), 0, nn::expr::input(0), nn::expr::input(1), {}}, + {static_cast(nn::ExprOp::Mul), 1, nn::expr::input(2), nn::expr::input(3), {}}, + {static_cast(nn::ExprOp::Add), 2, nn::expr::reg(0), nn::expr::reg(1), {}}, + {static_cast(nn::ExprOp::Mul), 3, nn::expr::input(0), nn::expr::input(4), {}}, + {static_cast(nn::ExprOp::Add), 4, nn::expr::reg(2), nn::expr::reg(3), {}}}, + {nn::expr::linear(), nn::expr::linear(), nn::expr::linear(), nn::expr::linear(), + nn::expr::linear()}, {}); const std::uint32_t orig_regs = s.num_regs; auto c = nn::canonicalize_expr_spec(s); CHECK(c.num_regs < orig_regs, "regalloc: 寄存器复用降低 num_regs"); - CHECK(c.num_regs > 0 && c.num_regs <= nn::EXPR_MAX_REGS, - "regalloc: num_regs 合法"); + CHECK(c.num_regs > 0 && c.num_regs <= nn::EXPR_MAX_REGS, "regalloc: num_regs 合法"); auto v = nn::validate_expr_spec(c, 5); CHECK(v.has_value(), "regalloc: canonical 通过 validate"); } @@ -256,8 +229,8 @@ void test_determinism() // 2) softmax 混合（归约视图 + 归约指令） { nn::Tensor x = make_tensor(eng, 4, 6); - auto expr = exp(leaf(x) - row_reduce_max(x)) - / row_reduce_sum(exp(leaf(x) - row_reduce_max(x))); + auto expr = + exp(leaf(x) - row_reduce_max(x)) / row_reduce_sum(exp(leaf(x) - row_reduce_max(x))); specs.push_back(to_expr_spec(expr).first); } // 3) RMSNorm 风格（列归约 + 广播 + 标量） @@ -270,14 +243,10 @@ void test_determinism() // 4) 手工构造：常量折叠 + CSE + 复用机会 { specs.push_back(build_spec( - {{static_cast(nn::ExprOp::Mul), 0, nn::expr::input(0), - nn::expr::cst(0), {}}, - {static_cast(nn::ExprOp::Add), 1, nn::expr::input(0), - nn::expr::input(1), {}}, - {static_cast(nn::ExprOp::Mul), 2, nn::expr::reg(1), - nn::expr::cst(1), {}}, - {static_cast(nn::ExprOp::Add), 3, nn::expr::reg(2), - nn::expr::input(0), {}}}, + {{static_cast(nn::ExprOp::Mul), 0, nn::expr::input(0), nn::expr::cst(0), {}}, + {static_cast(nn::ExprOp::Add), 1, nn::expr::input(0), nn::expr::input(1), {}}, + {static_cast(nn::ExprOp::Mul), 2, nn::expr::reg(1), nn::expr::cst(1), {}}, + {static_cast(nn::ExprOp::Add), 3, nn::expr::reg(2), nn::expr::input(0), {}}}, {nn::expr::linear(), nn::expr::linear()}, {Scalar{0}, Scalar{1}})); } @@ -286,8 +255,7 @@ void test_determinism() const nn::ExprSpec c1 = nn::canonicalize_expr_spec(specs[i]); const nn::ExprSpec c2 = nn::canonicalize_expr_spec(specs[i]); CHECK(nn::expr_spec_equal(c1, c2), "determinism: 两次 canonicalize 结果一致"); - CHECK(nn::expr_spec_key(c1) == nn::expr_spec_key(c2), - "determinism: canonical key 稳定"); + CHECK(nn::expr_spec_key(c1) == nn::expr_spec_key(c2), "determinism: canonical key 稳定"); // views 不变量：canonical 不改变输入视图 CHECK(c1.views == specs[i].views, "determinism: views 保持不变"); // canonical 输出合法 @@ -309,8 +277,8 @@ void test_semantic_equivalence() { const std::size_t R = 3, C = 7; nn::Tensor x = make_tensor(eng, R, C, -0.8f, 0.11f); - auto expr = exp(leaf(x) - row_reduce_max(x)) - / row_reduce_sum(exp(leaf(x) - row_reduce_max(x))); + auto expr = + exp(leaf(x) - row_reduce_max(x)) / row_reduce_sum(exp(leaf(x) - row_reduce_max(x))); auto [spec, inputs] = to_expr_spec(expr); const nn::ExprSpec canon = nn::canonicalize_expr_spec(spec); @@ -420,14 +388,10 @@ void test_pressure() const nn::ExprSpec c = nn::canonicalize_expr_spec(s); auto v = nn::validate_expr_spec(c, 8); CHECK(v.has_value(), "pressure: canonical 通过 validate"); - CHECK(c.instrs.size() <= nn::EXPR_MAX_INSTRS, - "pressure: 指令数 ≤ EXPR_MAX_INSTRS"); - CHECK(c.num_regs <= nn::EXPR_MAX_REGS, - "pressure: 寄存器数 ≤ EXPR_MAX_REGS"); - CHECK(c.views.size() <= nn::EXPR_MAX_INPUTS, - "pressure: 输入数 ≤ EXPR_MAX_INPUTS"); - CHECK(c.consts.size() <= nn::EXPR_MAX_CONSTS, - "pressure: 常量数 ≤ EXPR_MAX_CONSTS"); + CHECK(c.instrs.size() <= nn::EXPR_MAX_INSTRS, "pressure: 指令数 ≤ EXPR_MAX_INSTRS"); + CHECK(c.num_regs <= nn::EXPR_MAX_REGS, "pressure: 寄存器数 ≤ EXPR_MAX_REGS"); + CHECK(c.views.size() <= nn::EXPR_MAX_INPUTS, "pressure: 输入数 ≤ EXPR_MAX_INPUTS"); + CHECK(c.consts.size() <= nn::EXPR_MAX_CONSTS, "pressure: 常量数 ≤ EXPR_MAX_CONSTS"); } } @@ -437,8 +401,7 @@ void test_output_invariant() // 输出 = x*0：即使可化简为常量 0，输出指令必须保留（glsl_gen 输出 r） { auto s = build_spec( - {{static_cast(nn::ExprOp::Mul), 0, nn::expr::input(0), - nn::expr::cst(0), {}}}, + {{static_cast(nn::ExprOp::Mul), 0, nn::expr::input(0), nn::expr::cst(0), {}}}, {nn::expr::linear()}, {Scalar{0}}); auto c = nn::canonicalize_expr_spec(s); CHECK(!c.instrs.empty(), "output: 输出指令保留（不为空）"); @@ -456,8 +419,7 @@ void test_reduce_preserved() auto eng_owner = nn::CpuEngine::create(); nn::CpuEngine& eng = *eng_owner; nn::Tensor x = make_tensor(eng, 3, 7, -0.8f, 0.11f); - auto expr = exp(leaf(x) - row_reduce_max(x)) - / row_reduce_sum(exp(leaf(x) - row_reduce_max(x))); + auto expr = exp(leaf(x) - row_reduce_max(x)) / row_reduce_sum(exp(leaf(x) - row_reduce_max(x))); auto [spec, inputs] = to_expr_spec(expr); const nn::ExprSpec c = nn::canonicalize_expr_spec(spec); auto v = nn::validate_expr_spec(c, inputs.size()); @@ -476,8 +438,7 @@ void test_reduce_preserved() for (const auto& r : c.instrs) if (nn::expr_op_is_reduce(static_cast(r.op))) for (const auto& e : c.instrs) - if (!nn::expr_op_is_reduce(static_cast(e.op)) && - r.dst == e.dst) + if (!nn::expr_op_is_reduce(static_cast(e.op)) && r.dst == e.dst) disjoint = false; CHECK(disjoint, "reduce: 归约 dst 与逐元素 dst 互斥（区段分离）"); } @@ -491,27 +452,34 @@ int main() std::printf("========================================\n"); test_fold(); - std::printf("[run] test_fold done\n"); std::fflush(stdout); + std::printf("[run] test_fold done\n"); + std::fflush(stdout); test_dce(); - std::printf("[run] test_dce done\n"); std::fflush(stdout); + std::printf("[run] test_dce done\n"); + std::fflush(stdout); test_cse(); - std::printf("[run] test_cse done\n"); std::fflush(stdout); + std::printf("[run] test_cse done\n"); + std::fflush(stdout); test_regalloc(); - std::printf("[run] test_regalloc done\n"); std::fflush(stdout); + std::printf("[run] test_regalloc done\n"); + std::fflush(stdout); test_determinism(); - std::printf("[run] test_determinism done\n"); std::fflush(stdout); + std::printf("[run] test_determinism done\n"); + std::fflush(stdout); test_semantic_equivalence(); - std::printf("[run] test_semantic_equivalence done\n"); std::fflush(stdout); + std::printf("[run] test_semantic_equivalence done\n"); + std::fflush(stdout); test_pressure(); - std::printf("[run] test_pressure done\n"); std::fflush(stdout); + std::printf("[run] test_pressure done\n"); + std::fflush(stdout); test_output_invariant(); - std::printf("[run] test_output_invariant done\n"); std::fflush(stdout); + std::printf("[run] test_output_invariant done\n"); + std::fflush(stdout); test_reduce_preserved(); - std::printf("[run] test_reduce_preserved done\n"); std::fflush(stdout); + std::printf("[run] test_reduce_preserved done\n"); + std::fflush(stdout); - if (g_fail == 0) - std::printf("\n[expr_opt_test] 全部通过 ✅\n"); - else - std::printf("\n[expr_opt_test] %d 项失败 ❌\n", g_fail); + if (g_fail == 0) std::printf("\n[expr_opt_test] 全部通过 ✅\n"); + else std::printf("\n[expr_opt_test] %d 项失败 ❌\n", g_fail); return g_fail == 0 ? 0 : 1; } \ No newline at end of file diff --git a/src/expr_reduce_test.cpp b/src/expr_reduce_test.cpp index 2c8330e8..9558bafe 100644 --- a/src/expr_reduce_test.cpp +++ b/src/expr_reduce_test.cpp @@ -27,7 +27,7 @@ // 测试写在全局作用域（非 namespace nn），避免 ADL 把匹配拉进 nn 命名空间。 using namespace nn::dsl; -namespace dsl = nn::dsl; // 别名：全局作用域可用 dsl::leaf 等限定调用 +namespace dsl = nn::dsl; // 别名：全局作用域可用 dsl::leaf 等限定调用 using nn::Scalar; namespace @@ -36,7 +36,6 @@ namespace int g_fail = 0; - // ── 测试 1：归约视图——按行求和广播 x / row_sum ───────────────────────── void test_row_sum_broadcast() { @@ -56,8 +55,10 @@ void test_row_sum_broadcast() for (std::size_t r = 0; r < R; ++r) { float sum = 0.0f; - for (std::size_t c = 0; c < C; ++c) sum += xs[r * C + c]; - for (std::size_t c = 0; c < C; ++c) rs[r * C + c] = xs[r * C + c] / sum; + for (std::size_t c = 0; c < C; ++c) + sum += xs[r * C + c]; + for (std::size_t c = 0; c < C; ++c) + rs[r * C + c] = xs[r * C + c] / sum; } check_close(eng, *out, upload(eng, ref), "row_reduce_sum 广播 (x / row_sum)", g_fail); } @@ -80,8 +81,10 @@ void test_row_max_shift() for (std::size_t r = 0; r < R; ++r) { float mx = xs[r * C]; - for (std::size_t c = 1; c < C; ++c) mx = std::max(mx, xs[r * C + c]); - for (std::size_t c = 0; c < C; ++c) rs[r * C + c] = xs[r * C + c] - mx; + for (std::size_t c = 1; c < C; ++c) + mx = std::max(mx, xs[r * C + c]); + for (std::size_t c = 0; c < C; ++c) + rs[r * C + c] = xs[r * C + c] - mx; } check_close(eng, *out, upload(eng, ref), "row_reduce_max 广播 (x - row_max)", g_fail); } @@ -96,8 +99,8 @@ void test_softmax_mixed() const std::size_t R = 3, C = 7; nn::Tensor x = make_tensor(eng, R, C, -0.8f, 0.11f); - auto expr = dsl::exp(dsl::leaf(x) - dsl::row_reduce_max(x)) - / dsl::row_reduce_sum(dsl::exp(dsl::leaf(x) - dsl::row_reduce_max(x))); + auto expr = dsl::exp(dsl::leaf(x) - dsl::row_reduce_max(x)) / + dsl::row_reduce_sum(dsl::exp(dsl::leaf(x) - dsl::row_reduce_max(x))); auto out = dsl::compute(eng, expr, R, C); CHECK(out.has_value(), "softmax_mixed: compute failed"); @@ -108,10 +111,13 @@ void test_softmax_mixed() for (std::size_t r = 0; r < R; ++r) { float mx = xs[r * C]; - for (std::size_t c = 1; c < C; ++c) mx = std::max(mx, xs[r * C + c]); + for (std::size_t c = 1; c < C; ++c) + mx = std::max(mx, xs[r * C + c]); float sum = 0.0f; - for (std::size_t c = 0; c < C; ++c) sum += std::exp(xs[r * C + c] - mx); - for (std::size_t c = 0; c < C; ++c) rs[r * C + c] = std::exp(xs[r * C + c] - mx) / sum; + for (std::size_t c = 0; c < C; ++c) + sum += std::exp(xs[r * C + c] - mx); + for (std::size_t c = 0; c < C; ++c) + rs[r * C + c] = std::exp(xs[r * C + c] - mx) / sum; } check_close(eng, *out, upload(eng, ref), "Softmax（视图 row_max + 指令 row_sum）", g_fail); } @@ -135,8 +141,10 @@ void test_reduce_instr_output() for (std::size_t r = 0; r < R; ++r) { float sum = 0.0f; - for (std::size_t c = 0; c < C; ++c) sum += std::exp(xs[r * C + c]); - for (std::size_t c = 0; c < C; ++c) rs[r * C + c] = sum; + for (std::size_t c = 0; c < C; ++c) + sum += std::exp(xs[r * C + c]); + for (std::size_t c = 0; c < C; ++c) + rs[r * C + c] = sum; } check_close(eng, *out, upload(eng, ref), "归约指令 row_sum(exp(x)) 输出广播", g_fail); } @@ -159,8 +167,10 @@ void test_col_max_broadcast() for (std::size_t c = 0; c < C; ++c) { float mx = xs[c]; - for (std::size_t r = 1; r < R; ++r) mx = std::max(mx, xs[r * C + c]); - for (std::size_t r = 0; r < R; ++r) rs[r * C + c] = xs[r * C + c] * mx; + for (std::size_t r = 1; r < R; ++r) + mx = std::max(mx, xs[r * C + c]); + for (std::size_t r = 0; r < R; ++r) + rs[r * C + c] = xs[r * C + c] * mx; } check_close(eng, *out, upload(eng, ref), "col_reduce_max 广播 (x * col_max)", g_fail); } @@ -179,11 +189,14 @@ void test_reduce_of_reduce_view() const nn::Matrix xm = download(eng, x); float total = 0.0f; - for (auto v : xm.span()) total += v; + for (auto v : xm.span()) + total += v; nn::Matrix ref(R, C); - for (auto& v : ref.span()) v = total; + for (auto& v : ref.span()) + v = total; - check_close(eng, *out, upload(eng, ref), "归约指令叠加归约视图 row_sum(col_sum(x)) = 全局和", g_fail); + check_close(eng, *out, upload(eng, ref), "归约指令叠加归约视图 row_sum(col_sum(x)) = 全局和", + g_fail); } // ── 测试 7：归约与标量混合（均值平移）x - row_mean ───────────────────── @@ -196,8 +209,8 @@ void test_reduce_with_scalar() const float inv = 1.0f / static_cast(C); // x - (row_sum(x) / C) —— 归约视图 + 标量除法 - auto out = dsl::compute(eng, - dsl::leaf(x) - (dsl::row_reduce_sum(x) * nn::dsl::ConstLeaf{inv}), R, C); + auto out = + dsl::compute(eng, dsl::leaf(x) - (dsl::row_reduce_sum(x) * nn::dsl::ConstLeaf{inv}), R, C); CHECK(out.has_value(), "reduce_with_scalar: compute failed"); nn::Matrix ref(R, C); @@ -207,9 +220,11 @@ void test_reduce_with_scalar() for (std::size_t r = 0; r < R; ++r) { float sum = 0.0f; - for (std::size_t c = 0; c < C; ++c) sum += xs[r * C + c]; + for (std::size_t c = 0; c < C; ++c) + sum += xs[r * C + c]; const float mean = sum * inv; - for (std::size_t c = 0; c < C; ++c) rs[r * C + c] = xs[r * C + c] - mean; + for (std::size_t c = 0; c < C; ++c) + rs[r * C + c] = xs[r * C + c] - mean; } check_close(eng, *out, upload(eng, ref), "归约视图 + 标量 (x - row_mean)", g_fail); } @@ -245,7 +260,8 @@ void test_fold_structure() "fold: 指令1 应为 RowSum 归约指令"); CHECK(spec.instrs[1].a.kind == static_cast(nn::ExprOperandKind::Reg), "fold: RowSum 源应为寄存器"); - CHECK(spec.views.size() == 1 && spec.views[0].kind == static_cast(nn::ExprViewKind::Linear), + CHECK(spec.views.size() == 1 && + spec.views[0].kind == static_cast(nn::ExprViewKind::Linear), "fold: row_sum(exp(x)) 视图应为 Linear"); } @@ -264,10 +280,16 @@ void test_validation() { nn::ExprSpec spec; spec.num_regs = 2; - spec.instrs.push_back({static_cast(nn::ExprOp::RowSum), 0, - nn::expr::input(0), {}, {}}); // r0 = RowSum(input0) - spec.instrs.push_back({static_cast(nn::ExprOp::Neg), 1, - nn::expr::reg(0), {}, {}}); // r1 = -r0 ← 非法：Reg 引用归约 dst + spec.instrs.push_back({static_cast(nn::ExprOp::RowSum), + 0, + nn::expr::input(0), + {}, + {}}); // r0 = RowSum(input0) + spec.instrs.push_back({static_cast(nn::ExprOp::Neg), + 1, + nn::expr::reg(0), + {}, + {}}); // r1 = -r0 ← 非法：Reg 引用归约 dst spec.views.push_back(nn::expr::linear()); auto v = nn::validate_expr_spec(spec, 1); CHECK(!v.has_value(), "validate: Reg 引用归约 dst 应被拒绝"); @@ -277,10 +299,16 @@ void test_validation() { nn::ExprSpec spec; spec.num_regs = 2; - spec.instrs.push_back({static_cast(nn::ExprOp::Neg), 0, - nn::expr::input(0), {}, {}}); // r0 = -input0（逐元素） - spec.instrs.push_back({static_cast(nn::ExprOp::Mul), 1, - nn::expr::reduce(0), nn::expr::reg(0), {}}); // Reduce(0) 但 r0 非归约 + spec.instrs.push_back({static_cast(nn::ExprOp::Neg), + 0, + nn::expr::input(0), + {}, + {}}); // r0 = -input0（逐元素） + spec.instrs.push_back({static_cast(nn::ExprOp::Mul), + 1, + nn::expr::reduce(0), + nn::expr::reg(0), + {}}); // Reduce(0) 但 r0 非归约 spec.views.push_back(nn::expr::linear()); auto v = nn::validate_expr_spec(spec, 1); CHECK(!v.has_value(), "validate: Reduce 引用非归约 dst 应被拒绝"); @@ -290,8 +318,11 @@ void test_validation() { nn::ExprSpec spec; spec.num_regs = 1; - spec.instrs.push_back({static_cast(nn::ExprOp::RowSum), 0, - nn::expr::reduce(0), {}, {}}); // RowSum 源引用自己 + spec.instrs.push_back({static_cast(nn::ExprOp::RowSum), + 0, + nn::expr::reduce(0), + {}, + {}}); // RowSum 源引用自己 spec.views.push_back(nn::expr::linear()); auto v = nn::validate_expr_spec(spec, 1); CHECK(!v.has_value(), "validate: Reduce 自引用应被拒绝"); @@ -305,21 +336,17 @@ void test_key_determinism() nn::CpuEngine& eng = *eng_owner; nn::Tensor x = eng.create_tensor(4, 5); - auto key_row = nn::expr_spec_key(to_expr_spec( - dsl::leaf(x) / dsl::row_reduce_sum(x)).first); - auto key_row2 = nn::expr_spec_key(to_expr_spec( - dsl::leaf(x) / dsl::row_reduce_sum(x)).first); - auto key_col = nn::expr_spec_key(to_expr_spec( - dsl::leaf(x) / dsl::col_reduce_sum(x)).first); + auto key_row = nn::expr_spec_key(to_expr_spec(dsl::leaf(x) / dsl::row_reduce_sum(x)).first); + auto key_row2 = nn::expr_spec_key(to_expr_spec(dsl::leaf(x) / dsl::row_reduce_sum(x)).first); + auto key_col = nn::expr_spec_key(to_expr_spec(dsl::leaf(x) / dsl::col_reduce_sum(x)).first); CHECK(key_row == key_row2, "key: 同结构同 key"); CHECK(key_row != key_col, "key: 行/列归约 key 不同"); // 归约指令 vs 归约视图：row_sum(exp(x)) 与 exp(x) 的 key 不同 - auto key_instr = nn::expr_spec_key(to_expr_spec( - dsl::row_reduce_sum(dsl::exp(dsl::leaf(x)))).first); - auto key_exp = nn::expr_spec_key(to_expr_spec( - dsl::exp(dsl::leaf(x))).first); + auto key_instr = + nn::expr_spec_key(to_expr_spec(dsl::row_reduce_sum(dsl::exp(dsl::leaf(x)))).first); + auto key_exp = nn::expr_spec_key(to_expr_spec(dsl::exp(dsl::leaf(x))).first); CHECK(key_instr != key_exp, "key: 归约指令参与哈希"); } @@ -338,9 +365,7 @@ int main() test_validation(); test_key_determinism(); - if (g_fail == 0) - std::printf("\n[expr_reduce_test] 全部通过 ✅\n"); - else - std::printf("\n[expr_reduce_test] %d 项失败 ❌\n", g_fail); + if (g_fail == 0) std::printf("\n[expr_reduce_test] 全部通过 ✅\n"); + else std::printf("\n[expr_reduce_test] %d 项失败 ❌\n", g_fail); return g_fail == 0 ? 0 : 1; } \ No newline at end of file diff --git a/src/f16_compute_test.cpp b/src/f16_compute_test.cpp index 62bf6500..bf12c094 100644 --- a/src/f16_compute_test.cpp +++ b/src/f16_compute_test.cpp @@ -34,8 +34,10 @@ void test_f16_matmul_vs_f32() nn::Matrix b32(8, 3); std::mt19937 rng(42); std::uniform_real_distribution dist(-1.0f, 1.0f); - for (auto& v : a32.span()) v = dist(rng); - for (auto& v : b32.span()) v = dist(rng); + for (auto& v : a32.span()) + v = dist(rng); + for (auto& v : b32.span()) + v = dist(rng); // f32 matmul 参考结果 nn::Matrix c32_ref(4, 3); @@ -87,7 +89,8 @@ void test_cast_roundtrip() nn::Matrix m32(3, 4); std::mt19937 rng(123); std::uniform_real_distribution dist(-10.0f, 10.0f); - for (auto& v : m32.span()) v = dist(rng); + for (auto& v : m32.span()) + v = dist(rng); auto t32 = upload(engine, m32); @@ -131,7 +134,8 @@ void test_f16_linear_training() nn::Matrix x_data(D, N); nn::Matrix w_true(D, 1); - for (auto& v : x_data.span()) v = ndist(rng); + for (auto& v : x_data.span()) + v = ndist(rng); w_true.set_value(0, 0, 2.0f); w_true.set_value(1, 0, -1.0f); w_true.set_value(2, 0, 0.5f); @@ -149,7 +153,8 @@ void test_f16_linear_training() // 初始化权重 w（随机） nn::Matrix w(D, 1); - for (auto& v : w.span()) v = ndist(rng) * 0.1f; + for (auto& v : w.span()) + v = ndist(rng) * 0.1f; float lr = 0.01f; float prev_loss = 1e10f; @@ -172,25 +177,25 @@ void test_f16_linear_training() CHECK(t_pred32_r->is_cpu(), "cast result is CPU"); // loss = mean((pred - y)^2) — 在 f32 融合世界中计算 - auto t_diff_r = nn::dsl::compute(engine, - nn::dsl::leaf(*t_pred32_r) - nn::dsl::leaf(t_y32), - t_pred32_r->rows(), t_pred32_r->cols()); + auto t_diff_r = nn::dsl::compute(engine, nn::dsl::leaf(*t_pred32_r) - nn::dsl::leaf(t_y32), + t_pred32_r->rows(), t_pred32_r->cols()); CHECK(t_diff_r.has_value(), "sub"); - auto t_loss_r = nn::dsl::compute(engine, - nn::dsl::leaf(*t_diff_r) * nn::dsl::leaf(*t_diff_r), - t_diff_r->rows(), t_diff_r->cols()); + auto t_loss_r = + nn::dsl::compute(engine, nn::dsl::leaf(*t_diff_r) * nn::dsl::leaf(*t_diff_r), + t_diff_r->rows(), t_diff_r->cols()); CHECK(t_loss_r.has_value(), "sqr"); // 简化：取所有元素的平均作为 loss（通过 to_matrix 读回 CPU） auto loss_mat = engine.to_matrix(*t_loss_r); CHECK(loss_mat.has_value(), "to_matrix for loss"); float loss = 0.0f; - for (auto& v : loss_mat->span()) loss += v; + for (auto& v : loss_mat->span()) + loss += v; loss /= static_cast(N); // 反向：grad_w = 2/N * X @ (pred - y) (f16 matmul) // scale 在 f32 融合世界完成，然后 cast 回 f16 做 matmul - auto t_scale_r = nn::dsl::compute(engine, - nn::dsl::leaf(*t_diff_r) * nn::dsl::rparam(2.0f / static_cast(N)), + auto t_scale_r = nn::dsl::compute( + engine, nn::dsl::leaf(*t_diff_r) * nn::dsl::rparam(2.0f / static_cast(N)), t_diff_r->rows(), t_diff_r->cols()); CHECK(t_scale_r.has_value(), "scale grad"); // grad_w = X @ scaled_diff: X(D,N)^T * scaled_diff(N,1) = (D,1) @@ -235,8 +240,10 @@ void test_f16_batched_matmul() nn::Matrix b32(batch * K, N); // 注意：batched_matmul 要求 b_rows_per = K std::mt19937 rng(77); std::uniform_real_distribution dist(-1.0f, 1.0f); - for (auto& v : a32.span()) v = dist(rng); - for (auto& v : b32.span()) v = dist(rng); + for (auto& v : a32.span()) + v = dist(rng); + for (auto& v : b32.span()) + v = dist(rng); // f32 batched_matmul 参考 auto t_a32 = upload(engine, a32); @@ -250,8 +257,8 @@ void test_f16_batched_matmul() auto t_b16 = engine.cast(t_b32, nn::Precision::F16); CHECK(t_b16.has_value(), "cast B to f16"); - auto t_c16_r = engine.batched_matmul(*t_a16, *t_b16, batch, false, false, 1.0f, - nn::Precision::F16); + auto t_c16_r = + engine.batched_matmul(*t_a16, *t_b16, batch, false, false, 1.0f, nn::Precision::F16); CHECK(t_c16_r.has_value(), "f16 batched_matmul"); // cast 回 f32 比较 diff --git a/src/f16_cpu_probe.cpp b/src/f16_cpu_probe.cpp index f4c46512..a79c2007 100644 --- a/src/f16_cpu_probe.cpp +++ b/src/f16_cpu_probe.cpp @@ -46,18 +46,18 @@ void scan(nn::ComputeEngine& eng, const nn::Tensor& t, const char* tag) double mx = 0.0; for (auto v : m->span()) mx = std::max(mx, std::fabs(static_cast(v))); - std::printf(" [%s] finite ok (%s) max|v|=%.6g%s\n", tag, - t.shape_str().c_str(), mx, mx > 100.0 ? " <<< 超大" : ""); + std::printf(" [%s] finite ok (%s) max|v|=%.6g%s\n", tag, t.shape_str().c_str(), mx, + mx > 100.0 ? " <<< 超大" : ""); std::fflush(stdout); } -nn::Matrix rnd(std::size_t rows, std::size_t cols, unsigned seed, - float lo = -0.5f, float hi = 0.5f) +nn::Matrix rnd(std::size_t rows, std::size_t cols, unsigned seed, float lo = -0.5f, float hi = 0.5f) { nn::Matrix m(rows, cols); std::mt19937 rng(seed); std::uniform_real_distribution d(lo, hi); - for (auto& v : m.span()) v = d(rng); + for (auto& v : m.span()) + v = d(rng); return m; } } // namespace @@ -67,7 +67,7 @@ int main() std::printf("f16_cpu_probe（profile_f16 组件级定位）\n"); auto raw_owner = nn::CpuEngine::create(); nn::CpuEngine& raw = *raw_owner; - nn::ComputeEngine& ad = raw; // P-1：边界 cast 已下沉基类（原 PrecisionEngine） + nn::ComputeEngine& ad = raw; // P-1：边界 cast 已下沉基类（原 PrecisionEngine） const auto prof = nn::profile_f16(); constexpr std::size_t D = 16, OUT = 32, SEQ = 4, BATCH = 2, TOTAL = SEQ * BATCH; @@ -87,15 +87,27 @@ int main() NN_EXIT(x16, 1, " from_matrix FAIL\n"); banner("A1c Linear.forward(x=f16)"); auto y = ln.forward(*x16); - if (!y) { std::printf(" forward ERR: %s\n", y.error().message.c_str()); ++g_bad; } - else { std::printf(" y precision=%s\n", y->precision() == nn::Precision::F16 ? "f16" : "f32"); - scan(ad, *y, "A1 out"); } + if (!y) + { + std::printf(" forward ERR: %s\n", y.error().message.c_str()); + ++g_bad; + } + else + { + std::printf(" y precision=%s\n", + y->precision() == nn::Precision::F16 ? "f16" : "f32"); + scan(ad, *y, "A1 out"); + } banner("A2 Linear.forward (W=f16, x=f32)"); auto x32 = ad.from_matrix(rnd(D, TOTAL, 2)); NN_EXIT(x32, 1, " from_matrix FAIL\n"); auto y2 = ln.forward(*x32); - if (!y2) { std::printf(" forward ERR: %s\n", y2.error().message.c_str()); ++g_bad; } + if (!y2) + { + std::printf(" forward ERR: %s\n", y2.error().message.c_str()); + ++g_bad; + } else scan(ad, *y2, "A2 out"); banner("A3 Linear.backward (grad=f16)"); @@ -106,9 +118,14 @@ int main() if (auto z = ad.zero(ln.param_gradients()[1].get()); !z) std::printf(" zero grad FAIL\n"); auto gi = ln.backward(*g); - if (!gi) { std::printf(" backward ERR: %s\n", gi.error().message.c_str()); ++g_bad; } + if (!gi) + { + std::printf(" backward ERR: %s\n", gi.error().message.c_str()); + ++g_bad; + } else scan(ad, *gi, "A3 grad_in"); - for (auto& gt : ln.param_gradients()) scan(ad, gt.get(), "A3 param_grad"); + for (auto& gt : ln.param_gradients()) + scan(ad, gt.get(), "A3 param_grad"); } // ── B. LayerNorm forward（stable=F32，输入 f16）─────────────────────── @@ -116,11 +133,19 @@ int main() banner("B LayerNorm.forward (gamma/beta=f16, x=f16, stable=f32)"); nn::LayerNorm ln(D); ln.set_precision_profile(prof); - if (auto r = ln.init(ad); !r) { std::printf(" init FAIL\n"); return 1; } + if (auto r = ln.init(ad); !r) + { + std::printf(" init FAIL\n"); + return 1; + } auto x16 = ad.from_matrix(rnd(D, TOTAL, 4), nn::Precision::F16); NN_EXIT(x16, 1, " from_matrix FAIL\n"); auto y = ln.forward(*x16); - if (!y) { std::printf(" forward ERR: %s\n", y.error().message.c_str()); ++g_bad; } + if (!y) + { + std::printf(" forward ERR: %s\n", y.error().message.c_str()); + ++g_bad; + } else scan(ad, *y, "B out"); } @@ -129,11 +154,19 @@ int main() banner("C FeedForward.forward (x=f16, compute=f16)"); nn::FeedForward ff(D, 2 * D); ff.set_precision_profile(prof); - if (auto r = ff.init(ad); !r) { std::printf(" init FAIL\n"); return 1; } + if (auto r = ff.init(ad); !r) + { + std::printf(" init FAIL\n"); + return 1; + } auto x16 = ad.from_matrix(rnd(D, TOTAL, 5), nn::Precision::F16); NN_EXIT(x16, 1, " from_matrix FAIL\n"); auto y = ff.forward(*x16); - if (!y) { std::printf(" forward ERR: %s\n", y.error().message.c_str()); ++g_bad; } + if (!y) + { + std::printf(" forward ERR: %s\n", y.error().message.c_str()); + ++g_bad; + } else scan(ad, *y, "C out"); } @@ -142,11 +175,19 @@ int main() banner("D CausalSelfAttention.forward (x=f16, compute=f16)"); nn::CausalSelfAttention attn(D, 2, 1024, SEQ, nn::PosEncodingType::Learned); attn.set_precision_profile(prof); - if (auto r = attn.init(ad); !r) { std::printf(" init FAIL\n"); return 1; } + if (auto r = attn.init(ad); !r) + { + std::printf(" init FAIL\n"); + return 1; + } auto x16 = ad.from_matrix(rnd(D, TOTAL, 6), nn::Precision::F16); NN_EXIT(x16, 1, " from_matrix FAIL\n"); auto y = attn.forward(*x16); - if (!y) { std::printf(" forward ERR: %s\n", y.error().message.c_str()); ++g_bad; } + if (!y) + { + std::printf(" forward ERR: %s\n", y.error().message.c_str()); + ++g_bad; + } else scan(ad, *y, "D out"); } @@ -154,30 +195,52 @@ int main() { banner("E CrossEntropyLoss fwd/bwd (logits=f32, stable=f32)"); nn::GptConfig cfg{}; - cfg.vocab_size = 32; cfg.d_model = D; cfg.seq_len = SEQ; - cfg.num_heads = 2; cfg.d_ff = 2 * D; cfg.num_layers = 1; + cfg.vocab_size = 32; + cfg.d_model = D; + cfg.seq_len = SEQ; + cfg.num_heads = 2; + cfg.d_ff = 2 * D; + cfg.num_layers = 1; cfg.precision = nn::profile_f32(); auto mr = nn::build_gpt_model(ad, cfg); NN_EXIT(mr, 1, " build FAIL\n"); nn::Matrix x(SEQ, BATCH); - for (std::size_t i = 0; i < SEQ * BATCH; ++i) x.span()[i] = static_cast(i % 32); + for (std::size_t i = 0; i < SEQ * BATCH; ++i) + x.span()[i] = static_cast(i % 32); nn::Matrix tgt(32, TOTAL, 0.0f); - for (std::size_t i = 0; i < TOTAL; ++i) tgt.set_value(i % 32, i, 1.0f); + for (std::size_t i = 0; i < TOTAL; ++i) + tgt.set_value(i % 32, i, 1.0f); auto xt = ad.from_matrix(x); auto tt = ad.from_matrix(tgt); - if (!xt || !tt) { std::printf(" from_matrix FAIL\n"); return 1; } + if (!xt || !tt) + { + std::printf(" from_matrix FAIL\n"); + return 1; + } auto logits = mr->forward(*xt); - if (!logits) { std::printf(" forward ERR\n"); ++g_bad; } + if (!logits) + { + std::printf(" forward ERR\n"); + ++g_bad; + } else { scan(ad, *logits, "E logits(f32 profile)"); nn::CrossEntropyLoss ce; ce.set_precision_profile(prof); auto loss = ce.forward(ad, *logits, *tt); - if (!loss) { std::printf(" ce fwd ERR: %s\n", loss.error().message.c_str()); ++g_bad; } + if (!loss) + { + std::printf(" ce fwd ERR: %s\n", loss.error().message.c_str()); + ++g_bad; + } else std::printf(" loss=%.6f\n", *loss); auto grad = ce.backward(); - if (!grad) { std::printf(" ce bwd ERR\n"); ++g_bad; } + if (!grad) + { + std::printf(" ce bwd ERR\n"); + ++g_bad; + } else scan(ad, *grad, "E ce_grad"); } } @@ -189,24 +252,47 @@ int main() nn::Precision::F32, nn::Precision::F16}; nn::Linear ln(D, OUT); ln.set_precision_profile(prof_f); - if (auto r = ln.init(ad); !r) { std::printf(" init FAIL\n"); return 1; } - auto opt = nn::create_optimizer("adam", ad, ln.parameters(), ln.param_gradients(), - 1e-3f, 0.0f, prof_f); + if (auto r = ln.init(ad); !r) + { + std::printf(" init FAIL\n"); + return 1; + } + auto opt = nn::create_optimizer("adam", ad, ln.parameters(), ln.param_gradients(), 1e-3f, + 0.0f, prof_f); NN_EXIT(opt, 1, " create_optimizer FAIL\n"); for (int s = 0; s < 3; ++s) { auto x16 = ad.from_matrix(rnd(D, TOTAL, 100 + (unsigned)s)); auto g = ad.from_matrix(rnd(OUT, TOTAL, 200 + (unsigned)s)); - if (!x16 || !g) { std::printf(" from_matrix FAIL\n"); return 1; } + if (!x16 || !g) + { + std::printf(" from_matrix FAIL\n"); + return 1; + } auto y = ln.forward(*x16); - if (!y) { std::printf(" fwd ERR\n"); ++g_bad; break; } + if (!y) + { + std::printf(" fwd ERR\n"); + ++g_bad; + break; + } auto gi = ln.backward(*g); - if (!gi) { std::printf(" bwd ERR\n"); ++g_bad; break; } + if (!gi) + { + std::printf(" bwd ERR\n"); + ++g_bad; + break; + } if (auto r = opt->step(); !r) - { std::printf(" opt step ERR: %s\n", r.error().message.c_str()); ++g_bad; break; } + { + std::printf(" opt step ERR: %s\n", r.error().message.c_str()); + ++g_bad; + break; + } std::printf(" step %d: ", s); std::fflush(stdout); - for (auto& p : ln.parameters()) scan(ad, p.get(), "F param"); + for (auto& p : ln.parameters()) + scan(ad, p.get(), "F param"); } } @@ -214,7 +300,11 @@ int main() // （debug 构建下 NN_ASSERT 会把残余的精度错配点直接钉到行号） { using P = nn::Precision; - struct Case { const char* name; nn::PrecisionProfile prof; }; + struct Case + { + const char* name; + nn::PrecisionProfile prof; + }; const Case cases[] = { {"param=f16", {P::F16, P::F32, P::F32, P::F32}}, {"compute=f16", {P::F32, P::F16, P::F32, P::F32}}, @@ -224,11 +314,20 @@ int main() { banner(cs.name); nn::GptConfig cfg{}; - cfg.vocab_size = 32; cfg.d_model = D; cfg.seq_len = SEQ; - cfg.num_heads = 2; cfg.d_ff = 2 * D; cfg.num_layers = 2; + cfg.vocab_size = 32; + cfg.d_model = D; + cfg.seq_len = SEQ; + cfg.num_heads = 2; + cfg.d_ff = 2 * D; + cfg.num_layers = 2; cfg.precision = cs.prof; auto mr = nn::build_gpt_model(ad, cfg); - if (!mr) { std::printf(" build FAIL\n"); ++g_bad; continue; } + if (!mr) + { + std::printf(" build FAIL\n"); + ++g_bad; + continue; + } // 与 f16_precision_test::train_tiny_gpt 同构：固定初值（σ=0.05）+ adamw { std::mt19937 rng(12345); @@ -236,30 +335,51 @@ int main() for (auto& p : mr->parameters()) { nn::Matrix m(p.get().rows(), p.get().cols()); - for (auto& v : m.span()) v = nd(rng); + for (auto& v : m.span()) + v = nd(rng); if (auto r = ad.copy_from(p.get(), m); !r) - { std::printf(" copy_from FAIL\n"); ++g_bad; } + { + std::printf(" copy_from FAIL\n"); + ++g_bad; + } } } nn::Matrix x(SEQ, BATCH); - for (std::size_t i = 0; i < SEQ * BATCH; ++i) x.span()[i] = static_cast(i % 32); + for (std::size_t i = 0; i < SEQ * BATCH; ++i) + x.span()[i] = static_cast(i % 32); nn::Matrix tgt(32, TOTAL, 0.0f); - for (std::size_t i = 0; i < TOTAL; ++i) tgt.set_value(i % 32, i, 1.0f); + for (std::size_t i = 0; i < TOTAL; ++i) + tgt.set_value(i % 32, i, 1.0f); auto xt = ad.from_matrix(x); auto tt = ad.from_matrix(tgt); - if (!xt || !tt) { std::printf(" from_matrix FAIL\n"); return 1; } - auto opt = nn::create_optimizer("adamw", ad, mr->parameters(), - mr->param_gradients(), 3e-3f, 0.0f, cs.prof); - if (!opt) { std::printf(" optimizer FAIL\n"); ++g_bad; continue; } + if (!xt || !tt) + { + std::printf(" from_matrix FAIL\n"); + return 1; + } + auto opt = nn::create_optimizer("adamw", ad, mr->parameters(), mr->param_gradients(), + 3e-3f, 0.0f, cs.prof); + if (!opt) + { + std::printf(" optimizer FAIL\n"); + ++g_bad; + continue; + } nn::CrossEntropyLoss ce; ce.set_precision_profile(cs.prof); for (int s = 0; s < 6; ++s) { - if (auto r = mr->zero_grad(); !r) { std::printf(" zero_grad FAIL\n"); ++g_bad; break; } + if (auto r = mr->zero_grad(); !r) + { + std::printf(" zero_grad FAIL\n"); + ++g_bad; + break; + } // step0：zero_grad 后 grad 应全 0 —— 若已非 0/非有限则 zero 或存量有问题 if (s == 0) { - double pre_max = 0.0; bool pre_bad = false; + double pre_max = 0.0; + bool pre_bad = false; for (auto& p : mr->param_gradients()) { auto m = ad.to_matrix(p.get(), nn::Precision::F32); @@ -275,28 +395,53 @@ int main() std::fflush(stdout); } auto logits = mr->forward(*xt); - if (!logits) { std::printf(" fwd ERR: %s\n", logits.error().message.c_str()); ++g_bad; break; } + if (!logits) + { + std::printf(" fwd ERR: %s\n", logits.error().message.c_str()); + ++g_bad; + break; + } char tag[64]; std::snprintf(tag, sizeof(tag), "%s step%d logits", cs.name, s); scan(ad, *logits, tag); auto loss = ce.forward(ad, *logits, *tt); - if (!loss) { std::printf(" ce ERR\n"); ++g_bad; break; } + if (!loss) + { + std::printf(" ce ERR\n"); + ++g_bad; + break; + } auto grad = ce.backward(); - if (!grad) { std::printf(" ce bwd ERR\n"); ++g_bad; break; } + if (!grad) + { + std::printf(" ce bwd ERR\n"); + ++g_bad; + break; + } if (auto r = mr->backward(*grad); !r) - { std::printf(" bwd ERR: %s\n", r.error().message.c_str()); ++g_bad; break; } + { + std::printf(" bwd ERR: %s\n", r.error().message.c_str()); + ++g_bad; + break; + } bool bad = false; for (auto& p : mr->param_gradients()) { auto m = ad.to_matrix(p.get(), nn::Precision::F32); if (!m) continue; - double mx = 0.0; bool this_bad = false; std::size_t first = 0; + double mx = 0.0; + bool this_bad = false; + std::size_t first = 0; for (std::size_t i = 0; i < m->span().size(); ++i) { const float v = m->span()[i]; if (!std::isfinite(v)) { - if (!this_bad) { first = i; this_bad = true; } + if (!this_bad) + { + first = i; + this_bad = true; + } } else mx = std::max(mx, std::fabs(static_cast(v))); } @@ -313,8 +458,8 @@ int main() } if (mx > 50.0) { - std::printf(" [%s] step %d OUTLIER grad %s max|g|=%.6g\n", - cs.name, s, p.get().shape_str().c_str(), mx); + std::printf(" [%s] step %d OUTLIER grad %s max|g|=%.6g\n", cs.name, s, + p.get().shape_str().c_str(), mx); std::fflush(stdout); } } @@ -324,7 +469,11 @@ int main() ++g_bad; } if (auto r = opt->step(); !r) - { std::printf(" opt ERR: %s\n", r.error().message.c_str()); ++g_bad; break; } + { + std::printf(" opt ERR: %s\n", r.error().message.c_str()); + ++g_bad; + break; + } std::printf(" [%s] step %d loss=%.4f grads=%s\n", cs.name, s, *loss, bad ? "BAD" : "ok"); std::fflush(stdout); @@ -336,7 +485,11 @@ int main() // ── H0. GPTBlock 级：fwd+bwd 是否复现坏梯度 ────────────────────────── { using P = nn::Precision; - struct Case { const char* name; nn::PrecisionProfile prof; }; + struct Case + { + const char* name; + nn::PrecisionProfile prof; + }; const Case cases[] = { {"f32", nn::profile_f32()}, {"param=f16", {P::F16, P::F32, P::F32, P::F32}}, @@ -348,7 +501,12 @@ int main() banner((std::string("H0 GPTBlock ") + cs.name).c_str()); nn::GPTBlock blk(D, 2, 2 * D, 1024, SEQ, nn::PosEncodingType::Learned, nn::ActivationType::GeLU, nn::NormType::LayerNorm, cs.prof); - if (auto r = blk.init(ad); !r) { std::printf(" init FAIL\n"); ++g_bad; continue; } + if (auto r = blk.init(ad); !r) + { + std::printf(" init FAIL\n"); + ++g_bad; + continue; + } // 固定初值（同 train_tiny_gpt：σ=0.05, seed12345） { std::mt19937 rng(12345); @@ -356,36 +514,62 @@ int main() for (auto& p : blk.parameters()) { nn::Matrix m(p.get().rows(), p.get().cols()); - for (auto& v : m.span()) v = nd(rng); + for (auto& v : m.span()) + v = nd(rng); if (auto r = ad.copy_from(p.get(), m); !r) - { std::printf(" copy_from FAIL\n"); ++g_bad; } + { + std::printf(" copy_from FAIL\n"); + ++g_bad; + } } } - const nn::Precision in_p = (cs.prof.compute == nn::Precision::F16) - ? nn::Precision::F16 : nn::Precision::F32; + const nn::Precision in_p = + (cs.prof.compute == nn::Precision::F16) ? nn::Precision::F16 : nn::Precision::F32; auto x = ad.from_matrix(rnd(D, TOTAL, 42, -0.1f, 0.1f), in_p); NN_EXIT(x, 1, " from_matrix FAIL\n"); auto y = blk.forward(*x); - if (!y) { std::printf(" fwd ERR: %s\n", y.error().message.c_str()); ++g_bad; continue; } + if (!y) + { + std::printf(" fwd ERR: %s\n", y.error().message.c_str()); + ++g_bad; + continue; + } scan(ad, *y, (std::string(cs.name) + " block-y").c_str()); auto g = ad.from_matrix(rnd(D, TOTAL, 43, -0.1f, 0.1f), in_p); NN_EXIT(g, 1, " from_matrix FAIL\n"); if (auto r = blk.backward(*g); !r) - { std::printf(" bwd ERR: %s\n", r.error().message.c_str()); ++g_bad; continue; } + { + std::printf(" bwd ERR: %s\n", r.error().message.c_str()); + ++g_bad; + continue; + } int gi = 0; for (auto& p : blk.param_gradients()) { auto m = ad.to_matrix(p.get(), nn::Precision::F32); - if (!m) { ++gi; continue; } - double mx = 0.0; bool this_bad = false; std::size_t first = 0; + if (!m) + { + ++gi; + continue; + } + double mx = 0.0; + bool this_bad = false; + std::size_t first = 0; for (std::size_t i = 0; i < m->span().size(); ++i) { const float v = m->span()[i]; - if (!std::isfinite(v)) { if (!this_bad) { first = i; this_bad = true; } } + if (!std::isfinite(v)) + { + if (!this_bad) + { + first = i; + this_bad = true; + } + } else mx = std::max(mx, std::fabs(static_cast(v))); } - std::printf(" [%s] grad#%d %s [%s] %s idx=%zu max=%.6g\n", - cs.name, gi++, p.get().shape_str().c_str(), + std::printf(" [%s] grad#%d %s [%s] %s idx=%zu max=%.6g\n", cs.name, gi++, + p.get().shape_str().c_str(), p.get().precision() == nn::Precision::F16 ? "f16" : "f32", this_bad ? "BAD" : (mx > 20.0 ? "outlier" : "ok"), first, mx); if (this_bad) ++g_bad; @@ -400,46 +584,73 @@ int main() banner("I Linear.backward 拆解 (param=f16 compute=f32, x/g=f32)"); nn::Linear ln2(D, OUT); ln2.set_precision_profile({P::F16, P::F32, P::F32, P::F32}); - if (auto r = ln2.init(ad); !r) { std::printf(" init FAIL\n"); return 1; } - auto x = ad.from_matrix(rnd(D, TOTAL, 77, -0.1f, 0.1f)); // f32 - auto g = ad.from_matrix(rnd(OUT, TOTAL, 78, -0.1f, 0.1f)); // f32 - if (!x || !g) { std::printf(" from_matrix FAIL\n"); return 1; } + if (auto r = ln2.init(ad); !r) + { + std::printf(" init FAIL\n"); + return 1; + } + auto x = ad.from_matrix(rnd(D, TOTAL, 77, -0.1f, 0.1f)); // f32 + auto g = ad.from_matrix(rnd(OUT, TOTAL, 78, -0.1f, 0.1f)); // f32 + if (!x || !g) + { + std::printf(" from_matrix FAIL\n"); + return 1; + } auto y = ln2.forward(*x); - if (!y) { std::printf(" fwd ERR\n"); ++g_bad; } + if (!y) + { + std::printf(" fwd ERR\n"); + ++g_bad; + } auto params = ln2.parameters(); auto caches = ln2.activation_cache(); - const auto& w = params[0].get(); + const auto& w = params[0].get(); const auto& cache = caches[0].get(); std::printf(" w=%s %s | cache=%s %s\n", w.shape_str().c_str(), - w.precision() == nn::Precision::F16 ? "f16" : "f32", - cache.shape_str().c_str(), + w.precision() == nn::Precision::F16 ? "f16" : "f32", cache.shape_str().c_str(), cache.precision() == nn::Precision::F16 ? "f16" : "f32"); // ── 把两个临时 vector 显式化，避免 -Wdangling-reference ── - auto grads = ln2.param_gradients(); // 拷贝出一个 vector，里面的 reference_wrapper 仍指向真实张量 + auto grads = + ln2.param_gradients(); // 拷贝出一个 vector，里面的 reference_wrapper 仍指向真实张量 // step1: grad_input = matmul(w, g, transA, P=compute=F32) auto gi = ad.matmul(w, *g, true, false, P::F32); - if (!gi) { std::printf(" matmul ERR: %s\n", gi.error().message.c_str()); ++g_bad; } + if (!gi) + { + std::printf(" matmul ERR: %s\n", gi.error().message.c_str()); + ++g_bad; + } else scan(ad, *gi, "I grad_input"); // step2: grad_w += matmul(g, cache^T) （compute_into, dst=f16） auto& gw = grads[0].get(); - auto r2 = nn::dsl::compute_into(ad, - nn::dsl::leaf(gw) + nn::dsl::matmul(*g, cache, false, true), gw); - if (!r2) { std::printf(" compute_into ERR: %s\n", r2.error().message.c_str()); ++g_bad; } + auto r2 = nn::dsl::compute_into( + ad, nn::dsl::leaf(gw) + nn::dsl::matmul(*g, cache, false, true), gw); + if (!r2) + { + std::printf(" compute_into ERR: %s\n", r2.error().message.c_str()); + ++g_bad; + } else scan(ad, gw, "I grad_w (compute_into 后)"); // step3: grad_b += row_reduce_sum(g) auto gb = ad.row_reduce_sum(*g, P::F32); - if (!gb) { std::printf(" row_reduce ERR\n"); ++g_bad; } + if (!gb) + { + std::printf(" row_reduce ERR\n"); + ++g_bad; + } else { scan(ad, *gb, "I row_reduce(g)"); auto& gbb = grads[1].get(); if (auto r3 = ad.accumulate(gbb, *gb); !r3) - { std::printf(" accumulate ERR\n"); ++g_bad; } + { + std::printf(" accumulate ERR\n"); + ++g_bad; + } else scan(ad, gbb, "I grad_b (accumulate 后)"); } } diff --git a/src/f16_precision_test.cpp b/src/f16_precision_test.cpp index b7ab62e3..b67de614 100644 --- a/src/f16_precision_test.cpp +++ b/src/f16_precision_test.cpp @@ -26,16 +26,18 @@ namespace int g_failures = 0; -#define CHECK_NEAR(a, b, tol, msg) \ - do { \ - const double va = static_cast(a), vb = static_cast(b); \ - const double tol_ = (tol); \ - const double err = std::fabs(va - vb); \ - if (err > tol_ * (1.0 + std::fabs(vb))) { \ - std::fprintf(stderr, " FAIL %s:%d: %s (%.6g vs %.6g, err=%.3g)\n", \ - g_label, __LINE__, (msg), va, vb, err); \ - ++g_failures; \ - } \ +#define CHECK_NEAR(a, b, tol, msg) \ + do \ + { \ + const double va = static_cast(a), vb = static_cast(b); \ + const double tol_ = (tol); \ + const double err = std::fabs(va - vb); \ + if (err > tol_ * (1.0 + std::fabs(vb))) \ + { \ + std::fprintf(stderr, " FAIL %s:%d: %s (%.6g vs %.6g, err=%.3g)\n", g_label, __LINE__, \ + (msg), va, vb, err); \ + ++g_failures; \ + } \ } while (0) const char* g_label = "?"; @@ -55,13 +57,14 @@ bool env_flag(const char* name) #endif } -nn::Matrix make_rand(std::size_t rows, std::size_t cols, unsigned seed, - float lo = -1.0f, float hi = 1.0f) +nn::Matrix make_rand(std::size_t rows, std::size_t cols, unsigned seed, float lo = -1.0f, + float hi = 1.0f) { nn::Matrix m(rows, cols); std::mt19937 rng(seed); std::uniform_real_distribution d(lo, hi); - for (auto& v : m.span()) v = d(rng); + for (auto& v : m.span()) + v = d(rng); return m; } @@ -71,8 +74,7 @@ nn::Matrix read_f32(nn::ComputeEngine& eng, const nn::Tensor& t) auto r = eng.to_matrix(t, nn::Precision::F32); if (!r) { - std::fprintf(stderr, " FAIL %s: to_matrix 失败: %s\n", g_label, - r.error().message.c_str()); + std::fprintf(stderr, " FAIL %s: to_matrix 失败: %s\n", g_label, r.error().message.c_str()); ++g_failures; return nn::Matrix(1, 1); } @@ -108,7 +110,11 @@ void test_f32_zero_regression(nn::ComputeEngine& raw, nn::ComputeEngine& eng) bool identical = (mr.rows() == mg.rows() && mr.cols() == mg.cols()); if (identical) for (std::size_t i = 0; i < mr.size(); ++i) - if (mr.span()[i] != mg.span()[i]) { identical = false; break; } + if (mr.span()[i] != mg.span()[i]) + { + identical = false; + break; + } CHECK(identical, "f32 全 f32 配置逐字节一致（G5 零回归）"); CHECK(gots[k]->precision() == nn::Precision::F32, "全 f32 输出精度"); } @@ -124,12 +130,10 @@ void test_f32_zero_regression_dsl(nn::ComputeEngine& raw, nn::ComputeEngine& eng CHECK(ta && tb, "from_matrix"); if (!ta || !tb) return; - auto ref = nn::dsl::compute(raw, - nn::dsl::exp(nn::dsl::leaf(*ta) * nn::dsl::rparam(0.25f)) - + nn::dsl::leaf(*tb), 16, 8); - auto got = nn::dsl::compute(eng, - nn::dsl::exp(nn::dsl::leaf(*ta) * nn::dsl::rparam(0.25f)) - + nn::dsl::leaf(*tb), 16, 8); + auto ref = nn::dsl::compute( + raw, nn::dsl::exp(nn::dsl::leaf(*ta) * nn::dsl::rparam(0.25f)) + nn::dsl::leaf(*tb), 16, 8); + auto got = nn::dsl::compute( + eng, nn::dsl::exp(nn::dsl::leaf(*ta) * nn::dsl::rparam(0.25f)) + nn::dsl::leaf(*tb), 16, 8); CHECK(ref && got, "dsl::compute"); if (!ref || !got) return; @@ -139,7 +143,11 @@ void test_f32_zero_regression_dsl(nn::ComputeEngine& raw, nn::ComputeEngine& eng bool identical = (mr.rows() == mg.rows() && mr.cols() == mg.cols()); if (identical) for (std::size_t i = 0; i < mr.size(); ++i) - if (mr.span()[i] != mg.span()[i]) { identical = false; break; } + if (mr.span()[i] != mg.span()[i]) + { + identical = false; + break; + } CHECK(identical, "f32 DSL 逐字节一致（G5 零回归）"); } @@ -160,14 +168,14 @@ void test_f16_dsl(nn::ComputeEngine& raw, nn::ComputeEngine& eng) auto tb32 = eng.cast(*tb16, nn::Precision::F32); CHECK(ta32 && tb32, "cast f16→f32"); if (!ta32 || !tb32) return; - auto ref = nn::dsl::compute(raw, - nn::dsl::exp(nn::dsl::leaf(*ta32) * nn::dsl::rparam(0.25f)) - + nn::dsl::leaf(*tb32), 24, 10); + auto ref = nn::dsl::compute( + raw, nn::dsl::exp(nn::dsl::leaf(*ta32) * nn::dsl::rparam(0.25f)) + nn::dsl::leaf(*tb32), 24, + 10); CHECK(ref.has_value(), "f32 参考"); - auto got = nn::dsl::compute(eng, - nn::dsl::exp(nn::dsl::leaf(*ta16) * nn::dsl::rparam(0.25f)) - + nn::dsl::leaf(*tb16), 24, 10, nn::Precision::F16); + auto got = nn::dsl::compute( + eng, nn::dsl::exp(nn::dsl::leaf(*ta16) * nn::dsl::rparam(0.25f)) + nn::dsl::leaf(*tb16), 24, + 10, nn::Precision::F16); CHECK(got.has_value(), "f16 dsl::compute"); if (!ref || !got) return; @@ -176,12 +184,11 @@ void test_f16_dsl(nn::ComputeEngine& raw, nn::ComputeEngine& eng) const nn::Matrix mr = read_f32(raw, *ref); double max_err = 0.0; for (std::size_t i = 0; i < mr.size(); ++i) - max_err = std::max(max_err, std::fabs(static_cast(mg.span()[i]) - - static_cast(mr.span()[i]))); + max_err = std::max(max_err, std::fabs(static_cast(mg.span()[i]) - + static_cast(mr.span()[i]))); // f16 相对精度 2^-10 ≈ 1e-3，逐算子舍入累积 → 5e-3 容差 CHECK(max_err < 5e-3, "f16 DSL 与 f32 参考在容差内"); - if (max_err >= 5e-3) - std::fprintf(stderr, " max_err=%.3g\n", max_err); + if (max_err >= 5e-3) std::fprintf(stderr, " max_err=%.3g\n", max_err); } // ── 3. in-place：存储精度不可变 + 对象身份保持（§8.3）──────────────────── @@ -205,9 +212,13 @@ void test_f16_inplace(nn::ComputeEngine& raw, nn::ComputeEngine& eng) bool ok = true; for (std::size_t i = 0; i < a.size(); ++i) { - const float expect = static_cast( - static_cast(static_cast(a.span()[i]) + static_cast(b.span()[i]))); - if (std::fabs(out.span()[i] - expect) > 2e-3f) { ok = false; break; } + const float expect = static_cast(static_cast( + static_cast(a.span()[i]) + static_cast(b.span()[i]))); + if (std::fabs(out.span()[i] - expect) > 2e-3f) + { + ok = false; + break; + } } CHECK(ok, "f16 in-place 结果正确且别名可见"); @@ -222,7 +233,11 @@ void test_f16_inplace(nn::ComputeEngine& raw, nn::ComputeEngine& eng) const nn::Matrix z = read_f32(raw, *ta); bool all_zero = true; for (std::size_t i = 0; i < z.size(); ++i) - if (z.span()[i] != 0.0f) { all_zero = false; break; } + if (z.span()[i] != 0.0f) + { + all_zero = false; + break; + } CHECK(all_zero, "f16 zero 写零"); } @@ -246,7 +261,8 @@ void test_f16_move(nn::ComputeEngine& raw, nn::ComputeEngine& eng) bool ok = true; for (std::size_t r = 0; r < 3; ++r) for (std::size_t c = 0; c < 4; ++c) - if (std::fabs(ms.at(r, c) - static_cast(static_cast(a.at(r + 2, c)))) > 1e-3f) + if (std::fabs(ms.at(r, c) - + static_cast(static_cast(a.at(r + 2, c)))) > 1e-3f) ok = false; CHECK(ok, "slice_rows 数值正确"); } @@ -259,16 +275,18 @@ void test_f16_move(nn::ComputeEngine& raw, nn::ComputeEngine& eng) CHECK(ri.has_value(), "insert_rows(f16)"); CHECK(dst->precision() == nn::Precision::F16, "insert_rows 后 dst 精度不变"); const nn::Matrix md = read_f32(raw, *dst); - CHECK_NEAR(md.at(3, 1), static_cast(static_cast(a.at(3, 1))), - 1e-3, "insert_rows 数值"); + CHECK_NEAR(md.at(3, 1), static_cast(static_cast(a.at(3, 1))), 1e-3, + "insert_rows 数值"); } // gather_rows（表 f16、索引 f32） nn::Matrix tbl = make_rand(6, 3, 8); auto tt = eng.from_matrix(tbl, nn::Precision::F16); nn::Matrix idxm(4, 1); - idxm.set_value(0, 0, 3); idxm.set_value(1, 0, 0); - idxm.set_value(2, 0, 5); idxm.set_value(3, 0, 1); + idxm.set_value(0, 0, 3); + idxm.set_value(1, 0, 0); + idxm.set_value(2, 0, 5); + idxm.set_value(3, 0, 1); auto idx = raw.from_matrix(idxm); if (tt && idx) { @@ -277,8 +295,8 @@ void test_f16_move(nn::ComputeEngine& raw, nn::ComputeEngine& eng) if (g) { const nn::Matrix mg = read_f32(raw, *g); - CHECK_NEAR(mg.at(2, 2), static_cast(static_cast(tbl.at(5, 2))), - 1e-3, "gather_rows 数值"); + CHECK_NEAR(mg.at(2, 2), static_cast(static_cast(tbl.at(5, 2))), 1e-3, + "gather_rows 数值"); } } @@ -331,8 +349,8 @@ void test_f16_reduce_matmul(nn::ComputeEngine& raw, nn::ComputeEngine& eng) const nn::Matrix mr = read_f32(raw, *ref); double max_err = 0.0; for (std::size_t i = 0; i < mr.size(); ++i) - max_err = std::max(max_err, - std::fabs(static_cast(mg.span()[i]) - static_cast(mr.span()[i]))); + max_err = std::max(max_err, std::fabs(static_cast(mg.span()[i]) - + static_cast(mr.span()[i]))); CHECK(max_err < 5e-3, "f16 matmul 与 f32 参考在容差内"); if (max_err >= 5e-3) std::fprintf(stderr, " matmul max_err=%.3g\n", max_err); } @@ -346,12 +364,13 @@ void test_f16_end_to_end(nn::ComputeEngine& raw, nn::ComputeEngine& eng) { const std::size_t D = 8, N = 32; nn::Matrix x = make_rand(D, N, 21); - nn::Matrix w_true(D, 1, nn::Scalar{0.5f}); // 0.5 在 f16 中可精确表示 + nn::Matrix w_true(D, 1, nn::Scalar{0.5f}); // 0.5 在 f16 中可精确表示 nn::Matrix y(N, 1); for (std::size_t i = 0; i < N; ++i) { float s = 0.0f; - for (std::size_t d = 0; d < D; ++d) s += x.at(d, i) * w_true.at(d, 0); + for (std::size_t d = 0; d < D; ++d) + s += x.at(d, i) * w_true.at(d, 0); y.set_value(i, 0, s); } @@ -372,34 +391,62 @@ void test_f16_end_to_end(nn::ComputeEngine& raw, nn::ComputeEngine& eng) for (int step = 0; step < 40; ++step) { auto pred = eng.matmul(*tx, *tw, true, false, nn::Precision::F16); - if (!pred) { CHECK(false, "matmul forward"); return; } - auto diff = nn::dsl::compute(eng, - nn::dsl::leaf(*pred) - nn::dsl::leaf(*ty), N, 1, nn::Precision::F16); - if (!diff) { CHECK(false, "diff"); return; } + if (!pred) + { + CHECK(false, "matmul forward"); + return; + } + auto diff = nn::dsl::compute(eng, nn::dsl::leaf(*pred) - nn::dsl::leaf(*ty), N, 1, + nn::Precision::F16); + if (!diff) + { + CHECK(false, "diff"); + return; + } // loss（f32 参考读数，仅用于断言收敛） - auto sq = nn::dsl::compute(eng, - nn::dsl::leaf(*diff) * nn::dsl::leaf(*diff), N, 1, nn::Precision::F32); - if (!sq) { CHECK(false, "sq"); return; } + auto sq = nn::dsl::compute(eng, nn::dsl::leaf(*diff) * nn::dsl::leaf(*diff), N, 1, + nn::Precision::F32); + if (!sq) + { + CHECK(false, "sq"); + return; + } auto loss_sum = eng.col_reduce_sum(*sq, nn::Precision::F32); - if (!loss_sum) { CHECK(false, "loss reduce"); return; } + if (!loss_sum) + { + CHECK(false, "loss reduce"); + return; + } const nn::Matrix lm = read_f32(raw, *loss_sum); const float loss = lm.span()[0] / static_cast(N); if (step == 0) first_loss = loss; last_loss = loss; // grad_w = 2/N * x · diff - auto gdiff = nn::dsl::compute(eng, - nn::dsl::leaf(*diff) * nn::dsl::rparam(2.0f / static_cast(N)), - N, 1, nn::Precision::F16); - if (!gdiff) { CHECK(false, "gdiff"); return; } + auto gdiff = nn::dsl::compute( + eng, nn::dsl::leaf(*diff) * nn::dsl::rparam(2.0f / static_cast(N)), N, 1, + nn::Precision::F16); + if (!gdiff) + { + CHECK(false, "gdiff"); + return; + } auto grad = eng.matmul(*tx, *gdiff, false, false, nn::Precision::F16); - if (!grad) { CHECK(false, "grad matmul"); return; } + if (!grad) + { + CHECK(false, "grad matmul"); + return; + } // w -= lr*grad（f16 存储就地更新；先抬 f32 再落回） - auto upd = nn::dsl::compute_into(eng, - nn::dsl::leaf(*tw) - nn::dsl::leaf(*grad) * nn::dsl::rparam(lr), *tw); - if (!upd) { CHECK(false, "参数更新"); return; } + auto upd = nn::dsl::compute_into( + eng, nn::dsl::leaf(*tw) - nn::dsl::leaf(*grad) * nn::dsl::rparam(lr), *tw); + if (!upd) + { + CHECK(false, "参数更新"); + return; + } CHECK(tw->precision() == nn::Precision::F16, "训练后权重仍为 f16"); } CHECK(last_loss < first_loss * 0.5f, "全 f16 训练 loss 显著下降"); @@ -412,11 +459,11 @@ struct TrainResult bool ok = false; float first = 0.0f; float last = 0.0f; - std::vector losses; // 每步 loss（f32/f16 轨迹对拍用） + std::vector losses; // 每步 loss（f32/f16 轨迹对拍用） }; -TrainResult train_tiny_gpt(nn::ComputeEngine& eng, const nn::PrecisionProfile& prof, - int steps, float lr, bool verbose) +TrainResult train_tiny_gpt(nn::ComputeEngine& eng, const nn::PrecisionProfile& prof, int steps, + float lr, bool verbose) { TrainResult out; nn::GptConfig cfg{}; @@ -429,7 +476,11 @@ TrainResult train_tiny_gpt(nn::ComputeEngine& eng, const nn::PrecisionProfile& p cfg.precision = prof; auto mr = nn::build_gpt_model(eng, cfg); - if (!mr) { CHECK(false, "build_gpt_model"); return out; } + if (!mr) + { + CHECK(false, "build_gpt_model"); + return out; + } nn::Model& model = *mr; // 确定初值：把全部参数覆盖为固定伪随机值。Linear 的初值来自 @@ -441,7 +492,8 @@ TrainResult train_tiny_gpt(nn::ComputeEngine& eng, const nn::PrecisionProfile& p for (auto& p : model.parameters()) { nn::Matrix m(p.get().rows(), p.get().cols()); - for (auto& v : m.span()) v = nd(rng); + for (auto& v : m.span()) + v = nd(rng); if (auto r = eng.copy_from(p.get(), m); !r) { CHECK(false, "copy_from(固定初值)"); @@ -450,9 +502,13 @@ TrainResult train_tiny_gpt(nn::ComputeEngine& eng, const nn::PrecisionProfile& p } } - auto opt = nn::create_optimizer("adamw", eng, model.parameters(), - model.param_gradients(), lr, 0.0f, prof); - if (!opt) { CHECK(false, "create_optimizer"); return out; } + auto opt = nn::create_optimizer("adamw", eng, model.parameters(), model.param_gradients(), lr, + 0.0f, prof); + if (!opt) + { + CHECK(false, "create_optimizer"); + return out; + } const std::size_t batch = 2; const std::size_t total = batch * cfg.seq_len; @@ -464,7 +520,11 @@ TrainResult train_tiny_gpt(nn::ComputeEngine& eng, const nn::PrecisionProfile& p tgt.set_value(i % cfg.vocab_size, i, 1.0f); auto xt = eng.from_matrix(x); auto tt = eng.from_matrix(tgt); - if (!xt || !tt) { CHECK(false, "from_matrix"); return out; } + if (!xt || !tt) + { + CHECK(false, "from_matrix"); + return out; + } nn::CrossEntropyLoss ce; ce.set_precision_profile(prof); @@ -480,34 +540,60 @@ TrainResult train_tiny_gpt(nn::ComputeEngine& eng, const nn::PrecisionProfile& p if (!std::isfinite(v)) { std::fprintf(stderr, " [%s] step %d 首个非有限值: %s @%s 前 4 值 %.4g %.4g\n", - g_label, g_step, tag, t.shape_str().c_str(), - m->span()[0], m->span()[1]); + g_label, g_step, tag, t.shape_str().c_str(), m->span()[0], + m->span()[1]); return; } }; // 固定初值后先扫一遍参数：确认 NaN 不是初始化引入的 if (verbose) - for (auto& p : model.parameters()) scan("init-param", p.get()); + for (auto& p : model.parameters()) + scan("init-param", p.get()); for (int s = 0; s < steps; ++s) { g_step = s; - if (auto r = model.zero_grad(); !r) { CHECK(false, "zero_grad"); return out; } + if (auto r = model.zero_grad(); !r) + { + CHECK(false, "zero_grad"); + return out; + } auto logits = model.forward(*xt); - if (!logits) { CHECK(false, "model.forward"); return out; } + if (!logits) + { + CHECK(false, "model.forward"); + return out; + } scan("forward/logits", *logits); auto loss = ce.forward(eng, *logits, *tt); - if (!loss) { CHECK(false, "ce.forward"); return out; } + if (!loss) + { + CHECK(false, "ce.forward"); + return out; + } auto grad = ce.backward(); - if (!grad) { CHECK(false, "ce.backward"); return out; } + if (!grad) + { + CHECK(false, "ce.backward"); + return out; + } scan("ce/backward", *grad); auto bwd = model.backward(*grad); - if (!bwd) { CHECK(false, "model.backward"); return out; } + if (!bwd) + { + CHECK(false, "model.backward"); + return out; + } scan("model/backward-out", *bwd); if (verbose && s < 3) - for (auto& p : model.param_gradients()) scan("param-grad", p.get()); - if (auto r = opt->step(); !r) { CHECK(false, "optimizer.step"); return out; } + for (auto& p : model.param_gradients()) + scan("param-grad", p.get()); + if (auto r = opt->step(); !r) + { + CHECK(false, "optimizer.step"); + return out; + } if (s == 0) out.first = *loss; out.last = *loss; out.losses.push_back(*loss); @@ -530,8 +616,7 @@ TrainResult train_tiny_gpt(nn::ComputeEngine& eng, const nn::PrecisionProfile& p } if (bad || mx > 1e4) std::fprintf(stderr, " 参数异常: %s max=%.3g%s\n", - p.get().shape_str().c_str(), mx, - bad ? " (NaN/Inf)" : ""); + p.get().shape_str().c_str(), mx, bad ? " (NaN/Inf)" : ""); } break; } @@ -554,8 +639,7 @@ void test_f16_model_e2e(nn::ComputeEngine& raw, nn::ComputeEngine& eng) // 对照：原生引擎 + 全 f32 profile（现状路径） const auto r32 = train_tiny_gpt(raw, nn::profile_f32(), STEPS, LR, false); CHECK(r32.ok, "profile_f32 训练完成"); - if (r32.ok) - CHECK(r32.losses.back() <= r32.losses.front(), "profile_f32（原生引擎）loss 不增"); + if (r32.ok) CHECK(r32.losses.back() <= r32.losses.front(), "profile_f32（原生引擎）loss 不增"); // profile_f16：param + compute f16（CLI --f16 语义） // NN_F16_DEBUG=1 → 逐 step loss + 逐阶段非有限值扫描（定位 CPU f16 发散用） @@ -568,10 +652,10 @@ void test_f16_model_e2e(nn::ComputeEngine& raw, nn::ComputeEngine& eng) { // 硬失败判据：f16 训练 loss 必须全程有限（轨迹与 f32 对齐是本测试 // 前提）——出现非有限值即发散，直接失败，不做任何容忍分支。 - std::fprintf(stderr, " FAIL %s: f16 训练出现非有限 loss（回归！）" - " 轨迹: %.4f → %.4f\n", - g_label, r16.losses.empty() ? 0.0 : r16.losses.front(), - r16.losses.back()); + std::fprintf(stderr, + " FAIL %s: f16 训练出现非有限 loss（回归！）" + " 轨迹: %.4f → %.4f\n", + g_label, r16.losses.empty() ? 0.0 : r16.losses.front(), r16.losses.back()); ++g_failures; return; } @@ -586,16 +670,19 @@ void test_f16_model_e2e(nn::ComputeEngine& raw, nn::ComputeEngine& eng) max_rel = std::max(max_rel, rel); } CHECK(max_rel < 0.10, "f16 训练轨迹与 f32 逐 step 一致（<10%）"); - if (max_rel >= 0.10) - std::fprintf(stderr, " 轨迹最大相对偏差 %.3f\n", max_rel); - std::fprintf(stderr, " [%s] f32 %.4f → %.4f ; f16 %.4f → %.4f (max_rel=%.4f)\n", - g_label, r32.losses.front(), r32.losses.back(), - r16.losses.front(), r16.losses.back(), max_rel); + if (max_rel >= 0.10) std::fprintf(stderr, " 轨迹最大相对偏差 %.3f\n", max_rel); + std::fprintf(stderr, " [%s] f32 %.4f → %.4f ; f16 %.4f → %.4f (max_rel=%.4f)\n", g_label, + r32.losses.front(), r32.losses.back(), r16.losses.front(), r16.losses.back(), + max_rel); // 参数精度必须由 profile 驱动：profile_f16 下全部参数 f16、f32 下全 f32 { nn::GptConfig cfg{}; - cfg.vocab_size = 32; cfg.d_model = 16; cfg.seq_len = 4; - cfg.num_heads = 2; cfg.d_ff = 32; cfg.num_layers = 2; + cfg.vocab_size = 32; + cfg.d_model = 16; + cfg.seq_len = 4; + cfg.num_heads = 2; + cfg.d_ff = 32; + cfg.num_layers = 2; cfg.precision = nn::profile_f16(); auto m16 = nn::build_gpt_model(eng, cfg); CHECK(m16.has_value(), "build_gpt_model(profile_f16)"); @@ -641,15 +728,20 @@ void test_f16_model_e2e(nn::ComputeEngine& raw, nn::ComputeEngine& eng) void run_profile_matrix(nn::ComputeEngine& eng, bool verbose) { using P = nn::Precision; - struct Case { const char* name; nn::PrecisionProfile prof; bool healthy; }; + struct Case + { + const char* name; + nn::PrecisionProfile prof; + bool healthy; + }; const Case cases[] = { - {"f32", nn::profile_f32(), true}, - {"param=f16", {P::F16, P::F32, P::F32, P::F32}, true}, - {"compute=f16", {P::F32, P::F16, P::F32, P::F32}, true}, - {"stable=f16", {P::F32, P::F32, P::F16, P::F32}, true}, - {"opt=f16", {P::F32, P::F32, P::F32, P::F16}, false}, - {"all=f16", nn::profile_all_f16(), false}, - {"f16(CLI)", nn::profile_f16(), true}, + {"f32", nn::profile_f32(), true}, + {"param=f16", {P::F16, P::F32, P::F32, P::F32}, true}, + {"compute=f16", {P::F32, P::F16, P::F32, P::F32}, true}, + {"stable=f16", {P::F32, P::F32, P::F16, P::F32}, true}, + {"opt=f16", {P::F32, P::F32, P::F32, P::F16}, false}, + {"all=f16", nn::profile_all_f16(), false}, + {"f16(CLI)", nn::profile_f16(), true}, }; const char* saved = g_label; for (const auto& c : cases) @@ -692,15 +784,14 @@ void test_full_f16_training(nn::ComputeEngine& eng) const char* saved = g_label; g_label = "fullf16"; nn::PrecisionProfile full = nn::profile_f16(); - full.stable = nn::Precision::F16; + full.stable = nn::Precision::F16; full.optimizer = nn::Precision::F16; const TrainResult r = train_tiny_gpt(eng, full, /*steps=*/20, /*lr=*/3e-3f, /*verbose=*/false); CHECK(!r.losses.empty(), "full-f16：训练未产出 loss"); for (const float l : r.losses) CHECK(std::isfinite(l), "full-f16：loss 必须有限（NaN/inf = optimizer=f16 退化）"); - if (r.losses.size() >= 2) - CHECK(r.losses.back() < r.losses.front(), "full-f16：loss 应下降"); + if (r.losses.size() >= 2) CHECK(r.losses.back() < r.losses.front(), "full-f16：loss 应下降"); g_label = saved; } @@ -713,8 +804,8 @@ void test_full_f16_training(nn::ComputeEngine& eng) // 只让**求和**越过 65504 —— 从而把「输出溢出」与「元素级溢出」区分开。 void test_ce_loss_stable_f16_no_overflow(nn::ComputeEngine& eng) { - constexpr std::size_t C = 100; // 词表：lse ≈ ln(100) ≈ 4.6 - constexpr std::size_t N = 20000; // token：Σ ≈ -4.6·20000 ≈ -92000 > 65504 + constexpr std::size_t C = 100; // 词表：lse ≈ ln(100) ≈ 4.6 + constexpr std::size_t N = 20000; // token：Σ ≈ -4.6·20000 ≈ -92000 > 65504 nn::Matrix logits(C, N); std::mt19937 rng(777); std::uniform_real_distribution dist(-2.0f, 2.0f); @@ -736,14 +827,14 @@ void test_ce_loss_stable_f16_no_overflow(nn::ComputeEngine& eng) if (!ref) return; nn::CrossEntropyLoss ce16; - ce16.set_precision_profile({nn::Precision::F32, nn::Precision::F32, - nn::Precision::F16, nn::Precision::F32}); + ce16.set_precision_profile( + {nn::Precision::F32, nn::Precision::F32, nn::Precision::F16, nn::Precision::F32}); auto got = ce16.forward_sparse(eng, *t, labels, {}, C); CHECK(got.has_value(), "ce-f16-overflow: stable=f16 forward_sparse"); if (!got) return; - std::fprintf(stderr, " [%s] ce loss: f32 %.4f ; stable=f16 %.4f\n", - g_label, static_cast(*ref), static_cast(*got)); + std::fprintf(stderr, " [%s] ce loss: f32 %.4f ; stable=f16 %.4f\n", g_label, + static_cast(*ref), static_cast(*got)); CHECK(std::isfinite(static_cast(*got)), "stable=f16 的 loss 必须有限（(1,1) 归约输出不得落 f16 溢出为 ±inf）"); CHECK(std::fabs(static_cast(*got) - static_cast(*ref)) <= @@ -789,7 +880,7 @@ int main() { auto cpu_owner = nn::CpuEngine::create(); nn::CpuEngine& cpu = *cpu_owner; - nn::ComputeEngine& adapter = cpu; // P-1：原 PrecisionEngine 已下沉 + nn::ComputeEngine& adapter = cpu; // P-1：原 PrecisionEngine 已下沉 run_suite("cpu", cpu, adapter, /*allow_arbitrary_expr=*/true); if (env_flag("NN_F16_DEBUG")) { @@ -802,12 +893,12 @@ int main() // GPU（可选；设备不可用则跳过 GPU 部分而不是失败） { auto backend_owner = nn::GpuBackend::shared(); - nn::GpuBackend& backend = *backend_owner; + nn::GpuBackend& backend = *backend_owner; if (backend.initialize()) { auto gpu_owner = nn::GpuEngine::create(backend); nn::GpuEngine& gpu = *gpu_owner; - nn::ComputeEngine& adapter = gpu; // P-1：原 PrecisionEngine 已下沉 + nn::ComputeEngine& adapter = gpu; // P-1：原 PrecisionEngine 已下沉 run_suite("gpu", gpu, adapter, /*allow_arbitrary_expr=*/false); } else diff --git a/src/f16_writeback_probe.cpp b/src/f16_writeback_probe.cpp index 879bf205..a6ce38b6 100644 --- a/src/f16_writeback_probe.cpp +++ b/src/f16_writeback_probe.cpp @@ -27,7 +27,8 @@ using nn::Matrix; using nn::Precision; using nn::Scalar; -namespace { +namespace +{ int g_failures = 0; @@ -47,16 +48,18 @@ void probe_cast_rounding(nn::ComputeEngine& eng, const char* label) // 再加 32 个确定性伪随机值（任意值都落在两个网格点之间，都能区分两种舍入） constexpr int N_FIXED = 8, N_RAND = 32, N = N_FIXED + N_RAND; std::vector vals(N); - const Scalar fixed[N_FIXED] = {1.0f, 0.99997f, 1.00001f, 0.5f, - 0.077f, -0.077f, 0.036f, 2.0f}; - for (int i = 0; i < N_FIXED; ++i) vals[static_cast(i)] = fixed[i]; + const Scalar fixed[N_FIXED] = {1.0f, 0.99997f, 1.00001f, 0.5f, 0.077f, -0.077f, 0.036f, 2.0f}; + for (int i = 0; i < N_FIXED; ++i) + vals[static_cast(i)] = fixed[i]; { std::mt19937 rng(20261004); std::uniform_real_distribution d(-2.0f, 2.0f); - for (int i = N_FIXED; i < N; ++i) vals[static_cast(i)] = d(rng); + for (int i = N_FIXED; i < N; ++i) + vals[static_cast(i)] = d(rng); } Matrix m(1, static_cast(N)); - for (int i = 0; i < N; ++i) m.span()[static_cast(i)] = vals[static_cast(i)]; + for (int i = 0; i < N; ++i) + m.span()[static_cast(i)] = vals[static_cast(i)]; std::printf("[%s] cast f32→f16 舍入对拍（cast.comp，%d 值）:\n", label, N); @@ -79,8 +82,7 @@ void probe_cast_rounding(nn::ComputeEngine& eng, const char* label) if (got != expect) { if (n_diff < 8) - std::printf(" #%d f32=%.9g 期望(RNE)=%.9g 实际=%.9g\n", - i, v, expect, got); + std::printf(" #%d f32=%.9g 期望(RNE)=%.9g 实际=%.9g\n", i, v, expect, got); ++n_diff; } } @@ -94,7 +96,8 @@ nn::Result inplace_add_writeback(nn::ComputeEngine& eng, Precision targe { // 2×2 起始值全 1 Matrix m(2, 2); - for (auto& v : m.span()) v = 1.0f; + for (auto& v : m.span()) + v = 1.0f; auto p_r = eng.from_matrix(m, target_p); NN_TRY_CHECK(p_r); @@ -102,12 +105,12 @@ nn::Result inplace_add_writeback(nn::ComputeEngine& eng, Precision targe // delta 全 0.5，按 target_p 物化（对齐 optimizer 里 delta 精度 = p_.param） Matrix dm(2, 2); - for (auto& v : dm.span()) v = 0.5f; + for (auto& v : dm.span()) + v = 0.5f; auto d_r = eng.from_matrix(dm, target_p); NN_TRY_CHECK(d_r); // K3：compute_into(leaf(p) + leaf(delta), p) —— 目标传递 in-place - auto r = nn::dsl::compute_into(eng, - nn::dsl::leaf(p) + nn::dsl::leaf(*d_r), p); + auto r = nn::dsl::compute_into(eng, nn::dsl::leaf(p) + nn::dsl::leaf(*d_r), p); NN_TRY_CHECK(r); return eng.to_matrix(p); } @@ -121,8 +124,7 @@ void probe_inplace(nn::ComputeEngine& eng, const char* label, bool gpu) check(r32.has_value(), "f32 目标：执行成功"); if (r32) check(r32->at(0, 0) == 1.5f, - ("f32 目标：值 1.0+0.5=1.5，实际 " + - std::to_string(r32->at(0, 0))).c_str()); + ("f32 目标：值 1.0+0.5=1.5，实际 " + std::to_string(r32->at(0, 0))).c_str()); auto r16 = inplace_add_writeback(eng, Precision::F16); check(r16.has_value(), "f16 目标：执行成功"); @@ -130,8 +132,7 @@ void probe_inplace(nn::ComputeEngine& eng, const char* label, bool gpu) { const Scalar v = r16->at(0, 0); // f16 下 1.5 可精确表示 → 必须恰好 1.5 - check(v == 1.5f, - ("f16 目标：值 1.0+0.5=1.5，实际 " + std::to_string(v)).c_str()); + check(v == 1.5f, ("f16 目标：值 1.0+0.5=1.5，实际 " + std::to_string(v)).c_str()); } (void)gpu; @@ -165,7 +166,8 @@ void probe_optimizer_steps(nn::ComputeEngine& eng, const char* label, bool gpu) for (auto& p : model.parameters()) { nn::Matrix pm(p.get().rows(), p.get().cols()); - for (auto& v : pm.span()) v = nd(rng); + for (auto& v : pm.span()) + v = nd(rng); if (auto r = eng.copy_from(p.get(), pm); !r) { check(false, "copy_from(固定初值)"); @@ -174,19 +176,25 @@ void probe_optimizer_steps(nn::ComputeEngine& eng, const char* label, bool gpu) } } - auto opt = nn::create_optimizer("adam", eng, model.parameters(), - model.param_gradients(), /*lr=*/3e-3f, /*wd=*/0.0, cfg.precision); + auto opt = nn::create_optimizer("adam", eng, model.parameters(), model.param_gradients(), + /*lr=*/3e-3f, /*wd=*/0.0, cfg.precision); check(opt != nullptr, "create_optimizer(adam)"); if (!opt) return; // 参数快照（f16 下载往返后比较） - auto snapshot = [&]() -> std::vector { + auto snapshot = [&]() -> std::vector + { std::vector out; for (auto& t : model.parameters()) { auto m = eng.to_matrix(t.get(), nn::Precision::F32); - if (!m) { out.push_back(-1e30f); continue; } - for (auto v : m->span()) out.push_back(v); + if (!m) + { + out.push_back(-1e30f); + continue; + } + for (auto v : m->span()) + out.push_back(v); } return out; }; @@ -205,7 +213,11 @@ void probe_optimizer_steps(nn::ComputeEngine& eng, const char* label, bool gpu) tgt.set_value(i % cfg.vocab_size, i, 1.0f); auto xt = eng.from_matrix(x); auto tt = eng.from_matrix(tgt); - if (!xt || !tt) { check(false, "from_matrix(输入/目标)"); return; } + if (!xt || !tt) + { + check(false, "from_matrix(输入/目标)"); + return; + } nn::CrossEntropyLoss ce; ce.set_precision_profile(cfg.precision); @@ -213,19 +225,43 @@ void probe_optimizer_steps(nn::ComputeEngine& eng, const char* label, bool gpu) Scalar first_loss = -1e30f, last_loss = -1e30f; for (int step = 0; step < 8; ++step) { - if (auto r = model.zero_grad(); !r) { check(false, "zero_grad"); return; } + if (auto r = model.zero_grad(); !r) + { + check(false, "zero_grad"); + return; + } auto logits = model.forward(*xt); - if (!logits) { check(false, "model.forward"); return; } + if (!logits) + { + check(false, "model.forward"); + return; + } auto loss_r = ce.forward(eng, *logits, *tt); - if (!loss_r) { check(false, "ce.forward"); return; } + if (!loss_r) + { + check(false, "ce.forward"); + return; + } if (step == 0) first_loss = *loss_r; last_loss = *loss_r; auto grad = ce.backward(); - if (!grad) { check(false, "ce.backward"); return; } + if (!grad) + { + check(false, "ce.backward"); + return; + } auto bwd = model.backward(*grad); - if (!bwd) { check(false, "model.backward"); return; } + if (!bwd) + { + check(false, "model.backward"); + return; + } auto st = opt->step(); - if (!st) { check(false, "optimizer.step"); return; } + if (!st) + { + check(false, "optimizer.step"); + return; + } } auto after = snapshot(); @@ -242,8 +278,7 @@ void probe_optimizer_steps(nn::ComputeEngine& eng, const char* label, bool gpu) if (d > max_delta) max_delta = d; } std::printf(" loss %.4f → %.4f（8 步）\n", first_loss, last_loss); - std::printf(" 参数变化: %zu/%zu，最大位移 %.6g\n", - n_changed, n_total, max_delta); + std::printf(" 参数变化: %zu/%zu，最大位移 %.6g\n", n_changed, n_total, max_delta); check(n_changed > 0, "8 步 Adam 后参数有移动（f16 GPU 写回生效）"); } @@ -266,48 +301,52 @@ void probe_decay_writeback(nn::ComputeEngine& eng, const char* label) const Scalar factor = Scalar{1} - LR * WD; // 0.99997 const Scalar init[6] = {1.0f, 0.5f, 0.077f, -0.077f, 0.036f, 2.0f}; - const auto make_init = [&]() { + const auto make_init = [&]() + { Matrix m(1, 6); - for (int i = 0; i < 6; ++i) m.span()[i] = init[i]; + for (int i = 0; i < 6; ++i) + m.span()[i] = init[i]; return m; }; // 下载 f32 值序列（f16 自动升 cast） - const auto snapshot = [&](const nn::Tensor& t) { + const auto snapshot = [&](const nn::Tensor& t) + { std::vector out; auto m = eng.to_matrix(t, Precision::F32); - if (!m) { out.assign(1, -1e30f); return out; } + if (!m) + { + out.assign(1, -1e30f); + return out; + } out.assign(m->span().begin(), m->span().end()); return out; }; // 逐位比较，打印最大相对变化；返回发生变化的元素数 - const auto compare_bits = [&](const std::vector& before, - const std::vector& after, - const char* what) { + const auto compare_bits = + [&](const std::vector& before, const std::vector& after, const char* what) + { std::size_t n_diff = 0; Scalar max_rel = 0.0f; for (std::size_t i = 0; i < before.size() && i < after.size(); ++i) { if (before[i] == after[i]) continue; ++n_diff; - const Scalar denom = std::fabs(before[i]) > 1e-30f - ? std::fabs(before[i]) : 1e-30f; + const Scalar denom = std::fabs(before[i]) > 1e-30f ? std::fabs(before[i]) : 1e-30f; const Scalar rel = std::fabs(after[i] - before[i]) / denom; if (rel > max_rel) max_rel = rel; } - if (n_diff == 0) - check(true, what); + if (n_diff == 0) check(true, what); else { // 打印失败样例（首个变化元素 + 最大相对变化）便于定位 for (std::size_t i = 0; i < before.size() && i < after.size(); ++i) if (before[i] != after[i]) { - std::printf(" 首个变化元素 #%zu: %.7g → %.7g\n", - i, before[i], after[i]); + std::printf(" 首个变化元素 #%zu: %.7g → %.7g\n", i, before[i], after[i]); break; } - std::printf(" 变化 %zu/%zu 元素，最大相对变化 %.4g\n", - n_diff, before.size(), max_rel); + std::printf(" 变化 %zu/%zu 元素，最大相对变化 %.4g\n", n_diff, before.size(), + max_rel); check(false, what); } return n_diff; @@ -323,12 +362,19 @@ void probe_decay_writeback(nn::ComputeEngine& eng, const char* label) auto p = std::move(*p_r); for (int s = 0; s < STEPS; ++s) { - auto r = nn::dsl::compute_into(eng, - nn::dsl::leaf(p) * nn::dsl::rparam(factor), p); - if (!r) { check(false, "f32 对照：decay compute_into"); return; } + auto r = nn::dsl::compute_into(eng, nn::dsl::leaf(p) * nn::dsl::rparam(factor), p); + if (!r) + { + check(false, "f32 对照：decay compute_into"); + return; + } } auto m_after = eng.to_matrix(p); - if (!m_after) { check(false, "f32 对照：to_matrix"); return; } + if (!m_after) + { + check(false, "f32 对照：to_matrix"); + return; + } // 理论收缩 = 1 - factor^N ≈ N·lr·wd = 1.49%（N=500） const Scalar rel = 1.0f - m_after->at(0, 0) / init[0]; const bool ok = (rel > 0.005f && rel < 0.03f); @@ -348,13 +394,15 @@ void probe_decay_writeback(nn::ComputeEngine& eng, const char* label) check(before.size() == 6 && before[0] != -1e30f, "f16 表达式：快照可下载"); for (int s = 0; s < STEPS; ++s) { - auto r = nn::dsl::compute_into(eng, - nn::dsl::leaf(p) * nn::dsl::rparam(factor), p); - if (!r) { check(false, "f16 表达式：decay compute_into"); return; } + auto r = nn::dsl::compute_into(eng, nn::dsl::leaf(p) * nn::dsl::rparam(factor), p); + if (!r) + { + check(false, "f16 表达式：decay compute_into"); + return; + } } const auto after = snapshot(p); - compare_bits(before, after, - "f16 表达式：500 步 decay 写回逐位不变（RNE 舍回原网格点）"); + compare_bits(before, after, "f16 表达式：500 步 decay 写回逐位不变（RNE 舍回原网格点）"); } // ── (c) AdamW(wd>0) + 零梯度：唯一作用在权重上的是 decay 写回 ─────── @@ -363,15 +411,16 @@ void probe_decay_writeback(nn::ComputeEngine& eng, const char* label) check(p_r.has_value(), "AdamW decay：from_matrix(p)"); if (!p_r) return; Matrix gz(1, 6); - for (auto& v : gz.span()) v = 0.0f; // 零梯度 → Adam 更新量恰为 0 + for (auto& v : gz.span()) + v = 0.0f; // 零梯度 → Adam 更新量恰为 0 auto g_r = eng.from_matrix(gz, Precision::F16); check(g_r.has_value(), "AdamW decay：from_matrix(g=0)"); if (!g_r) return; std::vector params{std::ref(*p_r)}; std::vector grads{std::ref(*g_r)}; - const nn::PrecisionProfile prof{Precision::F16, Precision::F32, - Precision::F32, Precision::F32}; + const nn::PrecisionProfile prof{Precision::F16, Precision::F32, Precision::F32, + Precision::F32}; auto opt = nn::create_optimizer("adamw", eng, params, grads, LR, WD, prof); check(opt != nullptr, "AdamW decay：create_optimizer(adamw, wd=0.1)"); if (!opt) return; @@ -381,11 +430,15 @@ void probe_decay_writeback(nn::ComputeEngine& eng, const char* label) for (int s = 0; s < STEPS; ++s) { auto r = opt->step(); - if (!r) { check(false, "AdamW decay：optimizer.step"); return; } + if (!r) + { + check(false, "AdamW decay：optimizer.step"); + return; + } } const auto after = snapshot(*p_r); compare_bits(before, after, - "f16 + AdamW(wd>0) 零梯度：500 步后权重逐位不变（decay 写回 RNE）"); + "f16 + AdamW(wd>0) 零梯度：500 步后权重逐位不变（decay 写回 RNE）"); } } @@ -401,14 +454,17 @@ int main(int argc, char* argv[]) for (int i = 1; i < argc; ++i) { std::string arg = argv[i]; - if (auto dev = nn::cli::parse_gpu_option(argc, argv, i, true)) - gpu_device = *dev; + if (auto dev = nn::cli::parse_gpu_option(argc, argv, i, true)) gpu_device = *dev; else if (arg.rfind("--gpu", 0) == 0) { // --gpu 无值 = 自动选卡 if (arg == "--gpu") gpu_device.clear(); } - else { std::fprintf(stderr, "未知参数: %s\n", arg.c_str()); return 2; } + else + { + std::fprintf(stderr, "未知参数: %s\n", arg.c_str()); + return 2; + } } std::printf("f16_writeback_probe（issue #13 P0-②：GPU f16 权重冻结定位）\n"); @@ -428,7 +484,7 @@ int main(int argc, char* argv[]) #ifdef NN_HAS_VULKAN { auto backend_owner = nn::GpuBackend::shared(); - nn::GpuBackend& backend = *backend_owner; + nn::GpuBackend& backend = *backend_owner; auto init_r = backend.initialize(gpu_device); if (!init_r) { @@ -447,12 +503,17 @@ int main(int argc, char* argv[]) #else // 纯 CPU 构建：本探针的核心价值是 GPU f16 写回（issue #13 P0-②）， // 无 Vulkan 时无可测对象 → ctest SKIP（同 gpu_test 约定） - (void)argc; (void)argv; + (void)argc; + (void)argv; std::printf("[SKIP] 此程序需要 Vulkan SDK 支持（GPU f16 写回探针）。\n"); return 77; #endif - if (g_failures == 0) { std::printf("ALL PASSED\n"); return 0; } + if (g_failures == 0) + { + std::printf("ALL PASSED\n"); + return 0; + } std::printf("%d FAILURES\n", g_failures); return 1; } diff --git a/src/fused_gpu_test.cpp b/src/fused_gpu_test.cpp index 1a5d23c6..505791cd 100644 --- a/src/fused_gpu_test.cpp +++ b/src/fused_gpu_test.cpp @@ -41,11 +41,11 @@ int main() #else #include -using nn::Matrix; -using nn::Tensor; using nn::CpuEngine; -using nn::GpuEngine; using nn::GpuBackend; +using nn::GpuEngine; +using nn::Matrix; +using nn::Tensor; // ── 辅助：最大绝对误差 ──────────────────────────────────────────────────── Scalar max_abs_diff(const Matrix& a, const Matrix& b) @@ -71,7 +71,8 @@ int run_rope(CpuEngine& cpu, GpuEngine& gpu, std::size_t dk, bool backward) const std::size_t seq = 9; // 非 2 的幂，检验索引通用性 const std::size_t rows = 2 * dk; // d_k 的整数倍 Matrix q(rows, seq); - for (auto& v : q.span()) v = dist(rng); + for (auto& v : q.span()) + v = dist(rng); nn::RotaryEmbedding rope_cpu(dk); nn::RotaryEmbedding rope_gpu(dk); @@ -88,9 +89,8 @@ int run_rope(CpuEngine& cpu, GpuEngine& gpu, std::size_t dk, bool backward) const Scalar err = max_abs_diff(download(cpu, *cr), *gm); const bool ok = err < 1e-4f; - std::cout << "[" << (ok ? "PASS" : "FAIL") << "] rope " - << (backward ? "backward" : "forward") << " dk" << dk - << " err=" << std::scientific << std::setprecision(2) << err << "\n"; + std::cout << "[" << (ok ? "PASS" : "FAIL") << "] rope " << (backward ? "backward" : "forward") + << " dk" << dk << " err=" << std::scientific << std::setprecision(2) << err << "\n"; return ok ? 0 : 1; } @@ -102,11 +102,14 @@ int run_swiglu(CpuEngine& cpu, GpuEngine& gpu) std::mt19937 rng(555); std::uniform_real_distribution dist(-1.0f, 1.0f); Matrix input(2 * d_ff, batch), grad(d_ff, batch); - for (auto& v : input.span()) v = dist(rng); - for (auto& v : grad.span()) v = dist(rng); + for (auto& v : input.span()) + v = dist(rng); + for (auto& v : grad.span()) + v = dist(rng); nn::SwiGLU sg_cpu(d_ff), sg_gpu(d_ff); - (void)sg_cpu.init(cpu); (void)sg_gpu.init(gpu); // M6 段 C：层由 init 绑定引擎 + (void)sg_cpu.init(cpu); + (void)sg_gpu.init(gpu); // M6 段 C：层由 init 绑定引擎 const Tensor in_cpu = upload(cpu, input); const Tensor in_gpu = upload(gpu, input); auto fc = sg_cpu.forward(in_cpu); @@ -121,8 +124,7 @@ int run_swiglu(CpuEngine& cpu, GpuEngine& gpu) NN_EXIT(bc, 1, " CPU backward 失败: "); if (!bg) { - std::cerr << " GPU backward 失败（未命中融合 shader？）: " - << bg.error().message << "\n"; + std::cerr << " GPU backward 失败（未命中融合 shader？）: " << bg.error().message << "\n"; return 1; } auto gm = gpu.to_matrix(*bg); @@ -144,11 +146,14 @@ int run_gelu(CpuEngine& cpu, GpuEngine& gpu) std::mt19937 rng(999); std::uniform_real_distribution dist(-1.0f, 1.0f); Matrix x(R, C), grad(R, C); - for (auto& v : x.span()) v = dist(rng); - for (auto& v : grad.span()) v = dist(rng); + for (auto& v : x.span()) + v = dist(rng); + for (auto& v : grad.span()) + v = dist(rng); nn::GeLU gl_cpu, gl_gpu; - (void)gl_cpu.init(cpu); (void)gl_gpu.init(gpu); // M6 段 C：层由 init 绑定引擎 + (void)gl_cpu.init(cpu); + (void)gl_gpu.init(gpu); // M6 段 C：层由 init 绑定引擎 const Tensor in_cpu = upload(cpu, x); const Tensor in_gpu = upload(gpu, x); auto fc = gl_cpu.forward(in_cpu); @@ -156,8 +161,8 @@ int run_gelu(CpuEngine& cpu, GpuEngine& gpu) NN_EXIT(fc, 1, " CPU gelu forward 失败: "); if (!fg) { - std::cerr << " GPU gelu forward 失败（未命中融合 shader？）: " - << fg.error().message << "\n"; + std::cerr << " GPU gelu forward 失败（未命中融合 shader？）: " << fg.error().message + << "\n"; return 1; } const Scalar err_f = max_abs_diff(download(cpu, *fc), *gpu.to_matrix(*fg)); @@ -169,8 +174,8 @@ int run_gelu(CpuEngine& cpu, GpuEngine& gpu) NN_EXIT(bc, 1, " CPU gelu backward 失败: "); if (!bg) { - std::cerr << " GPU gelu backward 失败（未命中融合 shader？）: " - << bg.error().message << "\n"; + std::cerr << " GPU gelu backward 失败（未命中融合 shader？）: " << bg.error().message + << "\n"; return 1; } const Scalar err_b = max_abs_diff(download(cpu, *bc), *gpu.to_matrix(*bg)); @@ -194,11 +199,14 @@ int run_softmax(CpuEngine& cpu, GpuEngine& gpu) std::mt19937 rng(777); std::uniform_real_distribution dist(-1.0f, 1.0f); Matrix x(R, C), grad(R, C); - for (auto& v : x.span()) v = dist(rng); - for (auto& v : grad.span()) v = dist(rng); + for (auto& v : x.span()) + v = dist(rng); + for (auto& v : grad.span()) + v = dist(rng); nn::Softmax sm_cpu, sm_gpu; - (void)sm_cpu.init(cpu); (void)sm_gpu.init(gpu); // M6 段 C：层由 init 绑定引擎 + (void)sm_cpu.init(cpu); + (void)sm_gpu.init(gpu); // M6 段 C：层由 init 绑定引擎 const Tensor in_cpu = upload(cpu, x); const Tensor in_gpu = upload(gpu, x); auto fc = sm_cpu.forward(in_cpu); @@ -206,8 +214,8 @@ int run_softmax(CpuEngine& cpu, GpuEngine& gpu) NN_EXIT(fc, 1, " CPU softmax forward 失败: "); if (!fg) { - std::cerr << " GPU softmax forward 失败（未命中归约融合 shader？）: " - << fg.error().message << "\n"; + std::cerr << " GPU softmax forward 失败（未命中归约融合 shader？）: " << fg.error().message + << "\n"; return 1; } auto fgm = gpu.to_matrix(*fg); @@ -240,15 +248,16 @@ int run_softmax(CpuEngine& cpu, GpuEngine& gpu) // ── 归一化层通用：GPU（归约融合 shader）vs CPU ──────────────────────── // 驱动同一个 Layer（RMSNorm/LayerNorm）在 CPU 与 GPU 引擎上，对比 // forward/backward 输出。未命中融合 shader 时 GPU 端 eval_expr 硬报错。 -template -int run_norm(const char* name, CpuEngine& cpu, GpuEngine& gpu) +template int run_norm(const char* name, CpuEngine& cpu, GpuEngine& gpu) { const std::size_t F = 8, B = 5; std::mt19937 rng(888); 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); + for (auto& v : x.span()) + v = dist(rng); + for (auto& v : grad.span()) + v = dist(rng); NormT n_cpu(F), n_gpu(F); NN_EXIT(n_cpu.init(cpu), 1, " CPU "); @@ -260,8 +269,8 @@ int run_norm(const char* name, CpuEngine& cpu, GpuEngine& gpu) NN_EXIT(fc, 1, " CPU "); if (!fg) { - std::cerr << " GPU " << name << " forward 失败（未命中归约融合 shader？）: " - << fg.error().message << "\n"; + std::cerr << " GPU " << name + << " forward 失败（未命中归约融合 shader？）: " << fg.error().message << "\n"; return 1; } auto fgm = gpu.to_matrix(*fg); @@ -278,8 +287,8 @@ int run_norm(const char* name, CpuEngine& cpu, GpuEngine& gpu) NN_EXIT(bc, 1, " CPU "); if (!bg) { - std::cerr << " GPU " << name << " backward 失败（未命中归约融合 shader？）: " - << bg.error().message << "\n"; + std::cerr << " GPU " << name + << " backward 失败（未命中归约融合 shader？）: " << bg.error().message << "\n"; return 1; } auto bgm = gpu.to_matrix(*bg); @@ -301,21 +310,24 @@ int run_matmul(CpuEngine& cpu, GpuEngine& gpu) std::uniform_real_distribution dist(-1.0f, 1.0f); const std::size_t M = 5, K = 7, N = 4; Matrix A(M, K), B(K, N), bias(M, N); - for (auto& v : A.span()) v = dist(rng); - for (auto& v : B.span()) v = dist(rng); - for (auto& v : bias.span()) v = dist(rng); + for (auto& v : A.span()) + v = dist(rng); + for (auto& v : B.span()) + v = dist(rng); + for (auto& v : bias.span()) + v = dist(rng); // 出生绑定：双引擎对拍输入各上传一份（绑 cpu 的喂 gpu = 硬错误） const Tensor At_cpu = upload(cpu, A), At_gpu = upload(gpu, A); const Tensor Bt_cpu = upload(cpu, B), Bt_gpu = upload(gpu, B); const Tensor biast_cpu = upload(cpu, bias), biast_gpu = upload(gpu, bias); // matmul(A,B) + bias + relu（DSL 折叠 → matmul 段 + 尾链） - auto cr = nn::dsl::compute(cpu, - nn::dsl::max(nn::dsl::matmul(At_cpu, Bt_cpu) + nn::dsl::leaf(biast_cpu), Scalar{0}), - M, N); - auto gr = nn::dsl::compute(gpu, - nn::dsl::max(nn::dsl::matmul(At_gpu, Bt_gpu) + nn::dsl::leaf(biast_gpu), Scalar{0}), - M, N); + auto cr = nn::dsl::compute( + cpu, nn::dsl::max(nn::dsl::matmul(At_cpu, Bt_cpu) + nn::dsl::leaf(biast_cpu), Scalar{0}), M, + N); + auto gr = nn::dsl::compute( + gpu, nn::dsl::max(nn::dsl::matmul(At_gpu, Bt_gpu) + nn::dsl::leaf(biast_gpu), Scalar{0}), M, + N); NN_EXIT(cr, 1, " CPU matmul 融合求值失败: "); if (!gr) { @@ -335,10 +347,9 @@ int run_matmul(CpuEngine& cpu, GpuEngine& gpu) auto gp = nn::dsl::compute(gpu, nn::dsl::matmul(At_gpu, Bt_gpu), M, N); if (!cp || !gp) { - std::cerr << " 纯 matmul 求值失败（" - << (!cp ? ("CPU: " + cp.error().message) : "") - << (!cp && !gp ? "；" : "") - << (!gp ? ("GPU: " + gp.error().message) : "") << "）\n"; + std::cerr << " 纯 matmul 求值失败（" << (!cp ? ("CPU: " + cp.error().message) : "") + << (!cp && !gp ? "；" : "") << (!gp ? ("GPU: " + gp.error().message) : "") + << "）\n"; return 1; } auto gpm = gpu.to_matrix(*gp); @@ -351,17 +362,20 @@ int run_matmul(CpuEngine& cpu, GpuEngine& gpu) // 形状无关：不同 K（含非对齐）命中同一融合 shader（mm_k 运行时填充） const std::size_t M3 = 3, K3 = 6, N3 = 2; Matrix A3(M3, K3), B3(K3, N3), bias3(M3, N3); - for (auto& v : A3.span()) v = dist(rng); - for (auto& v : B3.span()) v = dist(rng); - for (auto& v : bias3.span()) v = dist(rng); + for (auto& v : A3.span()) + v = dist(rng); + for (auto& v : B3.span()) + v = dist(rng); + for (auto& v : bias3.span()) + v = dist(rng); const Tensor A3t_cpu = upload(cpu, A3), A3t_gpu = upload(gpu, A3); const Tensor B3t_cpu = upload(cpu, B3), B3t_gpu = upload(gpu, B3); const Tensor b3t_cpu = upload(cpu, bias3), b3t_gpu = upload(gpu, bias3); - auto c3 = nn::dsl::compute(cpu, - nn::dsl::max(nn::dsl::matmul(A3t_cpu, B3t_cpu) + nn::dsl::leaf(b3t_cpu), Scalar{0}), + auto c3 = nn::dsl::compute( + cpu, nn::dsl::max(nn::dsl::matmul(A3t_cpu, B3t_cpu) + nn::dsl::leaf(b3t_cpu), Scalar{0}), M3, N3); - auto g3 = nn::dsl::compute(gpu, - nn::dsl::max(nn::dsl::matmul(A3t_gpu, B3t_gpu) + nn::dsl::leaf(b3t_gpu), Scalar{0}), + auto g3 = nn::dsl::compute( + gpu, nn::dsl::max(nn::dsl::matmul(A3t_gpu, B3t_gpu) + nn::dsl::leaf(b3t_gpu), Scalar{0}), M3, N3); if (!c3 || !g3) { @@ -379,16 +393,19 @@ int run_matmul(CpuEngine& cpu, GpuEngine& gpu) // b (out,1) 行广播 → 输出 (out,batch)）。与 Linear::forward 完全相同的表达式。 const std::size_t in_f = 6, out_f = 4, batch = 3; Matrix Wm(out_f, in_f), Xm(in_f, batch), bm(out_f, 1); - for (auto& v : Wm.span()) v = dist(rng); - for (auto& v : Xm.span()) v = dist(rng); - for (auto& v : bm.span()) v = dist(rng); + for (auto& v : Wm.span()) + v = dist(rng); + for (auto& v : Xm.span()) + v = dist(rng); + for (auto& v : bm.span()) + v = dist(rng); const Tensor Wt_cpu = upload(cpu, Wm), Wt_gpu = upload(gpu, Wm); const Tensor Xt_cpu = upload(cpu, Xm), Xt_gpu = upload(gpu, Xm); const Tensor bt_cpu = upload(cpu, bm), bt_gpu = upload(gpu, bm); - auto cl = nn::dsl::compute(cpu, - nn::dsl::matmul(Wt_cpu, Xt_cpu) + nn::dsl::row_broadcast(bt_cpu), out_f, batch); - auto gl = nn::dsl::compute(gpu, - nn::dsl::matmul(Wt_gpu, Xt_gpu) + nn::dsl::row_broadcast(bt_gpu), out_f, batch); + auto cl = nn::dsl::compute( + cpu, nn::dsl::matmul(Wt_cpu, Xt_cpu) + nn::dsl::row_broadcast(bt_cpu), out_f, batch); + auto gl = nn::dsl::compute( + gpu, nn::dsl::matmul(Wt_gpu, Xt_gpu) + nn::dsl::row_broadcast(bt_gpu), out_f, batch); if (!cl || !gl) { std::cerr << " Linear 结构 matmul+row_broadcast 求值失败（CPU/GPU）\n"; @@ -415,19 +432,22 @@ int run_matmul_reduce(CpuEngine& cpu, GpuEngine& gpu) std::uniform_real_distribution dist(-1.0f, 1.0f); const std::size_t M = 5, K = 7, N = 4; // 非 16 倍数：分块/归约边界覆盖 Matrix Qm(M, K), Km(N, K), rm_m(M, 1); // K 存储 (N,K)（transB=1） - for (auto& v : Qm.span()) v = dist(rng); - for (auto& v : Km.span()) v = dist(rng); - for (auto& v : rm_m.span()) v = dist(rng); + for (auto& v : Qm.span()) + v = dist(rng); + for (auto& v : Km.span()) + v = dist(rng); + for (auto& v : rm_m.span()) + v = dist(rng); // 出生绑定：双引擎对拍输入各上传一份（绑 cpu 的喂 gpu = 硬错误） const Tensor Qt_cpu = upload(cpu, Qm), Qt_gpu = upload(gpu, Qm); const Tensor Kt_cpu = upload(cpu, Km), Kt_gpu = upload(gpu, Km); const Tensor rmt_cpu = upload(cpu, rm_m), rmt_gpu = upload(gpu, rm_m); // row_max(matmul(Q, K^T))：行归约 max（输出 (M,1)） - auto cr = nn::dsl::compute(cpu, - nn::dsl::row_reduce_max(nn::dsl::matmul(Qt_cpu, Kt_cpu, false, true)), M, N); - auto gr = nn::dsl::compute(gpu, - nn::dsl::row_reduce_max(nn::dsl::matmul(Qt_gpu, Kt_gpu, false, true)), M, N); + auto cr = nn::dsl::compute( + cpu, nn::dsl::row_reduce_max(nn::dsl::matmul(Qt_cpu, Kt_cpu, false, true)), M, N); + auto gr = nn::dsl::compute( + gpu, nn::dsl::row_reduce_max(nn::dsl::matmul(Qt_gpu, Kt_gpu, false, true)), M, N); if (!cr || !gr) { std::cerr << " row_max(matmul) 求值失败（CPU/GPU）: " @@ -442,13 +462,15 @@ int run_matmul_reduce(CpuEngine& cpu, GpuEngine& gpu) << " err=" << std::scientific << std::setprecision(2) << err1 << "\n"; // denom = row_sum(exp(matmul(Q,K^T) - rb(row_max)))：输出 (M,1) - auto cd = nn::dsl::compute(cpu, - nn::dsl::row_reduce_sum(nn::dsl::exp( - nn::dsl::matmul(Qt_cpu, Kt_cpu, false, true) - nn::dsl::row_broadcast(rmt_cpu))), + auto cd = nn::dsl::compute( + cpu, + nn::dsl::row_reduce_sum(nn::dsl::exp(nn::dsl::matmul(Qt_cpu, Kt_cpu, false, true) - + nn::dsl::row_broadcast(rmt_cpu))), M, N); - auto gd = nn::dsl::compute(gpu, - nn::dsl::row_reduce_sum(nn::dsl::exp( - nn::dsl::matmul(Qt_gpu, Kt_gpu, false, true) - nn::dsl::row_broadcast(rmt_gpu))), + auto gd = nn::dsl::compute( + gpu, + nn::dsl::row_reduce_sum(nn::dsl::exp(nn::dsl::matmul(Qt_gpu, Kt_gpu, false, true) - + nn::dsl::row_broadcast(rmt_gpu))), M, N); if (!cd || !gd) { @@ -467,10 +489,10 @@ int run_matmul_reduce(CpuEngine& cpu, GpuEngine& gpu) // 覆盖 matmul+列归约组合（生成器按元素 batch 分解），此对拍锁死 GPU 与 // CPU 语义一致（列归约遍历全部 rows，含所有 batch）。广播输出 // （dsl::compute → vector_out=0）： - auto ccol = nn::dsl::compute(cpu, - nn::dsl::col_reduce_max(nn::dsl::matmul(Qt_cpu, Kt_cpu, false, true)), M, N); - auto gcol = nn::dsl::compute(gpu, - nn::dsl::col_reduce_max(nn::dsl::matmul(Qt_gpu, Kt_gpu, false, true)), M, N); + auto ccol = nn::dsl::compute( + cpu, nn::dsl::col_reduce_max(nn::dsl::matmul(Qt_cpu, Kt_cpu, false, true)), M, N); + auto gcol = nn::dsl::compute( + gpu, nn::dsl::col_reduce_max(nn::dsl::matmul(Qt_gpu, Kt_gpu, false, true)), M, N); if (!ccol || !gcol) { std::cerr << " col_max(matmul) 求值失败（CPU/GPU）: " @@ -485,10 +507,10 @@ int run_matmul_reduce(CpuEngine& cpu, GpuEngine& gpu) << " err=" << std::scientific << std::setprecision(2) << err3 << "\n"; // 归约向量输出（dsl::compute_reduce → vector_out=1，输出 (1,N)） - auto ccvr = nn::dsl::compute_reduce(cpu, - nn::dsl::col_reduce_max(nn::dsl::matmul(Qt_cpu, Kt_cpu, false, true)), M, N); - auto gcvr = nn::dsl::compute_reduce(gpu, - nn::dsl::col_reduce_max(nn::dsl::matmul(Qt_gpu, Kt_gpu, false, true)), M, N); + auto ccvr = nn::dsl::compute_reduce( + cpu, nn::dsl::col_reduce_max(nn::dsl::matmul(Qt_cpu, Kt_cpu, false, true)), M, N); + auto gcvr = nn::dsl::compute_reduce( + gpu, nn::dsl::col_reduce_max(nn::dsl::matmul(Qt_gpu, Kt_gpu, false, true)), M, N); if (!ccvr || !gcvr) { std::cerr << " col_max(matmul) 归约向量求值失败（CPU/GPU）: " @@ -506,16 +528,20 @@ int run_matmul_reduce(CpuEngine& cpu, GpuEngine& gpu) // 遍历全部 2M 行，每行按 row/m_per 就地分解 batch { Matrix Q2m(M * 2, K), K2m(N * 2, K); - for (auto& v : Q2m.span()) v = dist(rng); - for (auto& v : K2m.span()) v = dist(rng); + for (auto& v : Q2m.span()) + v = dist(rng); + for (auto& v : K2m.span()) + v = dist(rng); const Tensor Q2t_cpu = upload(cpu, Q2m), Q2t_gpu = upload(gpu, Q2m); const Tensor K2t_cpu = upload(cpu, K2m), K2t_gpu = upload(gpu, K2m); - auto cb2 = nn::dsl::compute_reduce(cpu, - nn::dsl::col_reduce_max( - nn::dsl::matmul(Q2t_cpu, K2t_cpu, false, true, /*batch=*/2)), M * 2, N); - auto gb2 = nn::dsl::compute_reduce(gpu, - nn::dsl::col_reduce_max( - nn::dsl::matmul(Q2t_gpu, K2t_gpu, false, true, /*batch=*/2)), M * 2, N); + auto cb2 = nn::dsl::compute_reduce( + cpu, + nn::dsl::col_reduce_max(nn::dsl::matmul(Q2t_cpu, K2t_cpu, false, true, /*batch=*/2)), + M * 2, N); + auto gb2 = nn::dsl::compute_reduce( + gpu, + nn::dsl::col_reduce_max(nn::dsl::matmul(Q2t_gpu, K2t_gpu, false, true, /*batch=*/2)), + M * 2, N); if (!cb2 || !gb2) { std::cerr << " col_max(matmul) batch=2 求值失败（CPU/GPU）: " @@ -526,8 +552,7 @@ int run_matmul_reduce(CpuEngine& cpu, GpuEngine& gpu) NN_EXIT(gb2m, 1, " GPU col_max(matmul) batch=2 下载失败\n"); const Scalar err5 = max_abs_diff(download(cpu, *cb2), *gb2m); const bool ok5 = err5 < 1e-4f; - std::cout << "[" << (ok5 ? "PASS" : "FAIL") - << "] col_max(matmul) batch=2 归约向量 (1x4)" + std::cout << "[" << (ok5 ? "PASS" : "FAIL") << "] col_max(matmul) batch=2 归约向量 (1x4)" << " err=" << std::scientific << std::setprecision(2) << err5 << "\n"; NN_EXIT(ok5, 1); } @@ -545,23 +570,22 @@ int run_fallback(CpuEngine& cpu, GpuEngine& gpu) std::uniform_real_distribution dist(-2.0f, 2.0f); const std::size_t R = 5, C = 7; Matrix x(R, C), y(R, C); - for (auto& v : x.span()) v = dist(rng); - for (auto& v : y.span()) v = dist(rng); + for (auto& v : x.span()) + v = dist(rng); + for (auto& v : y.span()) + v = dist(rng); const Tensor xt_cpu = upload(cpu, x), xt_gpu = upload(gpu, x); const Tensor yt_cpu = upload(cpu, y), yt_gpu = upload(gpu, y); const nn::Scalar three{3}; // CPU 求值正常（模板路径） - auto cr = nn::dsl::compute(cpu, - nn::dsl::leaf(xt_cpu) * nn::dsl::leaf(yt_cpu) + three, R, C); + auto cr = nn::dsl::compute(cpu, nn::dsl::leaf(xt_cpu) * nn::dsl::leaf(yt_cpu) + three, R, C); NN_EXIT(cr, 1, " CPU 求值失败: "); // 未扫描表达式在 GPU 上必须**硬报错**（闭合世界，绝不静默回退） - auto gr = nn::dsl::compute(gpu, - nn::dsl::leaf(xt_gpu) * nn::dsl::leaf(yt_gpu) + three, R, C); + auto gr = nn::dsl::compute(gpu, nn::dsl::leaf(xt_gpu) * nn::dsl::leaf(yt_gpu) + three, R, C); const bool ok = !gr; - if (!ok) - std::cerr << " 未扫描表达式在 GPU 上未报错（应硬报错而非静默回退/成功）\n"; + if (!ok) std::cerr << " 未扫描表达式在 GPU 上未报错（应硬报错而非静默回退/成功）\n"; std::cout << "[" << (ok ? "PASS" : "FAIL") << "] 未扫描表达式 → GPU 硬报错（不静默回退）\n"; return ok ? 0 : 1; } @@ -580,27 +604,49 @@ int run_reduce_consts(nn::ComputeEngine& cpu, nn::ComputeEngine& gpu) // 列 1：无并列（max=5 出现 1 次） // 列 2：全相等（max=7 出现 4 次） Matrix xm(kk, cols); - xm.span()[0] = 1; xm.span()[1] = 5; xm.span()[2] = 7; - xm.span()[3] = 1; xm.span()[4] = 4; xm.span()[5] = 7; - xm.span()[6] = 0; xm.span()[7] = 3; xm.span()[8] = 7; - xm.span()[9] = 0; xm.span()[10] = 2; xm.span()[11] = 7; + xm.span()[0] = 1; + xm.span()[1] = 5; + xm.span()[2] = 7; + xm.span()[3] = 1; + xm.span()[4] = 4; + xm.span()[5] = 7; + xm.span()[6] = 0; + xm.span()[7] = 3; + xm.span()[8] = 7; + xm.span()[9] = 0; + xm.span()[10] = 2; + xm.span()[11] = 7; Matrix mxm(1, cols); - mxm.span()[0] = 1; mxm.span()[1] = 5; mxm.span()[2] = 7; + mxm.span()[0] = 1; + mxm.span()[1] = 5; + mxm.span()[2] = 7; const Scalar expect[3] = {2, 1, 4}; - const auto eval = [&](nn::ComputeEngine& e, const char* tag, Matrix& out) -> bool { + const auto eval = [&](nn::ComputeEngine& e, const char* tag, Matrix& out) -> bool + { auto x = e.from_matrix(xm); auto mx = e.from_matrix(mxm); - if (!x || !mx) { std::cerr << " [" << tag << "] 上传失败\n"; return false; } - auto r = nn::dsl::compute_reduce(e, + if (!x || !mx) + { + std::cerr << " [" << tag << "] 上传失败\n"; + return false; + } + auto r = nn::dsl::compute_reduce( + e, nn::dsl::col_reduce_sum(nn::dsl::select( - nn::dsl::leaf(*x) == nn::dsl::col_broadcast(*mx), - nn::Scalar{1}, nn::Scalar{0})), + nn::dsl::leaf(*x) == nn::dsl::col_broadcast(*mx), nn::Scalar{1}, nn::Scalar{0})), kk, cols); - if (!r) { std::cerr << " [" << tag << "] compute_reduce 失败: " - << r.error().message << "\n"; return false; } + if (!r) + { + std::cerr << " [" << tag << "] compute_reduce 失败: " << r.error().message << "\n"; + return false; + } auto m = e.to_matrix(*r); - if (!m) { std::cerr << " [" << tag << "] to_matrix 失败\n"; return false; } + if (!m) + { + std::cerr << " [" << tag << "] to_matrix 失败\n"; + return false; + } out = *m; return true; }; @@ -614,8 +660,7 @@ int run_reduce_consts(nn::ComputeEngine& cpu, nn::ComputeEngine& gpu) for (std::size_t i = 0; ok && i < cols; ++i) ok = std::fabs(mc.span()[i] - expect[i]) < 1e-6f; if (!ok) - std::cerr << " CPU vs GPU max_diff=" << diff - << "（GPU 常量池错位的历史 bug 回归）\n"; + std::cerr << " CPU vs GPU max_diff=" << diff << "（GPU 常量池错位的历史 bug 回归）\n"; std::cout << "[" << (ok ? "PASS" : "FAIL") << "] 归约内联常量（push-constant 固定头长度）CPU/GPU 一致\n"; return ok ? 0 : 1; @@ -632,14 +677,14 @@ int run_fold_gpu(CpuEngine& cpu, GpuEngine& gpu) int fail = 0; std::mt19937 rng(9013); std::uniform_real_distribution dist(-2.0f, 2.0f); - for (const std::size_t K : {std::size_t{1}, std::size_t{7}, std::size_t{32}, - std::size_t{33}, std::size_t{100}, - std::size_t{256}, std::size_t{260}}) + for (const std::size_t K : {std::size_t{1}, std::size_t{7}, std::size_t{32}, std::size_t{33}, + std::size_t{100}, std::size_t{256}, std::size_t{260}}) { const std::size_t rows = 6; // 写路径：宿主 Matrix 填好再各上传一份（出生绑定，绑 cpu 的喂 gpu 硬错误） Matrix xm(rows, K); - for (auto& v : xm.span()) v = dist(rng); + for (auto& v : xm.span()) + v = dist(rng); const Tensor xc = upload(cpu, xm); const Tensor xg = upload(gpu, xm); const std::vector ins_cpu{xc}; @@ -651,35 +696,36 @@ int run_fold_gpu(CpuEngine& cpu, GpuEngine& gpu) auto g = gpu.eval_expr(spec, ins_gpu, rows, 1); if (!c) { - std::cout << " [FAIL] fold " << tag << " K=" << K - << " CPU: " << c.error().message << "\n"; - ++fail; return; + std::cout << " [FAIL] fold " << tag << " K=" << K << " CPU: " << c.error().message + << "\n"; + ++fail; + return; } if (!g) { - std::cout << " [FAIL] fold " << tag << " K=" << K - << " GPU: " << g.error().message << "\n"; - ++fail; return; + std::cout << " [FAIL] fold " << tag << " K=" << K << " GPU: " << g.error().message + << "\n"; + ++fail; + return; } auto gm = gpu.to_matrix(*g); if (!gm) { - std::cout << " [FAIL] fold " << tag << " K=" << K - << " GPU 下载失败\n"; - ++fail; return; + std::cout << " [FAIL] fold " << tag << " K=" << K << " GPU 下载失败\n"; + ++fail; + return; } const auto cm = download(cpu, *c); const auto cs = cm.span(); const auto gs = gm->span(); Scalar err = 0; for (std::size_t i = 0; i < rows; ++i) - err = std::fmax(err, std::fabs(cs[i] - gs[i]) / - std::fmax(Scalar{1}, std::fabs(cs[i]))); + err = std::fmax(err, + std::fabs(cs[i] - gs[i]) / std::fmax(Scalar{1}, std::fabs(cs[i]))); const bool ok = err <= tol; - std::cout << "[" << (ok ? "PASS" : "FAIL") << "] fold " << tag - << " K=" << K << " err=" << std::scientific - << std::setprecision(2) << err << std::defaultfloat - << (tol == 0 ? " (exact)" : "") << "\n"; + std::cout << "[" << (ok ? "PASS" : "FAIL") << "] fold " << tag << " K=" << K + << " err=" << std::scientific << std::setprecision(2) << err + << std::defaultfloat << (tol == 0 ? " (exact)" : "") << "\n"; if (!ok) { // 诊断：打印全部行的 CPU/GPU 实值（行同值/巨值形态定位） @@ -694,12 +740,10 @@ int run_fold_gpu(CpuEngine& cpu, GpuEngine& gpu) } }; - one("rowmax", nn::expr::make_fold_rowmax( - static_cast(K)), Scalar{0}); - one("rowsum", nn::expr::make_fold_rowsum( - static_cast(K)), Scalar{1e-6}); - one("softmax_denom", nn::expr::make_fold_softmax_denom( - static_cast(K)), Scalar{1e-5}); + one("rowmax", nn::expr::make_fold_rowmax(static_cast(K)), Scalar{0}); + one("rowsum", nn::expr::make_fold_rowsum(static_cast(K)), Scalar{1e-6}); + one("softmax_denom", nn::expr::make_fold_softmax_denom(static_cast(K)), + Scalar{1e-5}); } return fail; } @@ -716,19 +760,24 @@ int run_fold_attn_gpu(CpuEngine& cpu, GpuEngine& gpu) int fail = 0; std::mt19937 rng(9024); std::uniform_real_distribution dist(-1.5f, 1.5f); - struct Sh { std::uint32_t bh, seq, dk; }; - struct Mk { nn::expr::AttnMaskKind mask; bool score_bias; const char* name; }; + struct Sh + { std::uint32_t bh, seq, dk; }; + struct Mk + { + nn::expr::AttnMaskKind mask; + bool score_bias; + const char* name; + }; for (const Mk mk : {Mk{nn::expr::AttnMaskKind::Plain, false, "plain"}, - Mk{nn::expr::AttnMaskKind::Causal, false, "causal"}, - Mk{nn::expr::AttnMaskKind::Causal, true, "alibi"}, - Mk{nn::expr::AttnMaskKind::CausalDoc, false, "doc"}, - Mk{nn::expr::AttnMaskKind::CausalDoc, true, "alibidoc"}}) + Mk{nn::expr::AttnMaskKind::Causal, false, "causal"}, + Mk{nn::expr::AttnMaskKind::Causal, true, "alibi"}, + Mk{nn::expr::AttnMaskKind::CausalDoc, false, "doc"}, + Mk{nn::expr::AttnMaskKind::CausalDoc, true, "alibidoc"}}) { const char* mname = mk.name; // {2,133,4}：seq > EXPR_FOLD_BLOCK(128) → GPU 多块流式 + causal // 整块跳过分支（k0>qt）；{1,1,4}：GPU 最小 rows=1（clamp+row_ok） - for (const Sh sh : {Sh{2, 9, 4}, Sh{1, 33, 8}, Sh{3, 5, 2}, - Sh{2, 133, 4}, Sh{1, 1, 4}}) + for (const Sh sh : {Sh{2, 9, 4}, Sh{1, 33, 8}, Sh{3, 5, 2}, Sh{2, 133, 4}, Sh{1, 1, 4}}) { const std::size_t rows_out = static_cast(sh.bh) * sh.seq; const std::size_t q_rows = static_cast(sh.bh) * sh.dk; @@ -747,21 +796,22 @@ int run_fold_attn_gpu(CpuEngine& cpu, GpuEngine& gpu) // max init=-inf → dm=−inf−−inf=NaN 的分支；边界若落在 128 之内 // （如 seq/2），首块永留有效项，该分支测不到。 const std::uint32_t doc_boundary = - sseq > nn::EXPR_FOLD_BLOCK ? nn::EXPR_FOLD_BLOCK + 1u - : sseq / 2u; + sseq > nn::EXPR_FOLD_BLOCK ? nn::EXPR_FOLD_BLOCK + 1u : sseq / 2u; const auto doc_of = [doc_boundary](std::uint32_t pos) -> Scalar { return pos < doc_boundary ? Scalar{1} : Scalar{2}; }; for (std::size_t r = 0; r < rows_out; ++r) - doc_colm.span()[r] = - doc_of(static_cast(r % sh.seq)); + doc_colm.span()[r] = doc_of(static_cast(r % sh.seq)); for (std::uint32_t blk = 0; blk < sh.bh; ++blk) for (std::uint32_t j = 0; j < sh.seq; ++j) - doc_ids_tm.span() - [static_cast(blk) * sh.seq + j] = doc_of(j); - for (auto& v : Qm.span()) v = dist(rng); - for (auto& v : Km.span()) v = dist(rng); - for (auto& v : Vtm.span()) v = dist(rng); - for (auto& v : slopesm.span()) v = dist(rng) * Scalar{0.1}; + doc_ids_tm.span()[static_cast(blk) * sh.seq + j] = doc_of(j); + for (auto& v : Qm.span()) + v = dist(rng); + for (auto& v : Km.span()) + v = dist(rng); + for (auto& v : Vtm.span()) + v = dist(rng); + for (auto& v : slopesm.span()) + v = dist(rng) * Scalar{0.1}; const bool alibi = mk.score_bias; const bool docm = (mk.mask == nn::expr::AttnMaskKind::CausalDoc); // 出生绑定：绑 cpu 的张量喂 gpu 引擎是硬错误，输入各上传一份 @@ -783,33 +833,37 @@ int run_fold_attn_gpu(CpuEngine& cpu, GpuEngine& gpu) } if (docm) { - ins_cpu.push_back(doc_col_cpu); ins_cpu.push_back(doc_ids_cpu); - ins_gpu.push_back(doc_col_gpu); ins_gpu.push_back(doc_ids_gpu); + ins_cpu.push_back(doc_col_cpu); + ins_cpu.push_back(doc_ids_cpu); + ins_gpu.push_back(doc_col_gpu); + ins_gpu.push_back(doc_ids_gpu); } - nn::ExprSpec spec = nn::expr::make_fold_attn_o(sh.seq, sh.dk, sh.bh, mk.mask, mk.score_bias); + nn::ExprSpec spec = + nn::expr::make_fold_attn_o(sh.seq, sh.dk, sh.bh, mk.mask, mk.score_bias); auto c = cpu.eval_expr(spec, ins_cpu, rows_out, sh.dk); auto g = gpu.eval_expr(spec, ins_gpu, rows_out, sh.dk); - const std::string tag = std::string("attn-fold ") + mname + - " bh=" + std::to_string(sh.bh) + " seq=" + std::to_string(sh.seq) + - " dk=" + std::to_string(sh.dk); + const std::string tag = + std::string("attn-fold ") + mname + " bh=" + std::to_string(sh.bh) + + " seq=" + std::to_string(sh.seq) + " dk=" + std::to_string(sh.dk); if (!c) { - std::cout << " [FAIL] " << tag << " CPU: " - << c.error().message << "\n"; - ++fail; continue; + std::cout << " [FAIL] " << tag << " CPU: " << c.error().message << "\n"; + ++fail; + continue; } if (!g) { - std::cout << " [FAIL] " << tag << " GPU: " - << g.error().message << "\n"; - ++fail; continue; + std::cout << " [FAIL] " << tag << " GPU: " << g.error().message << "\n"; + ++fail; + continue; } auto gm = gpu.to_matrix(*g); if (!gm) { std::cout << " [FAIL] " << tag << " GPU 下载失败\n"; - ++fail; continue; + ++fail; + continue; } const auto cm = download(cpu, *c); const auto cs = cm.span(); @@ -820,22 +874,23 @@ int run_fold_attn_gpu(CpuEngine& cpu, GpuEngine& gpu) // NaN 守卫：GPU 输出 NaN → diff=NaN → IEEE fmax(err,NaN)=err // 会静默吞掉 → err 保持正常值照样 PASS，故此处记 inf // 必超容差（NaN 输出必须硬失败） - const Scalar diff = std::fabs(cs[i] - gs[i]) / - std::fmax(Scalar{1}, std::fabs(cs[i])); + const Scalar diff = + std::fabs(cs[i] - gs[i]) / std::fmax(Scalar{1}, std::fabs(cs[i])); err = std::isfinite(diff) ? std::fmax(err, diff) : std::numeric_limits::infinity(); if (!std::isfinite(err)) break; } const bool ok = err <= Scalar{1e-4}; - std::cout << "[" << (ok ? "PASS" : "FAIL") << "] " << tag - << " err=" << std::scientific << std::setprecision(2) - << err << std::defaultfloat << "\n"; + std::cout << "[" << (ok ? "PASS" : "FAIL") << "] " << tag << " err=" << std::scientific + << std::setprecision(2) << err << std::defaultfloat << "\n"; if (!ok) { std::cout << " cpu[0..3]:"; - for (std::size_t i = 0; i < 4; ++i) std::cout << " " << cs[i]; + for (std::size_t i = 0; i < 4; ++i) + std::cout << " " << cs[i]; std::cout << "\n gpu[0..3]:"; - for (std::size_t i = 0; i < 4; ++i) std::cout << " " << gs[i]; + for (std::size_t i = 0; i < 4; ++i) + std::cout << " " << gs[i]; std::cout << "\n"; ++fail; } @@ -866,9 +921,8 @@ int main() int fail = 0; // 形状无关融合：任意 d_k（含非 2 的幂 40/96）都命中同一个融合 shader - for (const std::size_t dk : {std::size_t{16}, std::size_t{40}, - std::size_t{64}, std::size_t{96}, - std::size_t{128}}) + for (const std::size_t dk : + {std::size_t{16}, std::size_t{40}, std::size_t{64}, std::size_t{96}, std::size_t{128}}) { fail += run_rope(*cpu_engine, *gpu_engine, dk, /*backward=*/false); fail += run_rope(*cpu_engine, *gpu_engine, dk, /*backward=*/true); diff --git a/src/gpt_checkpoint_test.cpp b/src/gpt_checkpoint_test.cpp index 26cc264f..b739169c 100644 --- a/src/gpt_checkpoint_test.cpp +++ b/src/gpt_checkpoint_test.cpp @@ -36,12 +36,12 @@ namespace int run_test() { - const std::size_t vocab = 257; - const std::size_t d_model = 64; - const std::size_t seq_len = 16; - const std::size_t batch = 2; + const std::size_t vocab = 257; + const std::size_t d_model = 64; + const std::size_t seq_len = 16; + const std::size_t batch = 2; const std::size_t num_heads = 4; - const std::size_t d_ff = 128; + const std::size_t d_ff = 128; const std::size_t num_layers = 6; auto eng_owner = nn::CpuEngine::create(); @@ -50,16 +50,13 @@ int run_test() std::cout << "========================================\n"; std::cout << " GPT 梯度检查点一致性测试 (L1)\n"; std::cout << "========================================\n"; - std::cout << " vocab=" << vocab << " d_model=" << d_model - << " seq=" << seq_len << " batch=" << batch - << " heads=" << num_heads << " d_ff=" << d_ff + std::cout << " vocab=" << vocab << " d_model=" << d_model << " seq=" << seq_len + << " batch=" << batch << " heads=" << num_heads << " d_ff=" << d_ff << " layers=" << num_layers << "\n"; // 单一模型实例（checkpoint 关闭初始化） - GPTModel model(vocab, d_model, seq_len, num_heads, d_ff, num_layers, - PosEncodingType::Learned, - ActivationType::SwiGLU, - NormType::RMSNorm); + GPTModel model(vocab, d_model, seq_len, num_heads, d_ff, num_layers, PosEncodingType::Learned, + ActivationType::SwiGLU, NormType::RMSNorm); NN_EXIT(model.init(eng), 1, "GPTModel init 失败: "); std::mt19937_64 rng(7); @@ -72,7 +69,8 @@ int run_test() // 固定输出梯度 (vocab, seq*batch) nn::Matrix go_m(vocab, seq_len * batch); - for (std::size_t i = 0; i < go_m.size(); ++i) go_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); NN_EXIT(x, 1, "from_matrix(x) failed\n"); @@ -152,6 +150,4 @@ int run_test() } // namespace int main() -{ - return run_test(); -} +{ return run_test(); } diff --git a/src/gpt_gradcheck.cpp b/src/gpt_gradcheck.cpp index 11c46ec2..94e0f030 100644 --- a/src/gpt_gradcheck.cpp +++ b/src/gpt_gradcheck.cpp @@ -25,26 +25,30 @@ #include #include "test_common.hpp" -using nn::Scalar; -using nn::Matrix; -using nn::Tensor; -using nn::TensorRef; +using nn::ActivationType; using nn::ComputeEngine; using nn::GPTModel; using nn::Layer; -using nn::PosEncodingType; -using nn::ActivationType; +using nn::Matrix; using nn::NormType; +using nn::PosEncodingType; +using nn::Scalar; +using nn::Tensor; +using nn::TensorRef; -namespace { +namespace +{ // 用当前参数前向，返回 L = Σ logits*go -Scalar eval_loss(ComputeEngine& engine, GPTModel& model, - const Tensor& input, const Tensor& go) +Scalar eval_loss(ComputeEngine& engine, GPTModel& model, const Tensor& input, const Tensor& go) { auto y = model.forward(input); - if (!y) { std::cerr << " [eval] forward failed: " << y.error().message << "\n"; std::abort(); } + if (!y) + { + std::cerr << " [eval] forward failed: " << y.error().message << "\n"; + std::abort(); + } auto y_m = engine.to_matrix(*y); auto go_m = engine.to_matrix(go); return dot(*y_m, *go_m); @@ -52,14 +56,10 @@ Scalar eval_loss(ComputeEngine& engine, GPTModel& model, // 验证单个参数张量梯度（逐元素中心差分） -bool check_grad_tensor( - ComputeEngine& engine, GPTModel& model, - const Tensor& input, const Tensor& go, - Tensor& param, - const Matrix& base, - const Matrix& grad_analytical, - const std::string& name, - Scalar eps, Scalar tol) +bool check_grad_tensor(ComputeEngine& engine, GPTModel& model, const Tensor& input, + const Tensor& go, Tensor& param, const Matrix& base, + const Matrix& grad_analytical, const std::string& name, Scalar eps, + Scalar tol) { bool ok = true; std::size_t bad = 0; @@ -73,18 +73,30 @@ bool check_grad_tensor( Matrix pp = base; pp.set_value_unchecked(r, c, orig + eps); auto r_p = engine.write(param, pp.span()); - if (!r_p) { std::cerr << "write(+) failed\n"; return false; } + if (!r_p) + { + std::cerr << "write(+) failed\n"; + return false; + } const Scalar lp = eval_loss(engine, model, input, go); Matrix pm = base; pm.set_value_unchecked(r, c, orig - eps); auto r_m = engine.write(param, pm.span()); - if (!r_m) { std::cerr << "write(-) failed\n"; return false; } + if (!r_m) + { + std::cerr << "write(-) failed\n"; + return false; + } const Scalar lm = eval_loss(engine, model, input, go); // 恢复原值 auto r_r = engine.write(param, base.span()); - if (!r_r) { std::cerr << "write(restore) failed\n"; return false; } + if (!r_r) + { + std::cerr << "write(restore) failed\n"; + return false; + } const Scalar num = (lp - lm) / (Scalar{2} * eps); const Scalar ana = grad_analytical.at_unchecked(r, c); @@ -94,16 +106,14 @@ bool check_grad_tensor( { ok = false; if (bad < 6) - std::cout << " [FAIL] " << name << "(" << r << "," << c - << ") num=" << num << " ana=" << ana - << " err=" << err << "\n"; + std::cout << " [FAIL] " << name << "(" << r << "," << c << ") num=" << num + << " ana=" << ana << " err=" << err << "\n"; ++bad; } } } - std::cout << " " << name << ": " - << (ok ? "OK" : "FAIL") - << " max_err=" << max_err << " bad=" << bad << "\n"; + std::cout << " " << name << ": " << (ok ? "OK" : "FAIL") << " max_err=" << max_err + << " bad=" << bad << "\n"; return ok; } @@ -119,19 +129,17 @@ int main(int argc, char* argv[]) for (int i = 1; i < argc; ++i) { std::string a = argv[i]; - if (a == "--tol" && i + 1 < argc) - tol = static_cast(std::atof(argv[++i])); + if (a == "--tol" && i + 1 < argc) tol = static_cast(std::atof(argv[++i])); else if (a == "--layers" && i + 1 < argc) num_layers = static_cast(std::atoi(argv[++i])); else if (a == "--batch" && i + 1 < argc) batch = static_cast(std::atoi(argv[++i])); - else if (a == "--pos-enc" && i + 1 < argc) - pos_enc_name = argv[++i]; - else if (a == "--gpu") - use_gpu = true; + else if (a == "--pos-enc" && i + 1 < argc) pos_enc_name = argv[++i]; + else if (a == "--gpu") use_gpu = true; else if (a == "--help") { - std::cout << "用法: gpt_gradcheck [--gpu] [--layers N] [--batch N] [--tol ] [--pos-enc learned|sinusoidal|alibi|rope]\n"; + std::cout << "用法: gpt_gradcheck [--gpu] [--layers N] [--batch N] [--tol ] " + "[--pos-enc learned|sinusoidal|alibi|rope]\n"; return 0; } } @@ -159,17 +167,14 @@ int main(int argc, char* argv[]) std::cout << "========================================\n"; std::cout << " GPTModel 整链数值梯度检查 (gradcheck)\n"; std::cout << "========================================\n"; - std::cout << " vocab=" << vocab << " d_model=" << d_model - << " seq=" << seq_len << " batch=" << batch - << " heads=" << num_heads << " d_ff=" << d_ff + std::cout << " vocab=" << vocab << " d_model=" << d_model << " seq=" << seq_len + << " batch=" << batch << " heads=" << num_heads << " d_ff=" << d_ff << " layers=" << num_layers << "\n"; std::cout << " activation=SwiGLU norm=RMSNorm pos=" << pos_enc_name << "\n"; std::cout << " tol=" << tol << "\n"; - GPTModel model(vocab, d_model, seq_len, num_heads, d_ff, num_layers, - pos_enc, - ActivationType::SwiGLU, - NormType::RMSNorm); + GPTModel model(vocab, d_model, seq_len, num_heads, d_ff, num_layers, pos_enc, + ActivationType::SwiGLU, NormType::RMSNorm); NN_EXIT(model.init(eng), 1, "GPTModel init 失败: "); std::mt19937_64 rng(123); @@ -182,7 +187,8 @@ int main(int argc, char* argv[]) // 输出权重 go (vocab, seq*batch) Matrix go_m(vocab, seq_len * batch); - for (std::size_t i = 0; i < go_m.size(); ++i) go_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); NN_EXIT(x, 1, "from_matrix(x) failed\n"); @@ -196,11 +202,11 @@ int main(int argc, char* argv[]) // ── 清零梯度并 backward ── auto params = model.parameters(); - auto grads = model.param_gradients(); + auto grads = model.param_gradients(); if (params.size() != grads.size()) { - std::cerr << "params/grads size mismatch: " - << params.size() << " vs " << grads.size() << "\n"; + std::cerr << "params/grads size mismatch: " << params.size() << " vs " << grads.size() + << "\n"; return 1; } for (auto& g : grads) @@ -224,22 +230,19 @@ int main(int argc, char* argv[]) NN_EXIT(gm, 1, "to_matrix(grad) failed\n"); if (pm->rows() != gm->rows() || pm->cols() != gm->cols()) { - std::cerr << "shape mismatch param[" << p << "]: " - << pm->rows() << "x" << pm->cols() + std::cerr << "shape mismatch param[" << p << "]: " << pm->rows() << "x" << pm->cols() << " vs grad " << gm->rows() << "x" << gm->cols() << "\n"; return 1; } std::string name = "param[" + std::to_string(pi++) + "]"; // 精简命名：按顺序标注 static constexpr std::array tags = { - "token_emb", "pos_emb", "blk0.wq", "blk0.wk", - "blk0.wv", "blk0.wo", "blk0.norm1.g", "blk0.fc1.w", "blk0.fc1.b", - "blk0.fc2.w", "blk0.fc2.b", "blk0.norm2.g"}; + "token_emb", "pos_emb", "blk0.wq", "blk0.wk", "blk0.wv", "blk0.wo", + "blk0.norm1.g", "blk0.fc1.w", "blk0.fc1.b", "blk0.fc2.w", "blk0.fc2.b", "blk0.norm2.g"}; if (pi - 1 < 12) - name = std::string("p[") + std::to_string(pi - 1) + "]= " - + std::string(tags[pi - 1]); - all_pass &= check_grad_tensor( - eng, model, *x, *go, params[p].get(), *pm, *gm, name, eps, tol); + name = std::string("p[") + std::to_string(pi - 1) + "]= " + std::string(tags[pi - 1]); + all_pass &= + check_grad_tensor(eng, model, *x, *go, params[p].get(), *pm, *gm, name, eps, tol); } std::cout << "----------------------------------------\n"; diff --git a/src/gpt_offload_test.cpp b/src/gpt_offload_test.cpp index 631c1eb8..27d88974 100644 --- a/src/gpt_offload_test.cpp +++ b/src/gpt_offload_test.cpp @@ -48,12 +48,12 @@ namespace int run_test() { - const std::size_t vocab = 257; - const std::size_t d_model = 64; - const std::size_t seq_len = 16; - const std::size_t batch = 2; + const std::size_t vocab = 257; + const std::size_t d_model = 64; + const std::size_t seq_len = 16; + const std::size_t batch = 2; const std::size_t num_heads = 4; - const std::size_t d_ff = 128; + const std::size_t d_ff = 128; const std::size_t num_layers = 6; auto backend_owner = GpuBackend::shared(); @@ -72,15 +72,12 @@ int run_test() std::cout << "========================================\n" << " GPT activation offload 一致性测试\n" << "========================================\n"; - std::cout << " vocab=" << vocab << " d_model=" << d_model - << " seq=" << seq_len << " batch=" << batch - << " heads=" << num_heads << " d_ff=" << d_ff + std::cout << " vocab=" << vocab << " d_model=" << d_model << " seq=" << seq_len + << " batch=" << batch << " heads=" << num_heads << " d_ff=" << d_ff << " layers=" << num_layers << "\n"; - GPTModel model(vocab, d_model, seq_len, num_heads, d_ff, num_layers, - PosEncodingType::Learned, - ActivationType::SwiGLU, - NormType::RMSNorm); + GPTModel model(vocab, d_model, seq_len, num_heads, d_ff, num_layers, PosEncodingType::Learned, + ActivationType::SwiGLU, NormType::RMSNorm); NN_EXIT(model.init(eng), 1, "GPTModel init 失败: "); std::mt19937_64 rng(7); @@ -90,7 +87,8 @@ int run_test() for (std::size_t i = 0; i < x_m.size(); ++i) x_m.span()[i] = static_cast(rng() % vocab); Matrix go_m(vocab, seq_len * batch); - for (std::size_t i = 0; i < go_m.size(); ++i) go_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); NN_EXIT(x, 1, "from_matrix(x) failed\n"); @@ -172,7 +170,5 @@ int run_test() } // namespace int main() -{ - return run_test(); -} +{ return run_test(); } #endif diff --git a/src/gpu_batch_scope_test.cpp b/src/gpu_batch_scope_test.cpp index 64d96414..1f9f667a 100644 --- a/src/gpu_batch_scope_test.cpp +++ b/src/gpu_batch_scope_test.cpp @@ -18,16 +18,18 @@ #include int g_fail = 0; -#define CHECK(cond, msg) \ - do { \ - if (!(cond)) { \ - std::cout << "[FAIL] " << __FILE__ << ":" << __LINE__ << ": " \ - << msg << "\n"; \ - ++g_fail; \ - } \ +#define CHECK(cond, msg) \ + do \ + { \ + if (!(cond)) \ + { \ + std::cout << "[FAIL] " << __FILE__ << ":" << __LINE__ << ": " << msg << "\n"; \ + ++g_fail; \ + } \ } while (0) -namespace { +namespace +{ // 多 Conv 小模型：issue #28 的触发形态（第二次 Conv 分配处 OOM） nn::Result make_conv_model(nn::ComputeEngine& eng) diff --git a/src/gpu_f16_test.cpp b/src/gpu_f16_test.cpp index a9c30a76..5308bf26 100644 --- a/src/gpu_f16_test.cpp +++ b/src/gpu_f16_test.cpp @@ -5,7 +5,8 @@ #include "neuralnet.cpp/compute_gpu_engine.hpp" #include "neuralnet.cpp/compute_cpu_engine.hpp" -int main() { +int main() +{ printf("=== GPU f16 Capability Check ===\n"); auto backend_owner = nn::GpuBackend::shared(); nn::GpuBackend& backend = *backend_owner; @@ -28,8 +29,10 @@ int main() { // Create f32 matrix nn::Matrix m(2, 2); - m.set_value(0, 0, 1.0f); m.set_value(0, 1, 2.0f); - m.set_value(1, 0, 3.0f); m.set_value(1, 1, 4.0f); + m.set_value(0, 0, 1.0f); + m.set_value(0, 1, 2.0f); + m.set_value(1, 0, 3.0f); + m.set_value(1, 1, 4.0f); // Upload as f32 auto t32 = engine.from_matrix(m); @@ -48,14 +51,17 @@ int main() { auto result = engine.to_matrix(*t32_back); NN_EXIT(result, 1, "to_matrix failed: %s\n"); - printf("Original: [%.1f %.1f; %.1f %.1f]\n", m.at(0,0), m.at(0,1), m.at(1,0), m.at(1,1)); - printf("GPU f16: [%.1f %.1f; %.1f %.1f]\n", result->at(0,0), result->at(0,1), result->at(1,0), result->at(1,1)); + printf("Original: [%.1f %.1f; %.1f %.1f]\n", m.at(0, 0), m.at(0, 1), m.at(1, 0), m.at(1, 1)); + printf("GPU f16: [%.1f %.1f; %.1f %.1f]\n", result->at(0, 0), result->at(0, 1), + result->at(1, 0), result->at(1, 1)); // Test f16 matmul printf("\n=== GPU f16 Matmul ===\n"); nn::Matrix m2(2, 2); - m2.set_value(0, 0, 5.0f); m2.set_value(0, 1, 6.0f); - m2.set_value(1, 0, 7.0f); m2.set_value(1, 1, 8.0f); + m2.set_value(0, 0, 5.0f); + m2.set_value(0, 1, 6.0f); + m2.set_value(1, 0, 7.0f); + m2.set_value(1, 1, 8.0f); auto t16_a = engine.cast(engine.from_matrix(m).value(), nn::Precision::F16); auto t16_b = engine.cast(engine.from_matrix(m2).value(), nn::Precision::F16); @@ -65,8 +71,8 @@ int main() { auto c32 = engine.cast(*t16_c, nn::Precision::F32); auto c_mat = engine.to_matrix(*c32); - printf("A*B (f16): [%.1f %.1f; %.1f %.1f]\n", - c_mat->at(0,0), c_mat->at(0,1), c_mat->at(1,0), c_mat->at(1,1)); + printf("A*B (f16): [%.1f %.1f; %.1f %.1f]\n", c_mat->at(0, 0), c_mat->at(0, 1), c_mat->at(1, 0), + c_mat->at(1, 1)); printf("Expected: [19.0 22.0; 43.0 50.0]\n"); // ── 算子 f16 变体对拍（GPU f16 vs CPU f32 参考）───────────── @@ -77,13 +83,13 @@ int main() { auto cpu_owner = nn::CpuEngine::create(); nn::CpuEngine& cpu = *cpu_owner; int fails = 0; - const auto check = [&](const char* name, const nn::Matrix& got, - const nn::Matrix& ref, float tol) + const auto check = + [&](const char* name, const nn::Matrix& got, const nn::Matrix& ref, float tol) { if (got.rows() != ref.rows() || got.cols() != ref.cols()) { - printf("[FAIL] %s: shape %zux%zu vs ref %zux%zu\n", name, - got.rows(), got.cols(), ref.rows(), ref.cols()); + printf("[FAIL] %s: shape %zux%zu vs ref %zux%zu\n", name, got.rows(), got.cols(), + ref.rows(), ref.cols()); ++fails; return; } @@ -100,8 +106,7 @@ int main() { printf("[FAIL] %s: max_rel=%.3e (tol %.1e)\n", name, maxr, tol); ++fails; } - else - printf("[ok] %s (max_rel=%.2e)\n", name, maxr); + else printf("[ok] %s (max_rel=%.2e)\n", name, maxr); }; // GPU f16 算子 → 结果 cast f32 → to_matrix const auto gpu_op = [&](auto&& fn) -> nn::Matrix @@ -114,9 +119,19 @@ int main() { return {}; } auto f32r = engine.cast(*r, nn::Precision::F32); - if (!f32r) { printf("[FAIL] cast: %s\n", f32r.error().message.c_str()); ++fails; return {}; } + if (!f32r) + { + printf("[FAIL] cast: %s\n", f32r.error().message.c_str()); + ++fails; + return {}; + } auto m = engine.to_matrix(*f32r); - if (!m) { printf("[FAIL] to_matrix: %s\n", m.error().message.c_str()); ++fails; return {}; } + if (!m) + { + printf("[FAIL] to_matrix: %s\n", m.error().message.c_str()); + ++fails; + return {}; + } return std::move(*m); }; const auto to_gpu_f16 = [&](const nn::Matrix& m) -> nn::Tensor @@ -128,9 +143,19 @@ int main() { const auto cpu_op = [&](auto&& fn) -> nn::Matrix { auto r = fn(); - if (!r) { printf("[FAIL] cpu ref: %s\n", r.error().message.c_str()); ++fails; return {}; } + if (!r) + { + printf("[FAIL] cpu ref: %s\n", r.error().message.c_str()); + ++fails; + return {}; + } auto m = cpu.to_matrix(*r); - if (!m) { printf("[FAIL] cpu to_matrix: %s\n", m.error().message.c_str()); ++fails; return {}; } + if (!m) + { + printf("[FAIL] cpu to_matrix: %s\n", m.error().message.c_str()); + ++fails; + return {}; + } return std::move(*m); }; const auto fill = [](nn::Matrix& m, auto&& gen) @@ -157,8 +182,10 @@ int main() { { nn::Matrix x(6, 2); fill(x, [](std::size_t i) { return 0.25f * static_cast((i * 5) % 11) - 1.0f; }); - auto ref_s = cpu_op([&] { return cpu.grouped_reduce_sum(cpu.from_matrix(x).value(), 2, 3); }); - auto ref_m = cpu_op([&] { return cpu.grouped_reduce_max(cpu.from_matrix(x).value(), 2, 3); }); + auto ref_s = + cpu_op([&] { return cpu.grouped_reduce_sum(cpu.from_matrix(x).value(), 2, 3); }); + auto ref_m = + cpu_op([&] { return cpu.grouped_reduce_max(cpu.from_matrix(x).value(), 2, 3); }); auto x16 = to_gpu_f16(x); auto got_s = gpu_op([&] { return engine.grouped_reduce_sum(x16, 2, 3); }); auto got_m = gpu_op([&] { return engine.grouped_reduce_max(x16, 2, 3); }); @@ -170,33 +197,45 @@ int main() { { nn::Matrix dst_m(4, 2), grad_m(2, 2), idx_m(2, 1); fill(dst_m, [](std::size_t) { return 0.0f; }); - grad_m.set_value(0, 0, 1.5f); grad_m.set_value(0, 1, -2.25f); - grad_m.set_value(1, 0, 0.75f); grad_m.set_value(1, 1, 3.0f); - idx_m.set_value(0, 0, 1.0f); idx_m.set_value(1, 0, 3.0f); + grad_m.set_value(0, 0, 1.5f); + grad_m.set_value(0, 1, -2.25f); + grad_m.set_value(1, 0, 0.75f); + grad_m.set_value(1, 1, 3.0f); + idx_m.set_value(0, 0, 1.0f); + idx_m.set_value(1, 0, 3.0f); // CPU f32 参考 auto d_cpu = cpu.from_matrix(dst_m).value(); auto sc = cpu.scatter_add_rows(d_cpu, cpu.from_matrix(idx_m).value(), cpu.from_matrix(grad_m).value()); - if (!sc) { printf("[FAIL] cpu scatter: %s\n", sc.error().message.c_str()); ++fails; } + if (!sc) + { + printf("[FAIL] cpu scatter: %s\n", sc.error().message.c_str()); + ++fails; + } auto ref = cpu.to_matrix(d_cpu).value(); // GPU f16（原生打包 half CAS 路径） auto d16 = to_gpu_f16(dst_m); auto g16 = to_gpu_f16(grad_m); auto idx = engine.from_matrix(idx_m).value(); auto r = engine.scatter_add_rows(d16, idx, g16); - if (!r) { printf("[FAIL] scatter_add_rows f16: %s\n", r.error().message.c_str()); ++fails; } - else check("scatter_add_rows f16", gpu_op([&] { return engine.cast(d16, nn::Precision::F32); }), ref, 1e-3f); + if (!r) + { + printf("[FAIL] scatter_add_rows f16: %s\n", r.error().message.c_str()); + ++fails; + } + else + check("scatter_add_rows f16", + gpu_op([&] { return engine.cast(d16, nn::Precision::F32); }), ref, 1e-3f); } // 4) im2col / col2im（窗口展开 f16 变体；C=2,H=W=3,k=2,stride=1,pad=0） { nn::Matrix xm(18, 2); // C*H*W=18, B=2 fill(xm, [](std::size_t i) { return 0.3f * static_cast((i * 11) % 17) - 2.0f; }); - auto ref_col_t = cpu.im2col(cpu.from_matrix(xm).value(), - 2, 3, 3, 2, 1, 0, 2, 2).value(); + auto ref_col_t = cpu.im2col(cpu.from_matrix(xm).value(), 2, 3, 3, 2, 1, 0, 2, 2).value(); auto ref_col = cpu.to_matrix(ref_col_t).value(); - auto ref_back = cpu_op([&] { - return cpu.col2im(cpu.from_matrix(ref_col).value(), 2, 3, 3, 2, 1, 0, 2, 2); }); + auto ref_back = cpu_op( + [&] { return cpu.col2im(cpu.from_matrix(ref_col).value(), 2, 3, 3, 2, 1, 0, 2, 2); }); auto x16 = to_gpu_f16(xm); auto got_col = gpu_op([&] { return engine.im2col(x16, 2, 3, 3, 2, 1, 0, 2, 2); }); check("im2col f16", got_col, ref_col, 1e-3f); @@ -213,14 +252,17 @@ int main() { const std::size_t dk = 4, seq = 5, rows = 4; auto gen = [](std::size_t seed) { - return [seed](std::size_t i) { - return 0.2f * static_cast((i * 7 + seed * 13) % 11) - 1.0f; - }; + return [seed](std::size_t i) + { return 0.2f * static_cast((i * 7 + seed * 13) % 11) - 1.0f; }; }; nn::Matrix K_m(rows, seq), V_m(rows, seq), P_m(rows, seq), R_m(rows, seq); nn::Matrix A0_m(dk, dk), B0_m(dk, dk), bnd_m(1, seq); - fill(K_m, gen(1)); fill(V_m, gen(2)); fill(P_m, gen(3)); fill(R_m, gen(4)); - fill(A0_m, gen(5)); fill(B0_m, gen(6)); + fill(K_m, gen(1)); + fill(V_m, gen(2)); + fill(P_m, gen(3)); + fill(R_m, gen(4)); + fill(A0_m, gen(5)); + fill(B0_m, gen(6)); fill(bnd_m, [](std::size_t i) { return (i == 2) ? 1.0f : 0.0f; }); // 文档边界 @t=2 const auto f16 = [&](const nn::Matrix& m) { return to_gpu_f16(m); }; auto K16 = f16(K_m), V16 = f16(V_m), P16 = f16(P_m), R16 = f16(R_m); @@ -230,44 +272,69 @@ int main() { auto run_cpu = [&](auto&& fn) -> nn::Matrix { auto t = fn(); - if (!t) { printf("[FAIL] cpu scan: %s\n", t.error().message.c_str()); ++fails; return {}; } + if (!t) + { + printf("[FAIL] cpu scan: %s\n", t.error().message.c_str()); + ++fails; + return {}; + } auto m = cpu.to_matrix(*t); - if (!m) { printf("[FAIL] cpu scan to_matrix: %s\n", m.error().message.c_str()); ++fails; return {}; } + if (!m) + { + printf("[FAIL] cpu scan to_matrix: %s\n", m.error().message.c_str()); + ++fails; + return {}; + } return std::move(*m); }; - auto ref_pre = run_cpu([&] { - return cpu.scan_prefix_outer( - cpu.from_matrix(K_m).value(), cpu.from_matrix(V_m).value(), - cpu.from_matrix(P_m).value(), cpu.from_matrix(R_m).value(), - cpu.from_matrix(A0_m).value(), cpu.from_matrix(B0_m).value(), - true, dk, 1u, true, cpu.from_matrix(bnd_m).value(), true); }); - auto got_pre = gpu_op([&] { - return engine.scan_prefix_outer(K16, V16, P16, R16, A016, B016, - true, dk, 1u, true, bnd16, true, - nn::Precision::F16); }); + auto ref_pre = run_cpu( + [&] + { + return cpu.scan_prefix_outer( + cpu.from_matrix(K_m).value(), cpu.from_matrix(V_m).value(), + cpu.from_matrix(P_m).value(), cpu.from_matrix(R_m).value(), + cpu.from_matrix(A0_m).value(), cpu.from_matrix(B0_m).value(), true, dk, 1u, + true, cpu.from_matrix(bnd_m).value(), true); + }); + auto got_pre = gpu_op( + [&] + { + return engine.scan_prefix_outer(K16, V16, P16, R16, A016, B016, true, dk, 1u, true, + bnd16, true, nn::Precision::F16); + }); // 扫描族容差 5e-3：输入双侧 f16 舍入（eps≈4.9e-4）+ 跨 seq 步累加 // 的舍入传播——2e-3 实测卡在 scan_prefix 2.03e-3 上，非缺陷 check("scan_prefix_outer f16", got_pre, ref_pre, 5e-3f); nn::Matrix D_m(rows * dk, seq), X_m(rows, seq), Y_m(rows, seq); - fill(D_m, gen(7)); fill(X_m, gen(8)); fill(Y_m, gen(9)); + fill(D_m, gen(7)); + fill(X_m, gen(8)); + fill(Y_m, gen(9)); auto D16 = f16(D_m), X16 = f16(X_m), Y16 = f16(Y_m); - auto ref_suf = run_cpu([&] { - return cpu.scan_suffix_outer( - cpu.from_matrix(D_m).value(), cpu.from_matrix(X_m).value(), - cpu.from_matrix(Y_m).value(), dk, 1u, true, - cpu.from_matrix(bnd_m).value(), true); }); - auto got_suf = gpu_op([&] { - return engine.scan_suffix_outer(D16, X16, Y16, dk, 1u, true, - bnd16, true, nn::Precision::F16); }); + auto ref_suf = run_cpu( + [&] + { + return cpu.scan_suffix_outer(cpu.from_matrix(D_m).value(), + cpu.from_matrix(X_m).value(), + cpu.from_matrix(Y_m).value(), dk, 1u, true, + cpu.from_matrix(bnd_m).value(), true); + }); + auto got_suf = gpu_op( + [&] + { + return engine.scan_suffix_outer(D16, X16, Y16, dk, 1u, true, bnd16, true, + nn::Precision::F16); + }); check("scan_suffix_outer f16", got_suf, ref_suf, 5e-3f); - auto ref_oc = run_cpu([&] { - return cpu.outer_col(cpu.from_matrix(P_m).value(), - cpu.from_matrix(R_m).value(), - cpu.from_matrix(V_m).value(), dk, true); }); - auto got_oc = gpu_op([&] { - return engine.outer_col(P16, R16, V16, dk, true, nn::Precision::F16); }); + auto ref_oc = run_cpu( + [&] + { + return cpu.outer_col(cpu.from_matrix(P_m).value(), cpu.from_matrix(R_m).value(), + cpu.from_matrix(V_m).value(), dk, true); + }); + auto got_oc = + gpu_op([&] { return engine.outer_col(P16, R16, V16, dk, true, nn::Precision::F16); }); check("outer_col f16", got_oc, ref_oc, 5e-3f); } @@ -281,7 +348,8 @@ int main() { } #else #include -int main() { +int main() +{ // 纯 CPU 构建：无法执行 GPU 用例 → 77 = ctest SKIP（不得计为 Passed） printf("[SKIP] Vulkan not available\n"); return 77; diff --git a/src/gpu_stability_probe.cpp b/src/gpu_stability_probe.cpp index af76201b..abadaad0 100644 --- a/src/gpu_stability_probe.cpp +++ b/src/gpu_stability_probe.cpp @@ -44,365 +44,342 @@ namespace { - struct RunResult - { - std::vector losses; - std::uint64_t param_hash = 0; - }; +struct RunResult +{ + std::vector losses; + std::uint64_t param_hash = 0; +}; - std::uint64_t fnv1a(std::uint64_t h, const void *data, std::size_t n) +std::uint64_t fnv1a(std::uint64_t h, const void* data, std::size_t n) +{ + const auto* p = static_cast(data); + for (std::size_t i = 0; i < n; ++i) { - const auto *p = static_cast(data); - for (std::size_t i = 0; i < n; ++i) - { - h ^= p[i]; - h *= 1099511628211ULL; - } - return h; + h ^= p[i]; + h *= 1099511628211ULL; } + return h; +} - // 固定公式覆写全部参数（init 的 rng 起点不可跨轮复现，覆写后才可比对） - nn::Result fix_weights(nn::ComputeEngine &eng, nn::Model &model) +// 固定公式覆写全部参数（init 的 rng 起点不可跨轮复现，覆写后才可比对） +nn::Result fix_weights(nn::ComputeEngine& eng, nn::Model& model) +{ + std::size_t idx = 0; + for (auto tr : model.parameters()) { - std::size_t idx = 0; - for (auto tr : model.parameters()) - { - auto &p = tr.get(); - auto pm = eng.to_matrix(p); - NN_TRY_CHECK(pm); - auto sp = pm->span(); - for (std::size_t i = 0; i < sp.size(); ++i) - sp[i] = 0.05f * std::sin(0.61803398875f * static_cast(i + 1) + - 0.017f * static_cast(idx)); - NN_TRY(r, eng.write(p, pm->span())); - ++idx; - } - return {}; + auto& p = tr.get(); + auto pm = eng.to_matrix(p); + NN_TRY_CHECK(pm); + auto sp = pm->span(); + for (std::size_t i = 0; i < sp.size(); ++i) + sp[i] = 0.05f * std::sin(0.61803398875f * static_cast(i + 1) + + 0.017f * static_cast(idx)); + NN_TRY(r, eng.write(p, pm->span())); + ++idx; } + return {}; +} - // 全参数 FNV-1a 校验和（经 to_host 语义回读后按位散列） - nn::Result hash_params(nn::ComputeEngine &eng, nn::Model &model) +// 全参数 FNV-1a 校验和（经 to_host 语义回读后按位散列） +nn::Result hash_params(nn::ComputeEngine& eng, nn::Model& model) +{ + std::uint64_t h = 14695981039346656037ULL; + for (auto tr : model.parameters()) { - std::uint64_t h = 14695981039346656037ULL; - for (auto tr : model.parameters()) - { - auto pm = eng.to_matrix(tr.get()); - NN_TRY_CHECK(pm); - auto sp = pm->span(); - h = fnv1a(h, sp.data(), sp.size() * sizeof(nn::Scalar)); - } - return h; + auto pm = eng.to_matrix(tr.get()); + NN_TRY_CHECK(pm); + auto sp = pm->span(); + h = fnv1a(h, sp.data(), sp.size() * sizeof(nn::Scalar)); } + return h; +} - // 设备引擎（GPU 在首个 GpuBackend::shared() 前已注入 NN_VULKAN_DEVICE） - nn::Result> make_engine(bool gpu) +// 设备引擎（GPU 在首个 GpuBackend::shared() 前已注入 NN_VULKAN_DEVICE） +nn::Result> make_engine(bool gpu) +{ + std::shared_ptr eng; + if (gpu) { - std::shared_ptr eng; - if (gpu) - { #ifdef NN_HAS_VULKAN - auto backend_owner = nn::GpuBackend::shared(); - nn::GpuBackend& backend = *backend_owner; - // 修复（16 §7-2 根因）：shared() 只建 backend，Vulkan 设备要 - // initialize() 才建立——本探针原先漏调，首个 vkCreateBuffer 拿到 - // VK_NULL_HANDLE 即 "Invalid device" 崩溃。与 cli_engine_factory - // 及全部 GPU 测试入口同法。无显式 selector 时 initialize() 内部 - // 仍按 "显式 > NN_VULKAN_DEVICE 环境变量 > 自动打分" 读设备选择。 - NN_TRY(ir, backend.initialize()); - eng = nn::GpuEngine::create(backend); + auto backend_owner = nn::GpuBackend::shared(); + nn::GpuBackend& backend = *backend_owner; + // 修复（16 §7-2 根因）：shared() 只建 backend，Vulkan 设备要 + // initialize() 才建立——本探针原先漏调，首个 vkCreateBuffer 拿到 + // VK_NULL_HANDLE 即 "Invalid device" 崩溃。与 cli_engine_factory + // 及全部 GPU 测试入口同法。无显式 selector 时 initialize() 内部 + // 仍按 "显式 > NN_VULKAN_DEVICE 环境变量 > 自动打分" 读设备选择。 + NN_TRY(ir, backend.initialize()); + eng = nn::GpuEngine::create(backend); #else - NN_FAIL("built without Vulkan (NN_HAS_VULKAN)"); + NN_FAIL("built without Vulkan (NN_HAS_VULKAN)"); #endif - } - else - { - eng = nn::CpuEngine::create(); - } - return eng; } - - // 一轮完整训练：建引擎 → 建模型（覆写权重）→ N 步 AdamW → 回读校验和 - nn::Result run_once(bool gpu, std::size_t steps) + else { - RunResult out; - - auto eng_r = make_engine(gpu); - NN_TRY_CHECK(eng_r); - nn::ComputeEngine &e = **eng_r; - - // 微型配置：两层 GPT，seq=16 batch=2，双轮各 40 次前反向，足够覆盖 - // matmul / attention / embedding / CE / 优化器全链路。 - const std::size_t vocab = 257, seq = 16, batch = 2; - nn::GptConfig cfg{}; - cfg.vocab_size = vocab; - cfg.d_model = 32; - cfg.seq_len = seq; - cfg.num_heads = 4; - cfg.d_ff = 128; - cfg.num_layers = 2; - auto model_r = nn::build_gpt_model(e, cfg); - NN_TRY_CHECK(model_r); - nn::Model model = std::move(*model_r); - - NN_TRY(r, fix_weights(e, model)); - NN_TRY(r2, model.zero_grad()); - auto opt = nn::create_optimizer("adamw", e, model.parameters(), - model.param_gradients(), 1e-3, 0.0); - nn::CrossEntropyLoss ce; - - nn::Matrix x_m(seq, batch); - std::vector labels(seq * batch); - for (std::size_t s = 0; s < steps; ++s) - { - // 每步输入/标签都是确定性公式（不依赖任何 rng） - for (std::size_t i = 0; i < x_m.size(); ++i) - x_m.span()[i] = static_cast((i * 7u + s * 13u + 1u) % vocab); - for (std::size_t i = 0; i < labels.size(); ++i) - labels[i] = (i * 11u + s * 17u + 3u) % vocab; - - NN_TRY(r, model.zero_grad()); - auto x = e.from_matrix(x_m); - NN_TRY_CHECK(x); - auto logits = model.forward(*x); - NN_TRY_CHECK(logits); - auto loss = ce.forward_sparse(e, *logits, labels, {}, vocab); - NN_TRY_CHECK(loss); - auto grad = ce.backward(); - NN_TRY_CHECK(grad); - auto in_grad = model.backward(*grad); - NN_TRY_CHECK(in_grad); - NN_TRY(r2, opt->step()); - out.losses.push_back(static_cast(*loss)); - } - - auto h = hash_params(e, model); - NN_TRY_CHECK(h); - out.param_hash = *h; - return out; + eng = nn::CpuEngine::create(); } + return eng; +} - // ── 初值确定性（M2 专属验收）：建模型后立刻哈希参数（不覆写）──────── - using InitHashes = std::vector>; - - nn::Result hash_model(nn::ComputeEngine &e, const char *name, - nn::Result &&mr, InitHashes &out) +// 一轮完整训练：建引擎 → 建模型（覆写权重）→ N 步 AdamW → 回读校验和 +nn::Result run_once(bool gpu, std::size_t steps) +{ + RunResult out; + + auto eng_r = make_engine(gpu); + NN_TRY_CHECK(eng_r); + nn::ComputeEngine& e = **eng_r; + + // 微型配置：两层 GPT，seq=16 batch=2，双轮各 40 次前反向，足够覆盖 + // matmul / attention / embedding / CE / 优化器全链路。 + const std::size_t vocab = 257, seq = 16, batch = 2; + nn::GptConfig cfg{}; + cfg.vocab_size = vocab; + cfg.d_model = 32; + cfg.seq_len = seq; + cfg.num_heads = 4; + cfg.d_ff = 128; + cfg.num_layers = 2; + auto model_r = nn::build_gpt_model(e, cfg); + NN_TRY_CHECK(model_r); + nn::Model model = std::move(*model_r); + + NN_TRY(r, fix_weights(e, model)); + NN_TRY(r2, model.zero_grad()); + auto opt = + nn::create_optimizer("adamw", e, model.parameters(), model.param_gradients(), 1e-3, 0.0); + nn::CrossEntropyLoss ce; + + nn::Matrix x_m(seq, batch); + std::vector labels(seq * batch); + for (std::size_t s = 0; s < steps; ++s) { - NN_TRY_CHECK(mr); - nn::Model model = std::move(*mr); - auto h = hash_params(e, model); - NN_TRY_CHECK(h); - out.emplace_back(name, *h); - return {}; + // 每步输入/标签都是确定性公式（不依赖任何 rng） + for (std::size_t i = 0; i < x_m.size(); ++i) + x_m.span()[i] = static_cast((i * 7u + s * 13u + 1u) % vocab); + for (std::size_t i = 0; i < labels.size(); ++i) + labels[i] = (i * 11u + s * 17u + 3u) % vocab; + + NN_TRY(r, model.zero_grad()); + auto x = e.from_matrix(x_m); + NN_TRY_CHECK(x); + auto logits = model.forward(*x); + NN_TRY_CHECK(logits); + auto loss = ce.forward_sparse(e, *logits, labels, {}, vocab); + NN_TRY_CHECK(loss); + auto grad = ce.backward(); + NN_TRY_CHECK(grad); + auto in_grad = model.backward(*grad); + NN_TRY_CHECK(in_grad); + NN_TRY(r2, opt->step()); + out.losses.push_back(static_cast(*loss)); } - // 五类模型 = 全部 init 路径（Linear/Conv/Norm 常数/token_emb/位置编码/ - // ones_row_）。同一引擎顺序构建：创建序号（InitSpec 混流）按构造顺序推进， - // 跨进程同序 → 初值逐字节确定。 - // （ZiPT 类已整体移除；锚点表见 AGENTS.md §12 / docs/history/history.md。） - nn::Result run_init_hash(bool gpu) - { - auto eng_r = make_engine(gpu); - NN_TRY_CHECK(eng_r); - nn::ComputeEngine &e = **eng_r; - InitHashes out; - - NN_TRY(r, hash_model(e, "mnist_mlp", nn::build_mnist_mlp_model(e), out)); - nn::CnnConfig ccfg; - ccfg.convs = nn::MNIST_CNN_CONVS; - ccfg.fc_dims = nn::MNIST_CNN_FC; - NN_TRY(r2, hash_model(e, "cnn", nn::build_cnn_model(e, ccfg), out)); - NN_TRY(r3, hash_model(e, "mnist_transformer", - nn::build_mnist_transformer_model(e), out)); - nn::GptConfig gcfg{}; - gcfg.vocab_size = 257; - gcfg.d_model = 32; - gcfg.seq_len = 16; - gcfg.num_heads = 4; - gcfg.d_ff = 128; - gcfg.num_layers = 2; - NN_TRY(r4, hash_model(e, "gpt", nn::build_gpt_model(e, gcfg), out)); - nn::RAPTConfig rcfg; - rcfg.vocab_size = 257; - rcfg.d_model = 32; - rcfg.seq_len = 16; - rcfg.num_heads = 4; - rcfg.d_ff = 64; - rcfg.num_layers = 2; - NN_TRY(r5, hash_model(e, "rapt", nn::build_rapt_model(e, rcfg), out)); - return out; - } + auto h = hash_params(e, model); + NN_TRY_CHECK(h); + out.param_hash = *h; + return out; +} - // ── M3 批量读写语义对拍（17 §5 M3 专属验收）────────────────────────── - // 每个用例覆盖：write → read 逐元素往返（GPU = 上传 copy_from + 下载 - // to_matrix 的 staging 链路）、read 与 to_matrix 两条下载路径逐位一致、 - // get_index/set_index 往返、越界与 U2 精度错配的错误路径。 - constexpr std::size_t kIoRows = 6, kIoCols = 7; +// ── 初值确定性（M2 专属验收）：建模型后立刻哈希参数（不覆写）──────── +using InitHashes = std::vector>; - nn::Result io_f32_case(nn::ComputeEngine &e, const char *tag) - { - auto t = e.create_tensor(kIoRows, kIoCols, nn::Precision::F32, - nn::InitSpec::zero()); - if (!t.valid()) - NN_FAIL("io: create_tensor 失败"); - std::vector pat(kIoRows * kIoCols); - for (std::size_t i = 0; i < pat.size(); ++i) - pat[i] = 0.001f * static_cast((i * 37u) % 1000u) - 0.2f; - - NN_TRY(r5, e.write(t, std::span(pat))); - std::vector back(pat.size()); - NN_TRY(r6, e.read(t, std::span(back))); - if (back != pat) - NN_FAIL(std::string(tag) + ": write/read 往返不一致"); - // read vs to_matrix：两条下载路径逐位一致 - auto m = e.to_matrix(t, nn::Precision::F32); - NN_TRY_CHECK(m); - if (m->span().size() != pat.size()) - NN_FAIL(std::string(tag) + ": to_matrix 尺寸不一致"); - for (std::size_t i = 0; i < pat.size(); ++i) - if (m->span()[i] != pat[i]) - NN_FAIL(std::string(tag) + ": read 与 to_matrix 结果不一致"); - // get_index 往返（含首尾角与中部） - const std::pair pts[] = { - {0, 0}, {kIoRows - 1, kIoCols - 1}, {2, 3}}; - for (const auto &rc : pts) - { - auto v = e.get_index(t, rc.first, rc.second); - NN_TRY_CHECK(v); - if (*v != pat[rc.first * kIoCols + rc.second]) - NN_FAIL(std::string(tag) + ": get_index 不一致"); - } +nn::Result hash_model(nn::ComputeEngine& e, const char* name, nn::Result&& mr, + InitHashes& out) +{ + NN_TRY_CHECK(mr); + nn::Model model = std::move(*mr); + auto h = hash_params(e, model); + NN_TRY_CHECK(h); + out.emplace_back(name, *h); + return {}; +} - // set_index → 批量 read 复核 - NN_TRY(r7, e.set_index(t, 3, 4, -0.125f)); - NN_TRY(r8, e.read(t, std::span(back))); - if (back[3 * kIoCols + 4] != -0.125f) - NN_FAIL(std::string(tag) + ": set_index 未落盘"); - // 错误路径：越界 / U2 精度错配 / 尺寸错配必须报错 - if (e.get_index(t, kIoRows, 0)) - NN_FAIL(std::string(tag) + ": get_index 越界未报错"); - if (e.set_index(t, 0, kIoCols, 1.0f)) - NN_FAIL(std::string(tag) + ": set_index 越界未报错"); - std::vector wrong16(pat.size()); - if (e.read(t, std::span(wrong16))) - NN_FAIL(std::string(tag) + ": read f16 错配未报错"); - if (e.write(t, std::span(wrong16))) - NN_FAIL(std::string(tag) + ": write f16 错配未报错"); - std::vector shortbuf(3); - if (e.read(t, std::span(shortbuf))) - NN_FAIL(std::string(tag) + ": read 尺寸错配未报错"); - return {}; - } +// 五类模型 = 全部 init 路径（Linear/Conv/Norm 常数/token_emb/位置编码/ +// ones_row_）。同一引擎顺序构建：创建序号（InitSpec 混流）按构造顺序推进， +// 跨进程同序 → 初值逐字节确定。 +// （ZiPT 类已整体移除；锚点表见 AGENTS.md §12 / docs/history/history.md。） +nn::Result run_init_hash(bool gpu) +{ + auto eng_r = make_engine(gpu); + NN_TRY_CHECK(eng_r); + nn::ComputeEngine& e = **eng_r; + InitHashes out; + + NN_TRY(r, hash_model(e, "mnist_mlp", nn::build_mnist_mlp_model(e), out)); + nn::CnnConfig ccfg; + ccfg.convs = nn::MNIST_CNN_CONVS; + ccfg.fc_dims = nn::MNIST_CNN_FC; + NN_TRY(r2, hash_model(e, "cnn", nn::build_cnn_model(e, ccfg), out)); + NN_TRY(r3, hash_model(e, "mnist_transformer", nn::build_mnist_transformer_model(e), out)); + nn::GptConfig gcfg{}; + gcfg.vocab_size = 257; + gcfg.d_model = 32; + gcfg.seq_len = 16; + gcfg.num_heads = 4; + gcfg.d_ff = 128; + gcfg.num_layers = 2; + NN_TRY(r4, hash_model(e, "gpt", nn::build_gpt_model(e, gcfg), out)); + nn::RAPTConfig rcfg; + rcfg.vocab_size = 257; + rcfg.d_model = 32; + rcfg.seq_len = 16; + rcfg.num_heads = 4; + rcfg.d_ff = 64; + rcfg.num_layers = 2; + NN_TRY(r5, hash_model(e, "rapt", nn::build_rapt_model(e, rcfg), out)); + return out; +} + +// ── M3 批量读写语义对拍（17 §5 M3 专属验收）────────────────────────── +// 每个用例覆盖：write → read 逐元素往返（GPU = 上传 copy_from + 下载 +// to_matrix 的 staging 链路）、read 与 to_matrix 两条下载路径逐位一致、 +// get_index/set_index 往返、越界与 U2 精度错配的错误路径。 +constexpr std::size_t kIoRows = 6, kIoCols = 7; - nn::Result io_f16_case(nn::ComputeEngine &e, const char *tag) +nn::Result io_f32_case(nn::ComputeEngine& e, const char* tag) +{ + auto t = e.create_tensor(kIoRows, kIoCols, nn::Precision::F32, nn::InitSpec::zero()); + if (!t.valid()) NN_FAIL("io: create_tensor 失败"); + std::vector pat(kIoRows * kIoCols); + for (std::size_t i = 0; i < pat.size(); ++i) + pat[i] = 0.001f * static_cast((i * 37u) % 1000u) - 0.2f; + + NN_TRY(r5, e.write(t, std::span(pat))); + std::vector back(pat.size()); + NN_TRY(r6, e.read(t, std::span(back))); + if (back != pat) NN_FAIL(std::string(tag) + ": write/read 往返不一致"); + // read vs to_matrix：两条下载路径逐位一致 + auto m = e.to_matrix(t, nn::Precision::F32); + NN_TRY_CHECK(m); + if (m->span().size() != pat.size()) NN_FAIL(std::string(tag) + ": to_matrix 尺寸不一致"); + for (std::size_t i = 0; i < pat.size(); ++i) + if (m->span()[i] != pat[i]) NN_FAIL(std::string(tag) + ": read 与 to_matrix 结果不一致"); + // get_index 往返（含首尾角与中部） + const std::pair pts[] = {{0, 0}, {kIoRows - 1, kIoCols - 1}, {2, 3}}; + for (const auto& rc : pts) { - auto t = e.create_tensor(kIoRows, kIoCols, nn::Precision::F16, - nn::InitSpec::zero()); - if (!t.valid()) - NN_FAIL("io: create_tensor(f16) 失败"); - std::vector pat(kIoRows * kIoCols); - for (std::size_t i = 0; i < pat.size(); ++i) - pat[i] = nn::f16(0.01f * static_cast((i * 13u) % 200u) - 1.f); - - NN_TRY(r9, e.write(t, std::span(pat))); - std::vector back(pat.size()); - NN_TRY(r10, e.read(t, std::span(back))); - for (std::size_t i = 0; i < pat.size(); ++i) - if (static_cast(back[i]) != static_cast(pat[i])) - NN_FAIL(std::string(tag) + ": f16 write/read 往返不一致"); - // to_matrix（升 f32 精确无损）与 read(f16) 交叉对拍 - auto m = e.to_matrix(t, nn::Precision::F32); - NN_TRY_CHECK(m); - for (std::size_t i = 0; i < pat.size(); ++i) - if (m->span()[i] != static_cast(pat[i])) - NN_FAIL(std::string(tag) + ": f16 read 与 to_matrix 不一致"); - // 索引往返：set_index 的 Scalar 经 RHE 落 f16（0.5 可精确表示） - auto v = e.get_index(t, 1, 2); + auto v = e.get_index(t, rc.first, rc.second); NN_TRY_CHECK(v); - if (*v != static_cast(pat[1 * kIoCols + 2])) - NN_FAIL(std::string(tag) + ": f16 get_index 不一致"); - NN_TRY(r11, e.set_index(t, 5, 6, 0.5f)); - auto v2 = e.get_index(t, 5, 6); - if (!v2 || *v2 != 0.5f) - NN_FAIL(std::string(tag) + ": f16 set_index 未落盘"); - // U2：f32 span 配 f16 张量必须报错 - std::vector wrong32(pat.size()); - if (e.read(t, std::span(wrong32))) - NN_FAIL(std::string(tag) + ": read f32 错配未报错"); - if (e.write(t, std::span(wrong32))) - NN_FAIL(std::string(tag) + ": write f32 错配未报错"); - return {}; + if (*v != pat[rc.first * kIoCols + rc.second]) + NN_FAIL(std::string(tag) + ": get_index 不一致"); } - // 录制窗口子测（17 §7-4 调用约定）：begin_batch 内 write（覆盖既有存储、 - // drain）→ read（隐含 flush + 同步）→ end_batch → 复读值不变；张量全程 - // 存活到 end_batch 之后（铁律 #6）。 - nn::Result io_batch_window_case(nn::ComputeEngine &e, const char *tag) + // set_index → 批量 read 复核 + NN_TRY(r7, e.set_index(t, 3, 4, -0.125f)); + NN_TRY(r8, e.read(t, std::span(back))); + if (back[3 * kIoCols + 4] != -0.125f) NN_FAIL(std::string(tag) + ": set_index 未落盘"); + // 错误路径：越界 / U2 精度错配 / 尺寸错配必须报错 + if (e.get_index(t, kIoRows, 0)) NN_FAIL(std::string(tag) + ": get_index 越界未报错"); + if (e.set_index(t, 0, kIoCols, 1.0f)) NN_FAIL(std::string(tag) + ": set_index 越界未报错"); + std::vector wrong16(pat.size()); + if (e.read(t, std::span(wrong16))) NN_FAIL(std::string(tag) + ": read f16 错配未报错"); + if (e.write(t, std::span(wrong16))) + NN_FAIL(std::string(tag) + ": write f16 错配未报错"); + std::vector shortbuf(3); + if (e.read(t, std::span(shortbuf))) + NN_FAIL(std::string(tag) + ": read 尺寸错配未报错"); + return {}; +} + +nn::Result io_f16_case(nn::ComputeEngine& e, const char* tag) +{ + auto t = e.create_tensor(kIoRows, kIoCols, nn::Precision::F16, nn::InitSpec::zero()); + if (!t.valid()) NN_FAIL("io: create_tensor(f16) 失败"); + std::vector pat(kIoRows * kIoCols); + for (std::size_t i = 0; i < pat.size(); ++i) + pat[i] = nn::f16(0.01f * static_cast((i * 13u) % 200u) - 1.f); + + NN_TRY(r9, e.write(t, std::span(pat))); + std::vector back(pat.size()); + NN_TRY(r10, e.read(t, std::span(back))); + for (std::size_t i = 0; i < pat.size(); ++i) + if (static_cast(back[i]) != static_cast(pat[i])) + NN_FAIL(std::string(tag) + ": f16 write/read 往返不一致"); + // to_matrix（升 f32 精确无损）与 read(f16) 交叉对拍 + auto m = e.to_matrix(t, nn::Precision::F32); + NN_TRY_CHECK(m); + for (std::size_t i = 0; i < pat.size(); ++i) + if (m->span()[i] != static_cast(pat[i])) + NN_FAIL(std::string(tag) + ": f16 read 与 to_matrix 不一致"); + // 索引往返：set_index 的 Scalar 经 RHE 落 f16（0.5 可精确表示） + auto v = e.get_index(t, 1, 2); + NN_TRY_CHECK(v); + if (*v != static_cast(pat[1 * kIoCols + 2])) + NN_FAIL(std::string(tag) + ": f16 get_index 不一致"); + NN_TRY(r11, e.set_index(t, 5, 6, 0.5f)); + auto v2 = e.get_index(t, 5, 6); + if (!v2 || *v2 != 0.5f) NN_FAIL(std::string(tag) + ": f16 set_index 未落盘"); + // U2：f32 span 配 f16 张量必须报错 + std::vector wrong32(pat.size()); + if (e.read(t, std::span(wrong32))) + NN_FAIL(std::string(tag) + ": read f32 错配未报错"); + if (e.write(t, std::span(wrong32))) + NN_FAIL(std::string(tag) + ": write f32 错配未报错"); + return {}; +} + +// 录制窗口子测（17 §7-4 调用约定）：begin_batch 内 write（覆盖既有存储、 +// drain）→ read（隐含 flush + 同步）→ end_batch → 复读值不变；张量全程 +// 存活到 end_batch 之后（铁律 #6）。 +nn::Result io_batch_window_case(nn::ComputeEngine& e, const char* tag) +{ + auto t = e.create_tensor(4, 5, nn::Precision::F32, nn::InitSpec::zero()); + if (!t.valid()) NN_FAIL("io: create_tensor 失败"); + std::vector pat(20); + for (std::size_t i = 0; i < pat.size(); ++i) + pat[i] = 0.5f + static_cast(i); + + NN_TRY(r12, e.begin_batch()); + NN_TRY(r13, e.write(t, std::span(pat))); + std::vector back(pat.size()); + NN_TRY(r14, e.read(t, std::span(back))); + if (back != pat) NN_FAIL(std::string(tag) + ": 窗口内 write/read 不一致"); + NN_TRY(r15, e.end_batch()); + NN_TRY(r16, e.read(t, std::span(back))); + if (back != pat) NN_FAIL(std::string(tag) + ": end_batch 后复读值变化"); + return {}; +} + +int run_io_roundtrip(bool gpu) +{ + auto eng_r = make_engine(gpu); + if (!eng_r) { - auto t = e.create_tensor(4, 5, nn::Precision::F32, nn::InitSpec::zero()); - if (!t.valid()) - NN_FAIL("io: create_tensor 失败"); - std::vector pat(20); - for (std::size_t i = 0; i < pat.size(); ++i) - pat[i] = 0.5f + static_cast(i); - - NN_TRY(r12, e.begin_batch()); - NN_TRY(r13, e.write(t, std::span(pat))); - std::vector back(pat.size()); - NN_TRY(r14, e.read(t, std::span(back))); - if (back != pat) - NN_FAIL(std::string(tag) + ": 窗口内 write/read 不一致"); - NN_TRY(r15, e.end_batch()); - NN_TRY(r16, e.read(t, std::span(back))); - if (back != pat) - NN_FAIL(std::string(tag) + ": end_batch 后复读值变化"); - return {}; + std::fprintf(stderr, "io-roundtrip 建引擎失败: %s\n", eng_r.error().message.c_str()); + return 2; } + nn::ComputeEngine& e = **eng_r; + using CaseFn = nn::Result (*)(nn::ComputeEngine&, const char*); + const struct + { + const char* name; + CaseFn fn; + } cases[] = { + {"f32_roundtrip", io_f32_case}, + {"f16_roundtrip", io_f16_case}, + {"batch_window", io_batch_window_case}, + }; - int run_io_roundtrip(bool gpu) + int failures = 0; + for (const auto& c : cases) { - auto eng_r = make_engine(gpu); - if (!eng_r) - { - std::fprintf(stderr, "io-roundtrip 建引擎失败: %s\n", - eng_r.error().message.c_str()); - return 2; - } - nn::ComputeEngine &e = **eng_r; - using CaseFn = nn::Result (*)(nn::ComputeEngine &, const char *); - const struct - { - const char *name; - CaseFn fn; - } cases[] = { - {"f32_roundtrip", io_f32_case}, - {"f16_roundtrip", io_f16_case}, - {"batch_window", io_batch_window_case}, - }; - - int failures = 0; - for (const auto &c : cases) + auto r = c.fn(e, c.name); + if (r) { std::printf("IO %s ok=1\n", c.name); } + else { - auto r = c.fn(e, c.name); - if (r) - { - std::printf("IO %s ok=1\n", c.name); - } - else - { - std::printf("IO %s ok=0\n", c.name); - std::fprintf(stderr, "IO %s 失败: %s\n", c.name, - r.error().message.c_str()); - ++failures; - } + std::printf("IO %s ok=0\n", c.name); + std::fprintf(stderr, "IO %s 失败: %s\n", c.name, r.error().message.c_str()); + ++failures; } - std::printf("RESULT io_cases=%zu io_failures=%d verdict=%s\n", - sizeof(cases) / sizeof(cases[0]), failures, - failures == 0 ? "PASS" : "FAIL"); - return failures == 0 ? 0 : 1; } + std::printf("RESULT io_cases=%zu io_failures=%d verdict=%s\n", sizeof(cases) / sizeof(cases[0]), + failures, failures == 0 ? "PASS" : "FAIL"); + return failures == 0 ? 0 : 1; +} } // namespace -static void print_usage(const char *prog) +static void print_usage(const char* prog) { nn::cli::Help help(std::cout, prog, "CPU / GPU 稳定性与初值一致性探针"); @@ -419,7 +396,7 @@ static void print_usage(const char *prog) help.text("0 = 两轮逐位一致；1 = 不一致；2 = 运行错误"); } -int main(int argc, char **argv) +int main(int argc, char** argv) { std::size_t steps = 20; bool gpu = false; @@ -481,17 +458,17 @@ int main(int argc, char **argv) if (gpu) { // MSVC CRT 弃用 getenv（-Werror）：按平台用 _dupenv_s / getenv。 - const char *env = nullptr; + const char* env = nullptr; #if defined(_MSC_VER) - char *buf = nullptr; + char* buf = nullptr; std::size_t len = 0; _dupenv_s(&buf, &len, "NN_VULKAN_DEVICE"); env = buf; #else env = std::getenv("NN_VULKAN_DEVICE"); #endif - std::printf("CONFIG device=%s (NN_VULKAN_DEVICE=%s) steps=%zu\n", - device_desc.c_str(), env ? env : "", steps); + std::printf("CONFIG device=%s (NN_VULKAN_DEVICE=%s) steps=%zu\n", device_desc.c_str(), + env ? env : "", steps); #if defined(_MSC_VER) std::free(buf); #endif @@ -502,8 +479,7 @@ int main(int argc, char **argv) } // ── --io-roundtrip 模式（M3）：批量读写语义对拍（GPU 即 staging 验收）── - if (io_roundtrip) - return run_io_roundtrip(gpu); + if (io_roundtrip) return run_io_roundtrip(gpu); // ── --init-hash 模式（M2）：只建模型比初值，不训练 ──────────────────── if (init_hash) @@ -513,13 +489,13 @@ int main(int argc, char **argv) auto b = run_init_hash(gpu); NN_EXIT(b, 2, "init-hash round2 失败: %s\n"); // 固定前缀行：跨进程比对用 `grep '^INIT1' <(两次启动输出) | diff` - for (const auto &kv : *a) + for (const auto& kv : *a) std::printf("INIT1 %s hash=%016" PRIx64 "\n", kv.first.c_str(), kv.second); - for (const auto &kv : *b) + for (const auto& kv : *b) std::printf("INIT2 %s hash=%016" PRIx64 "\n", kv.first.c_str(), kv.second); const bool match = (*a == *b); - std::printf("RESULT init_in_process_match=%d verdict=%s\n", - match ? 1 : 0, match ? "PASS" : "FAIL"); + std::printf("RESULT init_in_process_match=%d verdict=%s\n", match ? 1 : 0, + match ? "PASS" : "FAIL"); return match ? 0 : 1; } diff --git a/src/gpu_test.cpp b/src/gpu_test.cpp index 0e3e3cca..ca393e69 100644 --- a/src/gpu_test.cpp +++ b/src/gpu_test.cpp @@ -32,11 +32,11 @@ int main() #else #include -using nn::Matrix; using nn::ComputeEngine; using nn::CpuEngine; -using nn::GpuEngine; using nn::GpuBackend; +using nn::GpuEngine; +using nn::Matrix; using nn::Tensor; void print_usage(const char* prog) @@ -67,9 +67,8 @@ Scalar max_abs_diff(const Matrix& a, const Matrix& b) } // ── 辅助：通过 engine 执行 matmul 并下载结果 ─────────────────────────────── -nn::Result engine_matmul( - ComputeEngine& engine, const Matrix& A, const Matrix& B, - bool transA = false, bool transB = false) +nn::Result engine_matmul(ComputeEngine& engine, const Matrix& A, const Matrix& B, + bool transA = false, bool transB = false) { auto a_t = engine.from_matrix(A); NN_TRY_CHECK(a_t); @@ -81,8 +80,7 @@ nn::Result engine_matmul( } // 数字解析辅助：解析失败时打印错误并退出（替代会抛异常的 std::stoi/stod） -template -T parse_num_or_die(const char* s, const char* opt) +template T parse_num_or_die(const char* s, const char* opt) { auto v = nn::parse_number(s); if (!v) @@ -102,25 +100,40 @@ int main(int argc, char* argv[]) for (int i = 1; i < argc; ++i) { std::string arg = argv[i]; - if (arg == "--help") { print_usage(argv[0]); return 0; } + if (arg == "--help") + { + print_usage(argv[0]); + return 0; + } else if (arg == "--size" && i + 1 < argc) { int v = parse_num_or_die(argv[++i], "--size"); - if (v <= 0) { std::cerr << "--size 必须为正整数\n"; return 1; } + if (v <= 0) + { + std::cerr << "--size 必须为正整数\n"; + return 1; + } N = static_cast(v); } else if (arg == "--iters" && i + 1 < argc) { iters = parse_num_or_die(argv[++i], "--iters"); - if (iters <= 0) { std::cerr << "--iters 必须为正整数\n"; return 1; } + if (iters <= 0) + { + std::cerr << "--iters 必须为正整数\n"; + return 1; + } } - else if (auto dev = nn::cli::parse_gpu_option( - argc, argv, i, /*allow_name_value=*/true)) + else if (auto dev = nn::cli::parse_gpu_option(argc, argv, i, /*allow_name_value=*/true)) { // 本程序无位置参数，空格形式也可直接给名称子串 gpu_device = *dev; } - else { std::cerr << "未知参数: " << arg << "\n"; return 1; } + else + { + std::cerr << "未知参数: " << arg << "\n"; + return 1; + } } std::size_t failures = 0; // ✅/❌ 检查计数，决定退出码（CI 必须反映 GPU 状态） @@ -151,8 +164,7 @@ int main(int argc, char* argv[]) auto gpu_engine = GpuEngine::create(backend); // 打印所选物理设备：选错卡（例如把独显老卡当目标）会直接表现为 // 能力不足 / 驱动死锁，先让"跑在哪块卡上"可见。 - std::cout << " 成功 (CpuEngine + GpuEngine, GPU: " - << backend.device().device_name() << ")\n"; + std::cout << " 成功 (CpuEngine + GpuEngine, GPU: " << backend.device().device_name() << ")\n"; // ── 2. 生成随机测试数据 ──────────────────────────────────────── std::cout << "[2/6] 生成 " << N << "×" << N << " 随机矩阵..." << std::flush; @@ -160,8 +172,10 @@ int main(int argc, char* argv[]) std::uniform_real_distribution dist(-1.0f, 1.0f); Matrix A(N, N), B(N, N); - for (auto& v : A.span()) v = dist(rng); - for (auto& v : B.span()) v = dist(rng); + for (auto& v : A.span()) + v = dist(rng); + for (auto& v : B.span()) + v = dist(rng); std::cout << " 完成\n"; // ── 3. matmul 正确性验证 (CpuEngine vs GpuEngine) ───────────── @@ -188,7 +202,8 @@ int main(int argc, char* argv[]) Scalar nrmse = (sum_sq_ref > 0) ? std::sqrt(sum_sq_err / sum_sq_ref) : 0.0f; std::cout << "\n"; - std::cout << " 最大绝对误差: " << std::scientific << std::setprecision(4) << matmul_err << "\n"; + std::cout << " 最大绝对误差: " << std::scientific << std::setprecision(4) << matmul_err + << "\n"; std::cout << " RMSE: " << std::scientific << std::setprecision(4) << rmse << "\n"; std::cout << " NRMSE: " << std::scientific << std::setprecision(4) << nrmse << "\n"; @@ -210,7 +225,11 @@ int main(int argc, char* argv[]) std::cout << "\n A^T * B 最大误差: " << std::scientific << std::setprecision(4) << err << (err < 1e-2f ? " ✅" : " ❌") << "\n"; } - else { std::cout << "\n ❌ A^T * B 失败\n"; ++failures; } + else + { + std::cout << "\n ❌ A^T * B 失败\n"; + ++failures; + } // A * B^T auto cpu_r2 = engine_matmul(*cpu_engine, A, B, false, true); @@ -222,7 +241,11 @@ int main(int argc, char* argv[]) std::cout << " A * B^T 最大误差: " << std::scientific << std::setprecision(4) << err << (err < 1e-2f ? " ✅" : " ❌") << "\n"; } - else { std::cout << " ❌ A * B^T 失败\n"; ++failures; } + else + { + std::cout << " ❌ A * B^T 失败\n"; + ++failures; + } } // ── 3c. 小 N matmul（GEMV 快路径，N ≤ 8 → matmul_gemv）──────────── @@ -230,26 +253,32 @@ int main(int argc, char* argv[]) // transA=1（沿 m vec4）、transB=1（沿 k vec4）及其标量回退组合。 std::cout << "[3c] 小 N matmul（GEMV 路径）...\n"; { - struct GemvCase { std::size_t Am, Ak, Bk, Bn; bool tA, tB; }; + struct GemvCase + { + std::size_t Am, Ak, Bk, Bn; + bool tA, tB; + }; // 字段 = 存储形状（行×列）。维度反推：tA ⇒ M=A.cols, K=A.rows； // tB ⇒ N=B.rows, K=B.cols。故非转置需 Ak==Bk，tA 需 Am==Bk， // tB 需 Bn==Ak。以下用例均已按此校验。 const GemvCase cases[] = { - {256, 512, 512, 1, false, false}, // N=1 纯 GEMV（B 沿 k 扁平 vec4） - {256, 512, 512, 5, false, false}, // N%4≠0 → 扁平 vec4 拆分 - {256, 512, 512, 8, false, false}, // N%4==0 → 按行 vec4 - {255, 513, 513, 7, false, false}, // 奇数 M/K：行尾 + k 尾块标量 - {256, 64, 256, 8, true, false}, // tA：M=64,K=256（M%4==0 → 沿 m vec4） - {257, 65, 257, 6, true, false}, // tA：M=65,K=257 奇数 → 标量回退 - {128, 384, 8, 384, false, true}, // tB：N=8,K=384（K%4==0 → 沿 k vec4） - {130, 386, 3, 386, false, true}, // tB：N=3,K=386 奇数 K → 标量回退 - {257, 64, 4, 257, true, true}, // 双转置：M=64,K=257,N=4 + {256, 512, 512, 1, false, false}, // N=1 纯 GEMV（B 沿 k 扁平 vec4） + {256, 512, 512, 5, false, false}, // N%4≠0 → 扁平 vec4 拆分 + {256, 512, 512, 8, false, false}, // N%4==0 → 按行 vec4 + {255, 513, 513, 7, false, false}, // 奇数 M/K：行尾 + k 尾块标量 + {256, 64, 256, 8, true, false}, // tA：M=64,K=256（M%4==0 → 沿 m vec4） + {257, 65, 257, 6, true, false}, // tA：M=65,K=257 奇数 → 标量回退 + {128, 384, 8, 384, false, true}, // tB：N=8,K=384（K%4==0 → 沿 k vec4） + {130, 386, 3, 386, false, true}, // tB：N=3,K=386 奇数 K → 标量回退 + {257, 64, 4, 257, true, true}, // 双转置：M=64,K=257,N=4 }; for (const auto& c : cases) { Matrix a(c.Am, c.Ak), b(c.Bk, c.Bn); - for (auto& v : a.span()) v = dist(rng); - for (auto& v : b.span()) v = dist(rng); + for (auto& v : a.span()) + v = dist(rng); + for (auto& v : b.span()) + v = dist(rng); auto cpu_r = engine_matmul(*cpu_engine, a, b, c.tA, c.tB); auto gpu_r = engine_matmul(*gpu_engine, a, b, c.tA, c.tB); @@ -261,11 +290,15 @@ int main(int argc, char* argv[]) // N 由存储形状反推：tB ⇒ N=B.rows，否则 N=B.cols const std::size_t n_out = c.tB ? c.Bk : c.Bn; std::cout << " A(" << c.Am << "x" << c.Ak << ") B(" << c.Bk << "x" << c.Bn - << ") tA=" << c.tA << " tB=" << c.tB - << " N=" << n_out << " err=" << std::scientific << std::setprecision(4) << err + << ") tA=" << c.tA << " tB=" << c.tB << " N=" << n_out + << " err=" << std::scientific << std::setprecision(4) << err << (ok ? " ✅" : " ❌") << "\n"; } - else { std::cout << " ❌ GEMV case 执行失败\n"; ++failures; } + else + { + std::cout << " ❌ GEMV case 执行失败\n"; + ++failures; + } } } @@ -277,11 +310,12 @@ int main(int argc, char* argv[]) // 奇数边界（513、1000 非 tile 整倍数）/ 非方阵。 std::cout << "[3d] transpose 正确性 (GPU vs CPU 参考, 逐元素)...\n"; { - struct TpCase { std::size_t R, C; }; + struct TpCase + { std::size_t R, C; }; const TpCase cases[] = { - { 64, 64}, { 512, 512}, { 513, 513}, // 方阵 / 奇数边界 - { 256, 1024}, {1024, 256}, // 单边 ≤512 - { 512, 1024}, {1024, 512}, // 单边恰好 512 + {64, 64}, {512, 512}, {513, 513}, // 方阵 / 奇数边界 + {256, 1024}, {1024, 256}, // 单边 ≤512 + {512, 1024}, {1024, 512}, // 单边恰好 512 {1000, 2000}, {2000, 1000}, // 非 tile 整倍数双边 >512 {1024, 1024}, {1024, 4096}, {4096, 1024}, // 双边 >512 {2048, 2048}, {4096, 4096}, // 大方阵 @@ -289,15 +323,18 @@ int main(int argc, char* argv[]) for (const auto& c : cases) { Matrix a(c.R, c.C); - for (auto& v : a.span()) v = dist(rng); + for (auto& v : a.span()) + v = dist(rng); auto a_t = gpu_engine->from_matrix(a); - auto r_t = a_t ? gpu_engine->transpose(*a_t) : nn::Result{std::unexpected(a_t.error())}; - auto out_r = r_t ? gpu_engine->to_matrix(*r_t) : nn::Result{std::unexpected(r_t.error())}; + auto r_t = a_t ? gpu_engine->transpose(*a_t) + : nn::Result{std::unexpected(a_t.error())}; + auto out_r = r_t ? gpu_engine->to_matrix(*r_t) + : nn::Result{std::unexpected(r_t.error())}; if (!out_r) { - std::cout << " " << c.R << "x" << c.C << " 执行失败: " - << out_r.error().message << "\n"; + std::cout << " " << c.R << "x" << c.C << " 执行失败: " << out_r.error().message + << "\n"; ++failures; continue; } @@ -315,19 +352,26 @@ int main(int argc, char* argv[]) for (std::size_t col = 0; col < c.C; ++col) if (os[col * c.R + r] != as[r * c.C + col]) { - if (!found_first) { first_r = r; first_c = col; found_first = true; } + if (!found_first) + { + first_r = r; + first_c = col; + found_first = true; + } ++bad; } } const std::size_t total = c.R * c.C; - const double pct = total ? 100.0 * static_cast(bad) / static_cast(total) : 0.0; + const double pct = + total ? 100.0 * static_cast(bad) / static_cast(total) : 0.0; const bool ok = shape_ok && bad == 0; if (!ok) ++failures; std::cout << " " << c.R << "x" << c.C - << (shape_ok ? "" : " 形状错误(out=" + std::to_string(out.rows()) + "x" + std::to_string(out.cols()) + ")") + << (shape_ok ? "" + : " 形状错误(out=" + std::to_string(out.rows()) + "x" + + std::to_string(out.cols()) + ")") << " 失配 " << std::fixed << std::setprecision(1) << pct << "%"; - if (found_first) - std::cout << " 首个错误下标 r=" << first_r << " c=" << first_c; + if (found_first) std::cout << " 首个错误下标 r=" << first_r << " c=" << first_c; std::cout << (ok ? " ✅" : " ❌") << "\n"; } } @@ -344,13 +388,20 @@ int main(int argc, char* argv[]) { Scalar err = max_abs_diff(A, *a_back_r); if (err != 0.0f) ++failures; - std::cout << " Upload→Download roundtrip 最大误差: " - << std::scientific << std::setprecision(4) << err - << (err == 0.0f ? " ✅" : " ❌") << "\n"; + std::cout << " Upload→Download roundtrip 最大误差: " << std::scientific + << std::setprecision(4) << err << (err == 0.0f ? " ✅" : " ❌") << "\n"; + } + else + { + std::cout << " ❌ roundtrip download 失败: " << a_back_r.error().message << "\n"; + ++failures; } - else { std::cout << " ❌ roundtrip download 失败: " << a_back_r.error().message << "\n"; ++failures; } } - else { std::cout << " ❌ roundtrip upload 失败\n"; ++failures; } + else + { + std::cout << " ❌ roundtrip upload 失败\n"; + ++failures; + } // 4b. CPU/GPU 张量准备（元素级表达式对拍见 expr_gpu_test） auto a_cpu_t = cpu_engine->from_matrix(A); @@ -369,12 +420,16 @@ int main(int argc, char* argv[]) { Scalar err = max_abs_diff(*sc, *sg); if (err >= 1e-3f) ++failures; - std::cout << " col_reduce_sum(A) 最大误差: " - << std::scientific << std::setprecision(4) << err - << (err < 1e-3f ? " ✅" : " ❌") << "\n"; + std::cout << " col_reduce_sum(A) 最大误差: " << std::scientific + << std::setprecision(4) << err << (err < 1e-3f ? " ✅" : " ❌") + << "\n"; } } - else { std::cout << " ❌ col_reduce_sum 失败\n"; ++failures; } + else + { + std::cout << " ❌ col_reduce_sum 失败\n"; + ++failures; + } // 4e. 归约: col_reduce_max(A) → (1, N) auto cmax_cpu_r = cpu_engine->col_reduce_max(*a_cpu_t); @@ -387,12 +442,16 @@ int main(int argc, char* argv[]) { Scalar err = max_abs_diff(*mc2, *mg2); if (err >= 1e-5f) ++failures; - std::cout << " col_reduce_max(A) 最大误差: " - << std::scientific << std::setprecision(4) << err - << (err < 1e-5f ? " ✅" : " ❌") << "\n"; + std::cout << " col_reduce_max(A) 最大误差: " << std::scientific + << std::setprecision(4) << err << (err < 1e-5f ? " ✅" : " ❌") + << "\n"; } } - else { std::cout << " ❌ col_reduce_max 失败\n"; ++failures; } + else + { + std::cout << " ❌ col_reduce_max 失败\n"; + ++failures; + } } } @@ -407,7 +466,8 @@ int main(int argc, char* argv[]) NN_EXIT(r, 1, "CPU matmul 性能测试失败\n"); } auto t_cpu_end = std::chrono::high_resolution_clock::now(); - double cpu_ms = std::chrono::duration(t_cpu_end - t_cpu_start).count() / iters; + double cpu_ms = + std::chrono::duration(t_cpu_end - t_cpu_start).count() / iters; // GPU 性能（含上传/下载开销，反映端到端延迟） auto t_gpu_start = std::chrono::high_resolution_clock::now(); @@ -417,17 +477,22 @@ int main(int argc, char* argv[]) NN_EXIT(r, 1, "GPU matmul 性能测试失败\n"); } auto t_gpu_end = std::chrono::high_resolution_clock::now(); - double gpu_ms = std::chrono::duration(t_gpu_end - t_gpu_start).count() / iters; + double gpu_ms = + std::chrono::duration(t_gpu_end - t_gpu_start).count() / iters; // 防御：计时为 0 时（极小矩阵/低精度时钟）避免除零输出 inf const double speedup = (gpu_ms > 0.0) ? cpu_ms / gpu_ms : 0.0; const double gflops = (gpu_ms > 0.0) ? (2.0 * N * N * N) / (gpu_ms * 1e6) : 0.0; std::cout << "\n"; - std::cout << " CPU 时间 (含上传/下载): " << std::fixed << std::setprecision(2) << cpu_ms << " ms\n"; - std::cout << " GPU 时间 (含上传/下载): " << std::fixed << std::setprecision(2) << gpu_ms << " ms\n"; - std::cout << " 加速比: " << std::fixed << std::setprecision(2) << speedup << "×\n"; - std::cout << " GPU GFLOPS (含传输开销): " << std::fixed << std::setprecision(1) << gflops << "\n\n"; + std::cout << " CPU 时间 (含上传/下载): " << std::fixed << std::setprecision(2) << cpu_ms + << " ms\n"; + std::cout << " GPU 时间 (含上传/下载): " << std::fixed << std::setprecision(2) << gpu_ms + << " ms\n"; + std::cout << " 加速比: " << std::fixed << std::setprecision(2) << speedup + << "×\n"; + std::cout << " GPU GFLOPS (含传输开销): " << std::fixed << std::setprecision(1) << gflops + << "\n\n"; // ── 6. batch 模式测试 ───────────────────────────────────────── std::cout << "[6/6] batch 模式 + 链式 matmul 测试...\n"; @@ -451,17 +516,33 @@ int main(int argc, char* argv[]) { Scalar err = max_abs_diff(*cpu_c_r, *c_res); if (err >= 1e-2f) ++failures; - std::cout << " Batch+matmul 最大误差: " - << std::scientific << std::setprecision(4) << err + std::cout << " Batch+matmul 最大误差: " << std::scientific + << std::setprecision(4) << err << (err < 1e-2f ? " ✅" : " ❌") << "\n"; } - else { std::cout << " ❌ download failed: " << c_res.error().message << "\n"; ++failures; } + else + { + std::cout << " ❌ download failed: " << c_res.error().message << "\n"; + ++failures; + } + } + else + { + std::cout << " ❌ end_batch failed: " << end_r.error().message << "\n"; + ++failures; } - else { std::cout << " ❌ end_batch failed: " << end_r.error().message << "\n"; ++failures; } } - else { std::cout << " ❌ matmul failed: " << c_t.error().message << "\n"; ++failures; } + else + { + std::cout << " ❌ matmul failed: " << c_t.error().message << "\n"; + ++failures; + } + } + else + { + std::cout << " ❌ begin_batch failed: " << batch_r.error().message << "\n"; + ++failures; } - else { std::cout << " ❌ begin_batch failed: " << batch_r.error().message << "\n"; ++failures; } } // 链式: matmul(A, B) → col_reduce_sum @@ -491,15 +572,16 @@ int main(int argc, char* argv[]) Scalar c_min = std::numeric_limits::max(); Scalar c_max = std::numeric_limits::lowest(); Scalar c_abs_max = 0; - for (auto v : c_cpu_m->span()) { + for (auto v : c_cpu_m->span()) + { c_min = std::min(c_min, v); c_max = std::max(c_max, v); c_abs_max = std::max(c_abs_max, std::fabs(v)); } std::cout << " [1] matmul 输出范围: [" << std::fixed << std::setprecision(2) << c_min << ", " << c_max << "] |max|=" << c_abs_max << "\n"; - std::cout << " matmul CPU vs GPU 最大误差: " << std::scientific << std::setprecision(4) - << mm_err << "\n"; + std::cout << " matmul CPU vs GPU 最大误差: " << std::scientific + << std::setprecision(4) << mm_err << "\n"; // ── Step 2: 直接复用 matmul 输出 ── auto exp_t = c_t; @@ -515,15 +597,17 @@ int main(int argc, char* argv[]) Scalar e_min = std::numeric_limits::max(); Scalar e_max = std::numeric_limits::lowest(); Scalar e_abs_max = 0; - for (auto v : exp_cpu_m->span()) { + for (auto v : exp_cpu_m->span()) + { e_min = std::min(e_min, v); e_max = std::max(e_max, v); e_abs_max = std::max(e_abs_max, std::fabs(v)); } - std::cout << " [2] matmul 输出范围: [" << std::scientific << std::setprecision(4) - << e_min << ", " << e_max << "] |max|=" << e_abs_max << "\n"; - std::cout << " matmul CPU vs GPU 最大误差: " << std::scientific << std::setprecision(4) - << exp_err << "\n"; + std::cout << " [2] matmul 输出范围: [" << std::scientific + << std::setprecision(4) << e_min << ", " << e_max + << "] |max|=" << e_abs_max << "\n"; + std::cout << " matmul CPU vs GPU 最大误差: " << std::scientific + << std::setprecision(4) << exp_err << "\n"; // ── Step 3: col_reduce_sum(matmul) on GPU ── auto sum_t = gpu_engine->col_reduce_sum(*exp_t); @@ -537,23 +621,29 @@ int main(int argc, char* argv[]) // 统计 reduce 输出范围 Scalar s_min = std::numeric_limits::max(); Scalar s_max = std::numeric_limits::lowest(); - for (auto v : res_cpu->span()) { + for (auto v : res_cpu->span()) + { s_min = std::min(s_min, v); s_max = std::max(s_max, v); } - std::cout << " [3] col_reduce_sum 输出范围: [" << std::scientific << std::setprecision(4) - << s_min << ", " << s_max << "]\n"; - std::cout << " 最终绝对误差: " << std::scientific << std::setprecision(4) - << final_err << "\n"; + std::cout << " [3] col_reduce_sum 输出范围: [" + << std::scientific << std::setprecision(4) << s_min + << ", " << s_max << "]\n"; + std::cout << " 最终绝对误差: " << std::scientific + << std::setprecision(4) << final_err << "\n"; // 相对误差（更合理的评判标准） Scalar ref_abs_max = 0; - for (auto v : res_cpu->span()) ref_abs_max = std::max(ref_abs_max, std::fabs(v)); - Scalar rel_err = (ref_abs_max > 0) ? (final_err / ref_abs_max) : 0; - std::cout << " 相对误差: " << std::scientific << std::setprecision(4) - << rel_err << (rel_err < 1e-3f ? " OK" : " FAIL") << "\n"; - std::cout << " 结论: 归约链最终误差 " << std::scientific << std::setprecision(2) - << final_err << "（Exp 级已随逐元素算子移除）\n"; + for (auto v : res_cpu->span()) + ref_abs_max = std::max(ref_abs_max, std::fabs(v)); + Scalar rel_err = + (ref_abs_max > 0) ? (final_err / ref_abs_max) : 0; + std::cout << " 相对误差: " << std::scientific + << std::setprecision(4) << rel_err + << (rel_err < 1e-3f ? " OK" : " FAIL") << "\n"; + std::cout << " 结论: 归约链最终误差 " << std::scientific + << std::setprecision(2) << final_err + << "（Exp 级已随逐元素算子移除）\n"; } } } @@ -574,8 +664,20 @@ int main(int argc, char* argv[]) std::mt19937 rng(2026); std::uniform_real_distribution cop(-1.0f, 1.0f); std::uniform_real_distribution cog(-0.1f, 0.1f); - const Matrix p0 = [&]{ Matrix m(R, C); for (auto& v : m.span()) v = cop(rng); return m; }(); - const Matrix g0 = [&]{ Matrix m(R, C); for (auto& v : m.span()) v = cog(rng); return m; }(); + const Matrix p0 = [&] + { + Matrix m(R, C); + for (auto& v : m.span()) + v = cop(rng); + return m; + }(); + const Matrix g0 = [&] + { + Matrix m(R, C); + for (auto& v : m.span()) + v = cog(rng); + return m; + }(); for (const char* name : {"adam", "adamw"}) { @@ -588,12 +690,12 @@ int main(int argc, char* argv[]) std::unique_ptr opt_cpu, opt_gpu; if (ok) { - opt_cpu = nn::create_optimizer(name, *cpu_engine, - std::vector{*p_cpu_r}, std::vector{*g_cpu_r}, - Scalar{1e-3f}); - opt_gpu = nn::create_optimizer(name, *gpu_engine, - std::vector{*p_gpu_r}, std::vector{*g_gpu_r}, - Scalar{1e-3f}); + opt_cpu = + nn::create_optimizer(name, *cpu_engine, std::vector{*p_cpu_r}, + std::vector{*g_cpu_r}, Scalar{1e-3f}); + opt_gpu = + nn::create_optimizer(name, *gpu_engine, std::vector{*p_gpu_r}, + std::vector{*g_gpu_r}, Scalar{1e-3f}); } if (!ok || !opt_cpu || !opt_gpu) { @@ -607,15 +709,25 @@ int main(int argc, char* argv[]) // 每步用同一梯度（两引擎一致），先后向 zero_grad 再 step step_ok = opt_cpu->step() && opt_gpu->step(); } - if (!step_ok) { std::cout << " ❌ " << name << " step 失败\n"; ++failures; continue; } + if (!step_ok) + { + std::cout << " ❌ " << name << " step 失败\n"; + ++failures; + continue; + } auto p_cpu_m = cpu_engine->to_matrix(*p_cpu_r); auto p_gpu_m = gpu_engine->to_matrix(*p_gpu_r); - if (!p_cpu_m || !p_gpu_m) { std::cout << " ❌ " << name << " 下载失败\n"; ++failures; continue; } + if (!p_cpu_m || !p_gpu_m) + { + std::cout << " ❌ " << name << " 下载失败\n"; + ++failures; + continue; + } Scalar err = max_abs_diff(*p_cpu_m, *p_gpu_m); bool pass = err < 2e-3f; - std::cout << " " << (pass ? "[PASS]" : "[FAIL]") << " " << name << " p_err=" - << std::scientific << std::setprecision(2) << err << "\n"; + std::cout << " " << (pass ? "[PASS]" : "[FAIL]") << " " << name + << " p_err=" << std::scientific << std::setprecision(2) << err << "\n"; if (!pass) ++failures; } } @@ -632,33 +744,51 @@ int main(int argc, char* argv[]) std::mt19937 rng(777); std::normal_distribution cop(0.0f, 0.5f); - auto fill = [&](Matrix& m) { for (auto& v : m.span()) v = cop(rng); }; + auto fill = [&](Matrix& m) + { + for (auto& v : m.span()) + v = cop(rng); + }; Matrix Kv(rows, seq), Vv(rows, seq), Pv(rows, seq), Rv(rows, seq); Matrix Dm(BH * dk * dk, seq), Xs(rows, seq), Ys(rows, seq); Matrix A0v(H * dk, dk), B0v(H * dk, dk), bnd(1, B * seq); - fill(Kv); fill(Vv); fill(Pv); fill(Rv); fill(Dm); fill(Xs); fill(Ys); - (void)A0v; (void)B0v; (void)bnd; + fill(Kv); + fill(Vv); + fill(Pv); + fill(Rv); + fill(Dm); + fill(Xs); + fill(Ys); + (void)A0v; + (void)B0v; + (void)bnd; auto run_prefix = [&](ComputeEngine& e, const Matrix& Ki, const Matrix& Vi, - const Matrix& Pi, const Matrix& Ri) -> nn::Result { + const Matrix& Pi, const Matrix& Ri) -> nn::Result + { NN_TRY(k, e.from_matrix(Ki)); NN_TRY(v, e.from_matrix(Vi)); NN_TRY(p, e.from_matrix(Pi)); NN_TRY(r, e.from_matrix(Ri)); - auto a0 = e.from_matrix(A0v); const auto& a0r = *a0; - auto b0 = e.from_matrix(B0v); const auto& b0r = *b0; - auto bd = e.from_matrix(bnd); const auto& bdr = *bd; - auto out = e.scan_prefix_outer(*k, *v, *p, *r, a0r, b0r, false, - dk, H, /*causal=*/true, bdr, false); + auto a0 = e.from_matrix(A0v); + const auto& a0r = *a0; + auto b0 = e.from_matrix(B0v); + const auto& b0r = *b0; + auto bd = e.from_matrix(bnd); + const auto& bdr = *bd; + auto out = e.scan_prefix_outer(*k, *v, *p, *r, a0r, b0r, false, dk, H, /*causal=*/true, + bdr, false); NN_TRY_CHECK(out); return e.to_matrix(*out); }; auto run_suffix = [&](ComputeEngine& e, const Matrix& Di, const Matrix& Xi, - const Matrix& Yi) -> nn::Result { + const Matrix& Yi) -> nn::Result + { NN_TRY(d, e.from_matrix(Di)); NN_TRY(x, e.from_matrix(Xi)); NN_TRY(y, e.from_matrix(Yi)); - auto bd = e.from_matrix(bnd); const auto& bdr = *bd; + auto bd = e.from_matrix(bnd); + const auto& bdr = *bd; auto out = e.scan_suffix_outer(*d, *x, *y, dk, H, /*causal=*/true, bdr, false); NN_TRY_CHECK(out); return e.to_matrix(*out); @@ -678,8 +808,8 @@ int main(int argc, char* argv[]) Scalar err = max_abs_diff(*cpu_p, *gpu_p); bool pass = err < 2e-3f; std::cout << " " << (pass ? "[PASS]" : "[FAIL]") - << " prefix_outer(dk=128) err=" - << std::scientific << std::setprecision(2) << err << "\n"; + << " prefix_outer(dk=128) err=" << std::scientific << std::setprecision(2) + << err << "\n"; if (!pass) ++failures; } @@ -697,8 +827,8 @@ int main(int argc, char* argv[]) Scalar err = max_abs_diff(*cpu_s, *gpu_s); bool pass = err < 2e-3f; std::cout << " " << (pass ? "[PASS]" : "[FAIL]") - << " suffix_outer(dk=128) err=" - << std::scientific << std::setprecision(2) << err << "\n"; + << " suffix_outer(dk=128) err=" << std::scientific << std::setprecision(2) + << err << "\n"; if (!pass) ++failures; } } diff --git a/src/host_bridge_test.cpp b/src/host_bridge_test.cpp index 1657f187..b2da88ee 100644 --- a/src/host_bridge_test.cpp +++ b/src/host_bridge_test.cpp @@ -18,12 +18,14 @@ #include int g_fail = 0; -#define CHECK(cond, msg) \ - do { \ - if (!(cond)) { \ - std::printf("[FAIL] %s:%d: %s\n", __FILE__, __LINE__, (msg)); \ - ++g_fail; \ - } \ +#define CHECK(cond, msg) \ + do \ + { \ + if (!(cond)) \ + { \ + std::printf("[FAIL] %s:%d: %s\n", __FILE__, __LINE__, (msg)); \ + ++g_fail; \ + } \ } while (0) int main() @@ -56,7 +58,8 @@ int main() CHECK(t->precision() == nn::Precision::F16, "F16 目标精度"); // 对拍 from_matrix 同值（RHE 同路径） nn::Matrix m(2, 2); - for (std::size_t i = 0; i < src.size(); ++i) m.span()[i] = src[i]; + for (std::size_t i = 0; i < src.size(); ++i) + m.span()[i] = src[i]; auto ref = engine->from_matrix(m, nn::Precision::F16); CHECK(ref.has_value(), "from_matrix F16 应成功"); if (ref) @@ -77,8 +80,8 @@ int main() std::vector src{1.0f, 2.0f}; auto t = nn::upload_span(*engine, 2, 3, nn::Precision::F32, src); CHECK(!t.has_value(), "元素数不匹配应报错"); - auto t2 = nn::upload_span(*engine, 1, 2, nn::Precision::F32, - std::vector{1.0f, 2.0f}); + auto t2 = + nn::upload_span(*engine, 1, 2, nn::Precision::F32, std::vector{1.0f, 2.0f}); CHECK(t2.has_value(), "匹配尺寸应成功"); if (t2) { diff --git a/src/layer_bench.cpp b/src/layer_bench.cpp index 0025cd7b..6029264b 100644 --- a/src/layer_bench.cpp +++ b/src/layer_bench.cpp @@ -32,11 +32,11 @@ #include #include -using nn::Scalar; -using nn::Matrix; -using nn::Tensor; using nn::ComputeEngine; +using nn::Matrix; using nn::Model; +using nn::Scalar; +using nn::Tensor; namespace { @@ -80,9 +80,7 @@ struct BenchConfig // 工具 // ═══════════════════════════════════════════════════════════════════════════ double ms_since(std::chrono::steady_clock::time_point t0) -{ - return std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); -} +{ return std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); } // 随机填充 Matrix 并上传到引擎设备 Tensor make_input(ComputeEngine& e, std::size_t rows, std::size_t cols) @@ -94,8 +92,7 @@ Tensor make_input(ComputeEngine& e, std::size_t rows, std::size_t cols) for (std::size_t i = 0; i < s.size(); ++i) s[i] = d(g); auto t = e.from_matrix(std::move(m)); - if (!t) - std::abort(); + if (!t) std::abort(); return std::move(*t); } @@ -123,18 +120,24 @@ bool build_linear(Model& m, const BenchConfig& c) auto r = m.template add(c.in, c.out); return r.has_value(); } -std::pair shape_linear(const BenchConfig& c) { return {c.in, c.batch}; } -double flops_linear(const BenchConfig& c) { return 2.0 * c.in * c.out * c.batch; } -std::size_t tok_linear(const BenchConfig& c) { return c.batch; } +std::pair shape_linear(const BenchConfig& c) +{ return {c.in, c.batch}; } +double flops_linear(const BenchConfig& c) +{ return 2.0 * c.in * c.out * c.batch; } +std::size_t tok_linear(const BenchConfig& c) +{ return c.batch; } bool build_swiglu(Model& m, const BenchConfig& c) { auto r = m.template add(c.dff); return r.has_value(); } -std::pair shape_swiglu(const BenchConfig& c) { return {2 * c.dff, c.batch}; } -double flops_swiglu(const BenchConfig& c) { return 8.0 * 2 * c.dff * c.batch; } -std::size_t tok_swiglu(const BenchConfig& c) { return c.batch; } +std::pair shape_swiglu(const BenchConfig& c) +{ return {2 * c.dff, c.batch}; } +double flops_swiglu(const BenchConfig& c) +{ return 8.0 * 2 * c.dff * c.batch; } +std::size_t tok_swiglu(const BenchConfig& c) +{ return c.batch; } bool build_layernorm(Model& m, const BenchConfig& c) { @@ -146,18 +149,24 @@ bool build_rmsnorm(Model& m, const BenchConfig& c) auto r = m.template add(c.dmodel); return r.has_value(); } -std::pair shape_norm(const BenchConfig& c) { return {c.dmodel, c.batch * c.seq}; } -double flops_norm(const BenchConfig& c) { return 6.0 * c.dmodel * c.batch * c.seq; } -std::size_t tok_norm(const BenchConfig& c) { return c.batch * c.seq; } +std::pair shape_norm(const BenchConfig& c) +{ return {c.dmodel, c.batch * c.seq}; } +double flops_norm(const BenchConfig& c) +{ return 6.0 * c.dmodel * c.batch * c.seq; } +std::size_t tok_norm(const BenchConfig& c) +{ return c.batch * c.seq; } bool build_softmax(Model& m, [[maybe_unused]] const BenchConfig& c) { auto r = m.template add(); return r.has_value(); } -std::pair shape_softmax(const BenchConfig& c) { return {c.rows, c.cols}; } -double flops_softmax(const BenchConfig& c) { return 6.0 * c.rows * c.cols; } -std::size_t tok_softmax(const BenchConfig& c) { return c.rows > 0 ? 1u : 0u; } +std::pair shape_softmax(const BenchConfig& c) +{ return {c.rows, c.cols}; } +double flops_softmax(const BenchConfig& c) +{ return 6.0 * c.rows * c.cols; } +std::size_t tok_softmax(const BenchConfig& c) +{ return c.rows > 0 ? 1u : 0u; } bool build_mha(Model& m, const BenchConfig& c) { @@ -169,13 +178,15 @@ bool build_causal(Model& m, const BenchConfig& c) auto r = m.template add(c.dmodel, c.heads, c.seq, c.seq); return r.has_value(); } -std::pair shape_attn(const BenchConfig& c) { return {c.dmodel, c.batch * c.seq}; } +std::pair shape_attn(const BenchConfig& c) +{ return {c.dmodel, c.batch * c.seq}; } double flops_attn(const BenchConfig& c) { const double T = static_cast(c.batch) * c.seq; return 6.0 * c.dmodel * c.dmodel * T + 4.0 * T * T * c.dmodel; } -std::size_t tok_attn(const BenchConfig& c) { return c.batch * c.seq; } +std::size_t tok_attn(const BenchConfig& c) +{ return c.batch * c.seq; } bool build_feedforward(Model& m, const BenchConfig& c) { @@ -183,9 +194,7 @@ bool build_feedforward(Model& m, const BenchConfig& c) return r.has_value(); } double flops_feedforward(const BenchConfig& c) -{ - return 4.0 * c.dmodel * c.dff * c.batch * c.seq; -} +{ return 4.0 * c.dmodel * c.dff * c.batch * c.seq; } bool build_transformer(Model& m, const BenchConfig& c) { @@ -193,9 +202,7 @@ bool build_transformer(Model& m, const BenchConfig& c) return r.has_value(); } double flops_transformer(const BenchConfig& c) -{ - return flops_attn(c) + flops_feedforward(c) + flops_norm(c) * 2.0; -} +{ return flops_attn(c) + flops_feedforward(c) + flops_norm(c) * 2.0; } bool build_gpt_block(Model& m, const BenchConfig& c) { @@ -205,29 +212,32 @@ bool build_gpt_block(Model& m, const BenchConfig& c) bool build_conv(Model& m, const BenchConfig& c) { - auto r = m.template add(c.cin, c.cout, c.kernel, /*stride=*/1, /*pad=*/c.kernel / 2, c.hw, c.hw); + auto r = m.template add(c.cin, c.cout, c.kernel, /*stride=*/1, /*pad=*/c.kernel / 2, + c.hw, c.hw); return r.has_value(); } -std::pair shape_conv(const BenchConfig& c) { return {c.cin * c.hw * c.hw, c.batch}; } +std::pair shape_conv(const BenchConfig& c) +{ return {c.cin * c.hw * c.hw, c.batch}; } double flops_conv(const BenchConfig& c) -{ - return 2.0 * c.cout * c.cin * c.kernel * c.kernel * c.hw * c.hw * c.batch; -} -std::size_t tok_conv(const BenchConfig& c) { return c.batch; } +{ return 2.0 * c.cout * c.cin * c.kernel * c.kernel * c.hw * c.hw * c.batch; } +std::size_t tok_conv(const BenchConfig& c) +{ return c.batch; } bool build_maxpool(Model& m, const BenchConfig& c) { auto r = m.template add(c.cin, c.hw, c.hw, /*pool=*/2, /*stride=*/2); return r.has_value(); } -std::pair shape_maxpool(const BenchConfig& c) { return {c.cin * c.hw * c.hw, c.batch}; } +std::pair shape_maxpool(const BenchConfig& c) +{ return {c.cin * c.hw * c.hw, c.batch}; } double flops_maxpool(const BenchConfig& c) { const std::size_t oh = (c.hw - 2) / 2 + 1; const std::size_t ow = (c.hw - 2) / 2 + 1; return 1.0 * c.cin * oh * ow * c.batch; } -std::size_t tok_maxpool(const BenchConfig& c) { return c.batch; } +std::size_t tok_maxpool(const BenchConfig& c) +{ return c.batch; } // ── 层注册表 ────────────────────────────────────────────────────────────── const std::vector& layer_registry() @@ -297,25 +307,38 @@ void setup_matmul_at(ComputeEngine& e, const BenchConfig& c, OpCtx& ctx) void run_matmul(ComputeEngine& e, const BenchConfig&, OpCtx& ctx) { // 输出精度跟随标记：--f16 时不额外付一次边界 cast（同 BitLinear 的 f16 路径） - auto r = e.matmul(ctx.a, ctx.b, false, false, - ctx.a.precision() == nn::Precision::F16 ? nn::Precision::F16 - : nn::Precision::F32); - if (!r) { std::printf("matmul error: %s\n", r.error().message.c_str()); std::abort(); } + auto r = + e.matmul(ctx.a, ctx.b, false, false, + ctx.a.precision() == nn::Precision::F16 ? nn::Precision::F16 : nn::Precision::F32); + if (!r) + { + std::printf("matmul error: %s\n", r.error().message.c_str()); + std::abort(); + } ctx.c = std::move(*r); } void run_matmul_bt(ComputeEngine& e, const BenchConfig&, OpCtx& ctx) { auto r = e.matmul(ctx.a, ctx.b, false, true); - if (!r) { std::printf("matmul_bt error: %s\n", r.error().message.c_str()); std::abort(); } + if (!r) + { + std::printf("matmul_bt error: %s\n", r.error().message.c_str()); + std::abort(); + } ctx.c = std::move(*r); } void run_matmul_at(ComputeEngine& e, const BenchConfig&, OpCtx& ctx) { auto r = e.matmul(ctx.a, ctx.b, true, false); - if (!r) { std::printf("matmul_at error: %s\n", r.error().message.c_str()); std::abort(); } + if (!r) + { + std::printf("matmul_at error: %s\n", r.error().message.c_str()); + std::abort(); + } ctx.c = std::move(*r); } -double work_matmul(const BenchConfig& c) { return 2.0 * c.m * c.n * c.k; } +double work_matmul(const BenchConfig& c) +{ return 2.0 * c.m * c.n * c.k; } void setup_batched(ComputeEngine& e, const BenchConfig& c, OpCtx& ctx) { @@ -325,10 +348,15 @@ void setup_batched(ComputeEngine& e, const BenchConfig& c, OpCtx& ctx) void run_batched(ComputeEngine& e, const BenchConfig& c, OpCtx& ctx) { auto r = e.batched_matmul(ctx.a, ctx.b, c.batch); - if (!r) { std::printf("batched_matmul error: %s\n", r.error().message.c_str()); std::abort(); } + if (!r) + { + std::printf("batched_matmul error: %s\n", r.error().message.c_str()); + std::abort(); + } ctx.c = std::move(*r); } -double work_batched(const BenchConfig& c) { return 2.0 * c.batch * c.m * c.n * c.k; } +double work_batched(const BenchConfig& c) +{ return 2.0 * c.batch * c.m * c.n * c.k; } void setup_add(ComputeEngine& e, const BenchConfig& c, OpCtx& ctx) { @@ -336,41 +364,34 @@ void setup_add(ComputeEngine& e, const BenchConfig& c, OpCtx& ctx) ctx.b = make_input(e, c.m, c.n); } void run_add(ComputeEngine& e, const BenchConfig&, OpCtx& ctx) -{ - (void)e.add_inplace(ctx.a, ctx.b); -} -double bytes_add(const BenchConfig& c) { return 3.0 * c.m * c.n * sizeof(Scalar); } +{ (void)e.add_inplace(ctx.a, ctx.b); } +double bytes_add(const BenchConfig& c) +{ return 3.0 * c.m * c.n * sizeof(Scalar); } void setup_exp(ComputeEngine& e, const BenchConfig& c, OpCtx& ctx) -{ - ctx.a = make_input(e, c.m, c.n); -} -double bytes_binary(const BenchConfig& c) { return 2.0 * c.m * c.n * sizeof(Scalar); } +{ ctx.a = make_input(e, c.m, c.n); } +double bytes_binary(const BenchConfig& c) +{ return 2.0 * c.m * c.n * sizeof(Scalar); } void run_row_reduce(ComputeEngine& e, const BenchConfig&, OpCtx& ctx) -{ - ctx.c = *e.row_reduce_sum(ctx.a); -} +{ ctx.c = *e.row_reduce_sum(ctx.a); } void run_col_reduce(ComputeEngine& e, const BenchConfig&, OpCtx& ctx) -{ - ctx.c = *e.col_reduce_sum(ctx.a); -} +{ ctx.c = *e.col_reduce_sum(ctx.a); } // 归约实际搬运 = 读全矩阵 (m*n) + 写输出（row_reduce 输出 (m,1) → m 个； // col_reduce 输出 (1,n) → n 个）。**不能记 2*m*n**——输出只有向量级，按全量 // 记会把 GB/s 虚高 ~2×（5244² 下虚报 440 GB/s，超过 448 理论峰值，物理不可能）。 -double bytes_row_reduce(const BenchConfig& c) { return (c.m * c.n + c.m) * sizeof(Scalar); } -double bytes_col_reduce(const BenchConfig& c) { return (c.m * c.n + c.n) * sizeof(Scalar); } +double bytes_row_reduce(const BenchConfig& c) +{ return (c.m * c.n + c.m) * sizeof(Scalar); } +double bytes_col_reduce(const BenchConfig& c) +{ return (c.m * c.n + c.n) * sizeof(Scalar); } void run_transpose(ComputeEngine& e, const BenchConfig&, OpCtx& ctx) -{ - ctx.c = *e.transpose(ctx.a); -} -double bytes_transpose(const BenchConfig& c) { return 2.0 * c.m * c.n * sizeof(Scalar); } +{ ctx.c = *e.transpose(ctx.a); } +double bytes_transpose(const BenchConfig& c) +{ return 2.0 * c.m * c.n * sizeof(Scalar); } void run_scale(ComputeEngine& e, const BenchConfig&, OpCtx& ctx) -{ - (void)e.scale_inplace(ctx.a, Scalar{0.5f}); -} +{ (void)e.scale_inplace(ctx.a, Scalar{0.5f}); } // ── RLA / RAPT 扫描原语（P4-2b）：状态扫描的顺序依赖是靠 barrier 串起来的， // kernel 天生是**延迟敏感**而非吞吐敏感型 —— 所以口径记 ms（附 GFLOPS 仅作 @@ -387,15 +408,16 @@ Tensor make_zeros(ComputeEngine& e, std::size_t rows, std::size_t cols) return std::move(*t); } -Tensor make_scan_input(ComputeEngine& e, const BenchConfig& c, - std::size_t rows, std::size_t cols) -{ - return c.f16 ? make_input_f16(e, rows, cols) : make_input(e, rows, cols); -} +Tensor make_scan_input(ComputeEngine& e, const BenchConfig& c, std::size_t rows, std::size_t cols) +{ return c.f16 ? make_input_f16(e, rows, cols) : make_input(e, rows, cols); } void setup_scan_prefix(ComputeEngine& e, const BenchConfig& c, OpCtx& ctx) { - if (c.heads == 0 || c.dk == 0) { std::printf("scan_prefix: heads/dk 不能为 0\n"); std::abort(); } + if (c.heads == 0 || c.dk == 0) + { + std::printf("scan_prefix: heads/dk 不能为 0\n"); + std::abort(); + } const std::size_t rows = c.batch * c.heads * c.dk; ctx.a = make_scan_input(e, c, rows, c.seq); // K ctx.b = make_scan_input(e, c, rows, c.seq); // V @@ -407,11 +429,15 @@ void setup_scan_prefix(ComputeEngine& e, const BenchConfig& c, OpCtx& ctx) } void run_scan_prefix(ComputeEngine& e, const BenchConfig& c, OpCtx& ctx) { - auto r = e.scan_prefix_outer(ctx.a, ctx.b, ctx.extra[0], ctx.extra[1], - ctx.extra[2], ctx.extra[3], /*has_state=*/true, - c.dk, c.heads, /*causal=*/true, ctx.aux, /*has_bnd=*/true, - c.f16 ? nn::Precision::F16 : nn::Precision::F32); - if (!r) { std::printf("scan_prefix error: %s\n", r.error().message.c_str()); std::abort(); } + auto r = + e.scan_prefix_outer(ctx.a, ctx.b, ctx.extra[0], ctx.extra[1], ctx.extra[2], ctx.extra[3], + /*has_state=*/true, c.dk, c.heads, /*causal=*/true, ctx.aux, + /*has_bnd=*/true, c.f16 ? nn::Precision::F16 : nn::Precision::F32); + if (!r) + { + std::printf("scan_prefix error: %s\n", r.error().message.c_str()); + std::abort(); + } ctx.c = std::move(*r); } // 每 (t, 头) 的浮点工作量：[0)/[1)/[2) 三个 dk×dk 的 matvec = 3·2dk²，[3)/[4) @@ -419,13 +445,16 @@ void run_scan_prefix(ComputeEngine& e, const BenchConfig& c, OpCtx& ctx) double flops_scan_prefix(const BenchConfig& c) { return 6.0 * static_cast(c.dk) * static_cast(c.dk) * - static_cast(c.heads) * static_cast(c.batch) * - static_cast(c.seq); + static_cast(c.heads) * static_cast(c.batch) * static_cast(c.seq); } void setup_scan_suffix(ComputeEngine& e, const BenchConfig& c, OpCtx& ctx) { - if (c.heads == 0 || c.dk == 0) { std::printf("scan_suffix: heads/dk 不能为 0\n"); std::abort(); } + if (c.heads == 0 || c.dk == 0) + { + std::printf("scan_suffix: heads/dk 不能为 0\n"); + std::abort(); + } const std::size_t rows = c.batch * c.heads * c.dk; ctx.extra.push_back(make_scan_input(e, c, rows * c.dk, c.seq)); // D (B*H*dk², seq) ctx.a = make_scan_input(e, c, rows, c.seq); // X @@ -434,13 +463,18 @@ void setup_scan_suffix(ComputeEngine& e, const BenchConfig& c, OpCtx& ctx) } void run_scan_suffix(ComputeEngine& e, const BenchConfig& c, OpCtx& ctx) { - auto r = e.scan_suffix_outer(ctx.extra[0], ctx.a, ctx.b, - c.dk, c.heads, /*causal=*/true, ctx.aux, /*has_bnd=*/true, - c.f16 ? nn::Precision::F16 : nn::Precision::F32); - if (!r) { std::printf("scan_suffix error: %s\n", r.error().message.c_str()); std::abort(); } + auto r = + e.scan_suffix_outer(ctx.extra[0], ctx.a, ctx.b, c.dk, c.heads, /*causal=*/true, ctx.aux, + /*has_bnd=*/true, c.f16 ? nn::Precision::F16 : nn::Precision::F32); + if (!r) + { + std::printf("scan_suffix error: %s\n", r.error().message.c_str()); + std::abort(); + } ctx.c = std::move(*r); } -double flops_scan_suffix(const BenchConfig& c) { return flops_scan_prefix(c); } +double flops_scan_suffix(const BenchConfig& c) +{ return flops_scan_prefix(c); } // ── 三值打包权重 GEMM（P2 / F3）与稠密 GEMM 的配对 A/B ─────────────────── // 口径（docs 21 §5 的 P2 准入门槛）：同一 (m,k)·(k,n) 形状、同一份随机数造的 @@ -454,8 +488,7 @@ double flops_scan_suffix(const BenchConfig& c) { return flops_scan_prefix(c); } // --f16 时打包侧跑 f16 存储变体（B/C 均 half）= BitLinear 的 compute=f16 路径。 // 造 latent (rows,len) → τ 逐行 absmean → 引擎 trit_pack（即被测路径本身） -void make_packed(ComputeEngine& e, std::size_t rows, std::size_t len, - Tensor& packed, Tensor& tau) +void make_packed(ComputeEngine& e, std::size_t rows, std::size_t len, Tensor& packed, Tensor& tau) { if (!e.describe().ternary_gemm) { @@ -466,7 +499,8 @@ void make_packed(ComputeEngine& e, std::size_t rows, std::size_t len, std::mt19937 g(4321u); std::uniform_real_distribution d(-1.0f, 1.0f); auto s = w.span(); - for (std::size_t i = 0; i < s.size(); ++i) s[i] = d(g); + for (std::size_t i = 0; i < s.size(); ++i) + s[i] = d(g); Matrix tm(rows, 1); for (std::size_t r = 0; r < rows; ++r) { @@ -477,7 +511,11 @@ void make_packed(ComputeEngine& e, std::size_t rows, std::size_t len, } auto wt = e.from_matrix(std::move(w)); auto tt = e.from_matrix(std::move(tm)); - if (!wt || !tt) { std::printf("matmul_q setup: 上传失败\n"); std::abort(); } + if (!wt || !tt) + { + std::printf("matmul_q setup: 上传失败\n"); + std::abort(); + } auto pk = e.trit_pack(*wt, *tt); if (!pk) { @@ -495,7 +533,8 @@ Tensor make_input_f16(ComputeEngine& e, std::size_t rows, std::size_t cols) std::mt19937 g(1234u); std::uniform_real_distribution d(-1.0f, 1.0f); auto s = m.span(); - for (std::size_t i = 0; i < s.size(); ++i) s[i] = d(g); + for (std::size_t i = 0; i < s.size(); ++i) + s[i] = d(g); auto t = e.from_matrix(std::move(m), nn::Precision::F16); if (!t) std::abort(); return std::move(*t); @@ -509,7 +548,11 @@ void setup_matmul_q(ComputeEngine& e, const BenchConfig& c, OpCtx& ctx) void run_matmul_q(ComputeEngine& e, const BenchConfig& c, OpCtx& ctx) { auto r = e.matmul_q(ctx.a, ctx.aux, ctx.b, c.f16 ? nn::Precision::F16 : nn::Precision::F32); - if (!r) { std::printf("matmul_q error: %s\n", r.error().message.c_str()); std::abort(); } + if (!r) + { + std::printf("matmul_q error: %s\n", r.error().message.c_str()); + std::abort(); + } ctx.c = std::move(*r); } // A 存 (K,M)：τ 按**存储行** k 取（转发形态，反向 dX 走这条） @@ -520,14 +563,20 @@ void setup_matmul_q_t(ComputeEngine& e, const BenchConfig& c, OpCtx& ctx) } void run_matmul_q_t(ComputeEngine& e, const BenchConfig& c, OpCtx& ctx) { - auto r = e.matmul_q_t(ctx.a, ctx.aux, ctx.b, c.m, - c.f16 ? nn::Precision::F16 : nn::Precision::F32); - if (!r) { std::printf("matmul_q_t error: %s\n", r.error().message.c_str()); std::abort(); } + auto r = + e.matmul_q_t(ctx.a, ctx.aux, ctx.b, c.m, c.f16 ? nn::Precision::F16 : nn::Precision::F32); + if (!r) + { + std::printf("matmul_q_t error: %s\n", r.error().message.c_str()); + std::abort(); + } ctx.c = std::move(*r); } // 权重带宽口径：打包 A 的实际字节（m*k/4 B）+ τ（4 B/存储行） -double bytes_packed(const BenchConfig& c) { return c.m * c.k / 4.0 + 4.0 * c.m; } -double bytes_packed_t(const BenchConfig& c) { return c.m * c.k / 4.0 + 4.0 * c.k; } +double bytes_packed(const BenchConfig& c) +{ return c.m * c.k / 4.0 + 4.0 * c.m; } +double bytes_packed_t(const BenchConfig& c) +{ return c.m * c.k / 4.0 + 4.0 * c.k; } const std::vector& op_registry() { @@ -581,7 +630,11 @@ void bench_op(ComputeEngine& engine, const BenchConfig& cfg, const OpSpec& spec) if (cfg.hash) { auto out_m = engine.to_matrix(ctx.c); - if (!out_m) { std::printf(" 哈希：to_matrix 失败\n"); std::abort(); } + if (!out_m) + { + std::printf(" 哈希：to_matrix 失败\n"); + std::abort(); + } std::uint64_t h = 1469598103934665603ull; for (const auto& v : out_m->span()) { @@ -630,8 +683,7 @@ void bench_layer(ComputeEngine& engine, const BenchConfig& cfg, const LayerSpec& for (int i = 0; i < cfg.warmup; ++i) { (void)model.forward(x); - if (!cfg.forward_only) - (void)model.backward(grad); + if (!cfg.forward_only) (void)model.backward(grad); } // 计时 @@ -683,7 +735,8 @@ void print_help(const char* prog) help.opt("--f16", "打包算子 (matmul_q*) 跑 f16 存储变体 (B/C 均 half)"); help.opt("--iter ", "迭代次数 (默认: 10)"); help.opt("--warmup ", "预热次数 (默认: 2)"); - help.opt("--inner ", "每次计时内重复执行次数 (默认: 1；小 n 的 GEMM 用它摊薄\nsubmit/wait 固定开销 —— 否则测到的是启动开销不是内核)"); + help.opt("--inner ", "每次计时内重复执行次数 (默认: 1；小 n 的 GEMM 用它摊薄\nsubmit/wait " + "固定开销 —— 否则测到的是启动开销不是内核)"); help.opt("--in/--out ", "Linear 输入/输出维度"); help.opt("--dmodel/--heads/--dff ", "序列层维度"); help.opt("--batch/--seq ", "batch / 序列长度"); @@ -730,14 +783,27 @@ int main(int argc, char* argv[]) std::string cur; for (char ch : val) { - if (ch == ',') { if (!cur.empty()) cfg.layers.push_back(cur); cur.clear(); } + if (ch == ',') + { + if (!cur.empty()) cfg.layers.push_back(cur); + cur.clear(); + } else cur.push_back(ch); } if (!cur.empty()) cfg.layers.push_back(cur); } - else if (a == "--iter" && i + 1 < argc) { cfg.iter = static_cast(parse_size(argv[++i])); } - else if (a == "--warmup" && i + 1 < argc) { cfg.warmup = static_cast(parse_size(argv[++i])); } - else if (a == "--inner" && i + 1 < argc) { cfg.inner = static_cast(parse_size(argv[++i])); } + else if (a == "--iter" && i + 1 < argc) + { + cfg.iter = static_cast(parse_size(argv[++i])); + } + else if (a == "--warmup" && i + 1 < argc) + { + cfg.warmup = static_cast(parse_size(argv[++i])); + } + else if (a == "--inner" && i + 1 < argc) + { + cfg.inner = static_cast(parse_size(argv[++i])); + } else if (a == "--in") { cfg.in = parse_size(argv[++i]); } else if (a == "--out") { cfg.out = parse_size(argv[++i]); } else if (a == "--dmodel") { cfg.dmodel = parse_size(argv[++i]); } @@ -757,7 +823,11 @@ int main(int argc, char* argv[]) std::string cur; for (char ch : val) { - if (ch == ',') { if (!cur.empty()) cfg.ops.push_back(cur); cur.clear(); } + if (ch == ',') + { + if (!cur.empty()) cfg.ops.push_back(cur); + cur.clear(); + } else cur.push_back(ch); } if (!cur.empty()) cfg.ops.push_back(cur); @@ -767,8 +837,16 @@ int main(int argc, char* argv[]) else if (a == "--k") { cfg.k = parse_size(argv[++i]); } else if (a == "--dk") { cfg.dk = parse_size(argv[++i]); } else if (a == "--hash") { cfg.hash = true; } - else if (a == "--help" || a == "-h") { print_help(argv[0]); return 0; } - else { std::printf("未知参数: %s（--help 查看用法）\n", a.c_str()); return 2; } + else if (a == "--help" || a == "-h") + { + print_help(argv[0]); + return 0; + } + else + { + std::printf("未知参数: %s（--help 查看用法）\n", a.c_str()); + return 2; + } } // 选择引擎（CPU/GPU），失败硬报错不回退 @@ -777,8 +855,7 @@ int main(int argc, char* argv[]) NN_EXIT(engine_r, 2, "引擎创建失败: %s\n"); ComputeEngine& engine = **engine_r; const bool op_mode = !cfg.ops.empty(); - std::printf("=== layer_bench (%s%s) ===\n", - cfg.gpu ? "GPU" : "CPU", + std::printf("=== layer_bench (%s%s) ===\n", cfg.gpu ? "GPU" : "CPU", op_mode ? ", ops" : (cfg.forward_only ? ", forward-only" : ", fwd+bwd")); // 算子模式：--op @@ -790,13 +867,19 @@ int main(int argc, char* argv[]) { if (want == "all") { - for (const auto& s : oreg) to_run.push_back(&s); + for (const auto& s : oreg) + to_run.push_back(&s); continue; } bool found = false; for (const auto& s : oreg) { - if (want == s.name) { to_run.push_back(&s); found = true; break; } + if (want == s.name) + { + to_run.push_back(&s); + found = true; + break; + } } if (!found) std::printf("未知算子: %s\n", want.c_str()); } @@ -815,7 +898,8 @@ int main(int argc, char* argv[]) const auto& reg = layer_registry(); if (cfg.all) { - for (const auto& s : reg) to_run.push_back(&s); + for (const auto& s : reg) + to_run.push_back(&s); } else if (!cfg.layers.empty()) { @@ -824,7 +908,12 @@ int main(int argc, char* argv[]) bool found = false; for (const auto& s : reg) { - if (want == s.name) { to_run.push_back(&s); found = true; break; } + if (want == s.name) + { + to_run.push_back(&s); + found = true; + break; + } } if (!found) std::printf("未知层: %s\n", want.c_str()); } diff --git a/src/layer_broadcast_test.cpp b/src/layer_broadcast_test.cpp index 79669993..96f16ffb 100644 --- a/src/layer_broadcast_test.cpp +++ b/src/layer_broadcast_test.cpp @@ -14,33 +14,26 @@ #include int g_fail = 0; -#define CHECK(cond, msg) \ - do { \ - if (!(cond)) { \ - std::cout << "[FAIL] " << __FILE__ << ":" << __LINE__ << ": " \ - << msg << "\n"; \ - ++g_fail; \ - } \ +#define CHECK(cond, msg) \ + do \ + { \ + if (!(cond)) \ + { \ + std::cout << "[FAIL] " << __FILE__ << ":" << __LINE__ << ": " << msg << "\n"; \ + ++g_fail; \ + } \ } while (0) -namespace { +namespace +{ // 记录型叶子层：捕获三个广播开关的到达情况 class RecordingLayer final : public nn::Layer { public: - [[nodiscard]] nn::Result forward(const nn::Tensor& x) override - { - return x; - } - [[nodiscard]] nn::Result backward(const nn::Tensor& g) override - { - return g; - } - [[nodiscard]] const char* layer_name() const noexcept override - { - return "RecordingLayer"; - } + [[nodiscard]] nn::Result forward(const nn::Tensor& x) override { return x; } + [[nodiscard]] nn::Result backward(const nn::Tensor& g) override { return g; } + [[nodiscard]] const char* layer_name() const noexcept override { return "RecordingLayer"; } void set_training(bool t) override { @@ -63,38 +56,21 @@ class CompositeOnlyChildren final : public nn::Layer public: RecordingLayer a, b; - [[nodiscard]] nn::Result forward(const nn::Tensor& x) override - { - return x; - } - [[nodiscard]] nn::Result backward(const nn::Tensor& g) override - { - return g; - } + [[nodiscard]] nn::Result forward(const nn::Tensor& x) override { return x; } + [[nodiscard]] nn::Result backward(const nn::Tensor& g) override { return g; } [[nodiscard]] const char* layer_name() const noexcept override - { - return "CompositeOnlyChildren"; - } + { return "CompositeOnlyChildren"; } protected: - [[nodiscard]] std::vector param_children() noexcept override - { - return {&a, &b}; - } + [[nodiscard]] std::vector param_children() noexcept override { return {&a, &b}; } }; // 叶子复合层（无子层）：默认实现应为安全 no-op class NoChildrenLayer final : public nn::Layer { public: - [[nodiscard]] nn::Result forward(const nn::Tensor& x) override - { - return x; - } - [[nodiscard]] nn::Result backward(const nn::Tensor& g) override - { - return g; - } + [[nodiscard]] nn::Result forward(const nn::Tensor& x) override { return x; } + [[nodiscard]] nn::Result backward(const nn::Tensor& g) override { return g; } }; } // namespace @@ -117,8 +93,7 @@ int main() // clear_cache 广播 c.clear_cache(); - CHECK(c.a.clear_cache_calls == 1 && c.b.clear_cache_calls == 1, - "clear_cache 到达全部子层"); + CHECK(c.a.clear_cache_calls == 1 && c.b.clear_cache_calls == 1, "clear_cache 到达全部子层"); // set_precision_profile 广播（未 init 的子层接收不 fail-fast） nn::PrecisionProfile prof; @@ -131,8 +106,7 @@ int main() CHECK(c.b.precision_profile().param == nn::Precision::F16, "set_precision_profile 到达子层 b"); // 复合层自身也记录（基类语义不变） - CHECK(c.precision_profile().param == nn::Precision::F16, - "复合层自身 profile 已记录"); + CHECK(c.precision_profile().param == nn::Precision::F16, "复合层自身 profile 已记录"); } // [3] 无子层的叶子：默认实现安全 no-op diff --git a/src/layer_lifetime_test.cpp b/src/layer_lifetime_test.cpp index bbddf0d8..29b34c77 100644 --- a/src/layer_lifetime_test.cpp +++ b/src/layer_lifetime_test.cpp @@ -32,13 +32,9 @@ class FlakyLayer : public nn::Layer { public: [[nodiscard]] nn::Result forward(const nn::Tensor&) override - { - NN_FAIL("FlakyLayer::forward: 未实现"); - } + { NN_FAIL("FlakyLayer::forward: 未实现"); } [[nodiscard]] nn::Result backward(const nn::Tensor&) override - { - NN_FAIL("FlakyLayer::backward: 未实现"); - } + { NN_FAIL("FlakyLayer::backward: 未实现"); } [[nodiscard]] const char* layer_name() const noexcept override { return "FlakyLayer"; } int attempts = 0; @@ -47,16 +43,14 @@ class FlakyLayer : public nn::Layer [[nodiscard]] nn::Result init_impl(nn::ComputeEngine&) override { ++attempts; - if (attempts == 1) - NN_FAIL("first init fails（测试注入）"); + if (attempts == 1) NN_FAIL("first init fails（测试注入）"); return {}; } }; void test_domain_model() { - static_assert(std::is_same_v, - "DomainId = 不透明整数域标识"); + static_assert(std::is_same_v, "DomainId = 不透明整数域标识"); static_assert(nn::kHostDomain != 0, "宿主域非零（0 保留给空张量）"); static_assert(!std::is_default_constructible_v, "CpuEngine 不允许栈上构造（必须用 create()）"); @@ -150,10 +144,10 @@ void test_cross_engine_reinit_rejected() void test_optimizer_keeps_engine_alive() { auto engine = nn::CpuEngine::create(); - nn::Tensor p = engine->create_tensor( - 2, 2, nn::Precision::F32, nn::InitSpec::constant(nn::Scalar{1})); - nn::Tensor g = engine->create_tensor( - 2, 2, nn::Precision::F32, nn::InitSpec::constant(nn::Scalar{0.5})); + nn::Tensor p = + engine->create_tensor(2, 2, nn::Precision::F32, nn::InitSpec::constant(nn::Scalar{1})); + nn::Tensor g = + engine->create_tensor(2, 2, nn::Precision::F32, nn::InitSpec::constant(nn::Scalar{0.5})); std::vector params{std::ref(p)}; std::vector grads{std::ref(g)}; auto opt = nn::create_optimizer("sgd", *engine, params, grads, 0.1f, 0.0f); @@ -250,24 +244,20 @@ void test_layer_init_failure_rollback() auto engine = nn::CpuEngine::create(); FlakyLayer layer; - CHECK(layer.binding_state() == nn::Layer::BindingState::unbound, - "新层状态 = unbound"); + CHECK(layer.binding_state() == nn::Layer::BindingState::unbound, "新层状态 = unbound"); CHECK(!layer.engine_alive(), "新层未持有引擎"); auto r1 = layer.init(*engine); CHECK(!r1, "首次 init 按注入失败"); - CHECK(layer.binding_state() == nn::Layer::BindingState::failed, - "init 失败后状态 = failed"); + CHECK(layer.binding_state() == nn::Layer::BindingState::failed, "init 失败后状态 = failed"); CHECK(layer.engine_alive(), "失败后仍保留可诊断的引擎强引用"); layer.set_precision_profile(nn::profile_f16()); - CHECK(layer.precision_profile().compute == nn::Precision::F16, - "failed 状态允许修正 profile"); + CHECK(layer.precision_profile().compute == nn::Precision::F16, "failed 状态允许修正 profile"); auto r2 = layer.init(*engine); CHECK(r2, "失败后可重新 init"); - CHECK(layer.binding_state() == nn::Layer::BindingState::ready, - "重新 init 成功后状态 = ready"); + CHECK(layer.binding_state() == nn::Layer::BindingState::ready, "重新 init 成功后状态 = ready"); } } // namespace @@ -285,9 +275,7 @@ int main() test_cross_domain_rejected(); test_layer_init_failure_rollback(); - if (g_fail == 0) - std::printf("layer_lifetime_test: ALL PASS\n"); - else - std::printf("layer_lifetime_test: %d failure(s)\n", g_fail); + if (g_fail == 0) std::printf("layer_lifetime_test: ALL PASS\n"); + else std::printf("layer_lifetime_test: %d failure(s)\n", g_fail); return g_fail == 0 ? 0 : 1; } diff --git a/src/lifetime_contract_test.cpp b/src/lifetime_contract_test.cpp index 282409f2..3249d9b0 100644 --- a/src/lifetime_contract_test.cpp +++ b/src/lifetime_contract_test.cpp @@ -64,8 +64,7 @@ void test_offload_state_machine() auto r2 = off.export_activations(*engine, {a, b}); CHECK(!r2, "未 import 的重复 export 必须 fail-fast"); if (!r2) - CHECK(r2.error().message.find("尚未 import") != std::string::npos, - "错误信息点明未 import"); + CHECK(r2.error().message.find("尚未 import") != std::string::npos, "错误信息点明未 import"); // forward 入口的作废通道：discard 后状态复位，可重新导出 off.discard(); @@ -98,9 +97,7 @@ void test_offload_rejects_cpu_engine() off.set_enabled(true); auto r = off.export_activations(*engine, {a}); CHECK(!r, "CPU 引擎 export 必须 fail-fast"); - if (!r) - CHECK(r.error().message.find("仅支持") != std::string::npos, - "错误信息说明仅支持 GPU"); + if (!r) CHECK(r.error().message.find("仅支持") != std::string::npos, "错误信息说明仅支持 GPU"); CHECK(!off.offloaded(), "失败后状态保持 Idle"); CHECK(a.valid(), "失败路径不得改动激活张量"); } @@ -115,9 +112,7 @@ void test_offload_rejects_non_f32() off.set_enabled(true); auto r = off.export_activations(*engine, {h}); CHECK(!r, "非 F32 激活导出必须 fail-fast"); - if (!r) - CHECK(r.error().message.find("f16") != std::string::npos, - "错误信息带出实际精度"); + if (!r) CHECK(r.error().message.find("f16") != std::string::npos, "错误信息带出实际精度"); CHECK(!off.offloaded(), "失败后状态保持 Idle"); CHECK(h.valid(), "失败路径不得改动激活张量"); } @@ -153,9 +148,7 @@ int main() test_offload_rejects_non_f32(); test_optimizer_before_model_order(); - if (g_fail == 0) - std::printf("lifetime_contract_test: ALL PASS\n"); - else - std::printf("lifetime_contract_test: %d failure(s)\n", g_fail); + if (g_fail == 0) std::printf("lifetime_contract_test: ALL PASS\n"); + else std::printf("lifetime_contract_test: %d failure(s)\n", g_fail); return g_fail == 0 ? 0 : 1; } diff --git a/src/lr_scheduler_test.cpp b/src/lr_scheduler_test.cpp index 70286c34..646a4a03 100644 --- a/src/lr_scheduler_test.cpp +++ b/src/lr_scheduler_test.cpp @@ -27,10 +27,10 @@ namespace { int g_fail = 0; -using nn::Scalar; using nn::compute_epoch_lr; using nn::compute_step_lr; using nn::LrScheduleConfig; +using nn::Scalar; using nn::StepLrScheduleConfig; // 相对容差比较（float 调度曲线） diff --git a/src/maxpool_gradcheck.cpp b/src/maxpool_gradcheck.cpp index 9374ba3a..acd5ae76 100644 --- a/src/maxpool_gradcheck.cpp +++ b/src/maxpool_gradcheck.cpp @@ -24,13 +24,14 @@ #include #include -using nn::Scalar; -using nn::Matrix; -using nn::Tensor; using nn::ComputeEngine; +using nn::Matrix; using nn::MaxPool2D; +using nn::Scalar; +using nn::Tensor; -namespace { +namespace +{ // 独立参考：窗口 max + argmax 散射（不依赖层内任何实现）。 // 平局取"先出现者"，与层的 `if (v > best)` 判据一致。 @@ -40,9 +41,7 @@ struct PoolRef Matrix x; // (C*H*W, B) Scalar xv(std::size_t c, std::size_t ih, std::size_t iw, std::size_t b) const - { - return x.at_unchecked(c * H * W + ih * W + iw, b); - } + { return x.at_unchecked(c * H * W + ih * W + iw, b); } Matrix fwd() const { Matrix out(C * OH * OW, B); @@ -79,8 +78,7 @@ struct PoolRef if (xv(c, oh * S + dh, ow * S + dw, b) == best) ++cnt; if (cnt == 0) continue; const std::size_t orr = c * OH * OW + oh * OW + ow; - const Scalar share = - go.at_unchecked(orr, b) / static_cast(cnt); + const Scalar share = go.at_unchecked(orr, b) / static_cast(cnt); for (std::size_t dh = 0; dh < P; ++dh) for (std::size_t dw = 0; dw < P; ++dw) { @@ -129,7 +127,8 @@ int main(int argc, char* argv[]) std::cout << " MaxPool2D 正确性检查（独立参考实现比对）\n"; std::cout << "========================================\n"; - struct Case { std::size_t C, H, W, P, S, B; }; + struct Case + { std::size_t C, H, W, P, S, B; }; const Case cases[] = { {2, 6, 6, 2, 2, 2}, // 标准（stride == pool） {2, 6, 6, 2, 1, 2}, // 重叠窗口（梯度在同一位置累加） @@ -151,22 +150,29 @@ int main(int argc, char* argv[]) { const std::size_t OH = (cs.H - cs.P) / cs.S + 1; const std::size_t OW = (cs.W - cs.P) / cs.S + 1; - std::printf("── C=%zu H=%zu W=%zu pool=%zu stride=%zu B=%zu (out %zux%zu) ──\n", - cs.C, cs.H, cs.W, cs.P, cs.S, cs.B, OH, OW); + std::printf("── C=%zu H=%zu W=%zu pool=%zu stride=%zu B=%zu (out %zux%zu) ──\n", cs.C, cs.H, + cs.W, cs.P, cs.S, cs.B, OH, OW); MaxPool2D pool(cs.C, cs.H, cs.W, cs.P, cs.S); (void)pool.init(eng); // M6 段 C：层由 init 绑定引擎 Matrix xm(cs.C * cs.H * cs.W, cs.B); - for (std::size_t i = 0; i < xm.size(); ++i) xm.span()[i] = dist(rng); + for (std::size_t i = 0; i < xm.size(); ++i) + xm.span()[i] = dist(rng); Matrix gom(cs.C * OH * OW, cs.B); - for (std::size_t i = 0; i < gom.size(); ++i) gom.span()[i] = dist(rng); + for (std::size_t i = 0; i < gom.size(); ++i) + gom.span()[i] = dist(rng); PoolRef ref{cs.C, cs.H, cs.W, cs.P, cs.S, OH, OW, cs.B, xm}; auto x = eng.from_matrix(xm); auto go = eng.from_matrix(gom); - if (!x || !go) { std::cerr << " 上传失败\n"; return 1; } + if (!x || !go) + { + std::cerr << " 上传失败\n"; + return 1; + } - const auto report = [&](const char* what, const Matrix& got, const Matrix& want) { + const auto report = [&](const char* what, const Matrix& got, const Matrix& want) + { const Scalar e = max_abs_diff(got, want); const bool ok = e < tol; std::printf(" %-22s %s max_abs_diff=%.3e\n", what, ok ? "OK" : "FAIL", @@ -198,8 +204,7 @@ int main(int argc, char* argv[]) auto out3 = pool.forward(*x); NN_EXIT(out3, 1, " ckpt forward 失败\n"); auto gx3 = pool.backward(*go); - std::printf(" %-22s %s\n", "ckpt→backward", - gx3 ? "FAIL（未拦截）" : "OK（已拒绝）"); + std::printf(" %-22s %s\n", "ckpt→backward", gx3 ? "FAIL（未拦截）" : "OK（已拒绝）"); if (gx3) ++failures; auto rec = pool.forward_recompute(*x); @@ -228,12 +233,19 @@ int main(int argc, char* argv[]) MaxPool2D pool(C, H, W, P, S); (void)pool.init(eng); // M6 段 C：层由 init 绑定引擎 Matrix xm(C * H * W, B); - xm.span()[0] = 1; xm.span()[1] = 1; xm.span()[2] = 1; xm.span()[3] = 0; + xm.span()[0] = 1; + xm.span()[1] = 1; + xm.span()[2] = 1; + xm.span()[3] = 0; Matrix gom(C * 1 * 1, B); gom.span()[0] = 2; auto x = eng.from_matrix(xm); auto go = eng.from_matrix(gom); - if (!x || !go) { std::cerr << " 上传失败\n"; return 1; } + if (!x || !go) + { + std::cerr << " 上传失败\n"; + return 1; + } auto out = pool.forward(*x); if (!out) @@ -260,10 +272,10 @@ int main(int argc, char* argv[]) std::fabs(gm->at_unchecked(2, 0) - want) < 1e-6f && std::fabs(gm->at_unchecked(3, 0)) < 1e-6f; Scalar sum = 0; - for (std::size_t i = 0; i < gm->size(); ++i) sum += gm->span()[i]; + for (std::size_t i = 0; i < gm->size(); ++i) + sum += gm->span()[i]; const bool conserved = std::fabs(sum - Scalar{2}) < 1e-6f; - std::printf(" %-22s %s（期望 [2/3,2/3,2/3,0]）\n", "并列均分", - ok ? "OK" : "FAIL"); + std::printf(" %-22s %s（期望 [2/3,2/3,2/3,0]）\n", "并列均分", ok ? "OK" : "FAIL"); std::printf(" %-22s %s（Σgrad=%.6f，期望 2）\n", "总梯度守恒", conserved ? "OK" : "FAIL", double(sum)); if (!ok || !conserved) ++failures; diff --git a/src/mem_probe.cpp b/src/mem_probe.cpp index e154cef9..2b904dda 100644 --- a/src/mem_probe.cpp +++ b/src/mem_probe.cpp @@ -60,11 +60,10 @@ namespace [[nodiscard]] long long smi_used_mb() { #ifdef _WIN32 - FILE* f = _popen( - "nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>nul", "r"); + FILE* f = _popen("nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>nul", "r"); #else - FILE* f = popen( - "nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null", "r"); + FILE* f = + popen("nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null", "r"); #endif if (!f) return -1; long long best = -1; @@ -84,9 +83,7 @@ namespace } [[nodiscard]] double to_mb(std::size_t bytes) -{ - return static_cast(bytes) / (1024.0 * 1024.0); -} +{ return static_cast(bytes) / (1024.0 * 1024.0); } // ── 采样点 ──────────────────────────────────────────────────────────────── struct Sample @@ -120,7 +117,12 @@ void print_top_allocs(const char* name, const std::vector& sizes, { bool found = false; for (auto& a : agg) - if (a.first == sz) { ++a.second; found = true; break; } + if (a.first == sz) + { + ++a.second; + found = true; + break; + } if (!found) agg.push_back({sz, 1}); } std::sort(agg.begin(), agg.end(), @@ -134,14 +136,25 @@ void print_top_allocs(const char* name, const std::vector& sizes, std::size_t b1 = 0, b2 = 0, b3 = 0, c1 = 0, c2 = 0, c3 = 0; for (auto sz : sizes) { - if (sz >= 100ull * 1024 * 1024) { b1 += sz; ++c1; } - else if (sz >= 10ull * 1024 * 1024) { b2 += sz; ++c2; } - else { b3 += sz; ++c3; } + if (sz >= 100ull * 1024 * 1024) + { + b1 += sz; + ++c1; + } + else if (sz >= 10ull * 1024 * 1024) + { + b2 += sz; + ++c2; + } + else + { + b3 += sz; + ++c3; + } } - std::cout << "\n " << name << " 桶: ≥100MB " << c1 << "项/" - << (b1 / (1024.0 * 1024.0)) << "MB 10-100MB " << c2 << "项/" - << (b2 / (1024.0 * 1024.0)) << "MB <10MB " << c3 << "项/" - << (b3 / (1024.0 * 1024.0)) << "MB\n"; + std::cout << "\n " << name << " 桶: ≥100MB " << c1 << "项/" << (b1 / (1024.0 * 1024.0)) + << "MB 10-100MB " << c2 << "项/" << (b2 / (1024.0 * 1024.0)) << "MB <10MB " << c3 + << "项/" << (b3 / (1024.0 * 1024.0)) << "MB\n"; std::cout << std::setprecision(1); } } // namespace @@ -192,9 +205,9 @@ void take_sample(const std::string& label) { std::cout << " " << std::left << std::setw(9) << name << std::right << " committed=" << std::fixed << std::setprecision(1) - << to_mb(static_cast(ps.total_bytes)) - << "MB (dev " << to_mb(static_cast(ps.device_bytes)) - << " + host " << to_mb(static_cast(ps.host_bytes)) << ")" + << to_mb(static_cast(ps.total_bytes)) << "MB (dev " + << to_mb(static_cast(ps.device_bytes)) << " + host " + << to_mb(static_cast(ps.host_bytes)) << ")" << " live=" << to_mb(static_cast(ps.allocated_bytes)) << "MB free=" << to_mb(static_cast(ps.free_bytes)) << "MB frag=" << std::setprecision(2) << ps.fragmentation @@ -203,16 +216,14 @@ void take_sample(const std::string& label) { std::cout << " Δlive=" << std::showpos << std::setprecision(1) << (d_live / (1024.0 * 1024.0)) - << "MB Δcommit=" << (d_commit / (1024.0 * 1024.0)) - << "MB" << std::noshowpos; + << "MB Δcommit=" << (d_commit / (1024.0 * 1024.0)) << "MB" + << std::noshowpos; } std::cout << "\n"; // 池账本计数器（累计值；相邻采样点相减得该阶段的开销） std::cout << " calls=" << ps.c_alloc_calls - << " blk_new/free=" << ps.c_blocks_created << "/" - << ps.c_blocks_released - << " scans=" << ps.c_block_scans << "blk/" - << ps.c_region_scans << "reg" + << " blk_new/free=" << ps.c_blocks_created << "/" << ps.c_blocks_released + << " scans=" << ps.c_block_scans << "blk/" << ps.c_region_scans << "reg" << " vkalloc=" << std::setprecision(1) << ps.vkalloc_ms << "ms\n"; }; pool_line("persist", s.persist, d_live, d_commit, has_prev); @@ -220,13 +231,12 @@ void take_sample(const std::string& label) std::cout << std::resetiosflags(std::ios_base::floatfield); if (s.pending_bytes > 0) std::cout << std::fixed << std::setprecision(0) - << " 延迟销毁未归还(pending) = " - << (s.pending_bytes / (1024.0 * 1024.0)) << " MB\n"; + << " 延迟销毁未归还(pending) = " << (s.pending_bytes / (1024.0 * 1024.0)) + << " MB\n"; print_top_allocs("persist", s.p_allocs, 16); print_top_allocs("transient", s.t_allocs, 16); } - if (s.smi >= 0) - std::cout << " nvidia-smi used = " << s.smi << " MiB\n"; + if (s.smi >= 0) std::cout << " nvidia-smi used = " << s.smi << " MiB\n"; g_samples.push_back(std::move(s)); } @@ -241,13 +251,13 @@ void take_sample(const std::string& label) [[nodiscard]] std::size_t sum_tensor_bytes(std::vector& refs) { std::size_t total = 0; - for (auto& r : refs) total += tensor_bytes(r.get()); + for (auto& r : refs) + total += tensor_bytes(r.get()); return total; } // 解析命令行数字 -template -[[nodiscard]] T parse_num(const char* s, const char* opt) +template [[nodiscard]] T parse_num(const char* s, const char* opt) { auto v = nn::parse_number(s); if (!v) @@ -278,8 +288,7 @@ void print_usage(const char* prog) help.opt("--no-kv", "跳过 KV cache 推理阶段"); help.opt("--kv-steps ", "KV 阶段 forward_step 次数 (默认: 32)"); // RV-04：名单文本由 nn::optimizer_table()（唯一真源）生成（此前此处漏了 adamwquant） - help.opt("--optimizer ", - "优化器: " + nn::optimizer_names("/") + " (默认: adam)"); + help.opt("--optimizer ", "优化器: " + nn::optimizer_names("/") + " (默认: adam)"); help.opt("--help, -h", "显示此帮助信息"); help.section("模型尺寸"); @@ -317,7 +326,8 @@ int main(int argc, char* argv[]) for (int i = 1; i < argc; ++i) { const std::string arg = argv[i]; - const auto need = [&](const char* name) -> const char* { + const auto need = [&](const char* name) -> const char* + { if (i + 1 >= argc) { std::cerr << "--" << name << " 缺少值\n"; @@ -325,9 +335,17 @@ int main(int argc, char* argv[]) } return argv[++i]; }; - if (arg == "--help" || arg == "-h") { print_usage(argv[0]); return 0; } + if (arg == "--help" || arg == "-h") + { + print_usage(argv[0]); + return 0; + } else if (arg == "--f16") f16 = true; - else if (arg == "--f16-all") { f16 = true; f16_all = true; } + else if (arg == "--f16-all") + { + f16 = true; + f16_all = true; + } else if (arg == "--activation-offload") activation_offload = true; else if (arg == "--no-kv") do_kv = false; else if (arg == "--doc-mask") doc_mask = true; @@ -335,13 +353,20 @@ int main(int argc, char* argv[]) else if (arg == "--seq") seq = parse_num(need("seq"), "--seq"); else if (arg == "--vocab") vocab = parse_num(need("vocab"), "--vocab"); else if (arg == "--d-model") d_model = parse_num(need("d-model"), "--d-model"); - else if (arg == "--num-heads") heads = parse_num(need("num-heads"), "--num-heads"); - else if (arg == "--num-layers") layers = parse_num(need("num-layers"), "--num-layers"); + else if (arg == "--num-heads") + heads = parse_num(need("num-heads"), "--num-heads"); + else if (arg == "--num-layers") + layers = parse_num(need("num-layers"), "--num-layers"); else if (arg == "--d-ff") d_ff = parse_num(need("d-ff"), "--d-ff"); - else if (arg == "--steps") steps = std::max(1, parse_num(need("steps"), "--steps")); - else if (arg == "--kv-steps") kv_steps = parse_num(need("kv-steps"), "--kv-steps"); - else if (arg == "--checkpoint-every") checkpoint_every = parse_num(need("checkpoint-every"), "--checkpoint-every"); - else if (arg == "--flush-interval") flush_interval = parse_num(need("flush-interval"), "--flush-interval"); + else if (arg == "--steps") + steps = std::max(1, parse_num(need("steps"), "--steps")); + else if (arg == "--kv-steps") + kv_steps = parse_num(need("kv-steps"), "--kv-steps"); + else if (arg == "--checkpoint-every") + checkpoint_every = + parse_num(need("checkpoint-every"), "--checkpoint-every"); + else if (arg == "--flush-interval") + flush_interval = parse_num(need("flush-interval"), "--flush-interval"); else if (arg == "--optimizer") opt_name = need("optimizer"); else if (auto g = nn::cli::parse_gpu_option(argc, argv, i, /*allow_name_value=*/true)) { @@ -359,10 +384,10 @@ int main(int argc, char* argv[]) << " mem_probe — 运行时显存分项探针\n" << "========================================\n" << " 配置: vocab=" << vocab << " d_model=" << d_model << " heads=" << heads - << " layers=" << layers << " d_ff=" << d_ff << " seq=" << seq - << " batch=" << batch << "\n" - << " 优化器: " << opt_name << " 精度: " - << (f16 ? (f16_all ? "全f16" : "f16存储") : "f32") + << " layers=" << layers << " d_ff=" << d_ff << " seq=" << seq << " batch=" << batch + << "\n" + << " 优化器: " << opt_name + << " 精度: " << (f16 ? (f16_all ? "全f16" : "f16存储") : "f32") << " checkpoint-every=" << checkpoint_every << " offload=" << (activation_offload ? "on" : "off") << "\n" << "========================================\n\n"; @@ -400,10 +425,10 @@ int main(int argc, char* argv[]) nn::ComputeEngine* engine = raw_engine.get(); // ── 构建模型 ──────────────────────────────────────────────────────── - auto model_r = nn::build_gpt_model(*engine, nn::GptConfig{ - vocab, d_model, seq, heads, d_ff, layers, - nn::PosEncodingType::Learned, nn::ActivationType::GeLU, - nn::NormType::LayerNorm, prof}); + auto model_r = nn::build_gpt_model(*engine, nn::GptConfig{vocab, d_model, seq, heads, d_ff, + layers, nn::PosEncodingType::Learned, + nn::ActivationType::GeLU, + nn::NormType::LayerNorm, prof}); NN_EXIT(model_r, 1, "构建模型失败: "); nn::Model model = std::move(*model_r); model.set_flush_interval(flush_interval); // 与 text_train 默认一致 @@ -417,7 +442,8 @@ int main(int argc, char* argv[]) const std::size_t param_bytes = sum_tensor_bytes(params); const std::size_t grad_bytes = sum_tensor_bytes(grads); std::size_t param_elems = 0; - for (auto& r : params) param_elems += r.get().rows() * r.get().cols(); + for (auto& r : params) + param_elems += r.get().rows() * r.get().cols(); // ── 优化器（状态在构造期分配）────────────────────────────────────── auto optimizer = nn::create_optimizer(opt_name, *engine, params, grads, @@ -432,18 +458,17 @@ int main(int argc, char* argv[]) const auto& b = g_samples.back(); // optimizer-created if (a.pools_valid && b.pools_valid) { - const long long d = static_cast( - (b.persist.allocated_bytes + b.transient.allocated_bytes) - - (a.persist.allocated_bytes + a.transient.allocated_bytes)); + const long long d = + static_cast((b.persist.allocated_bytes + b.transient.allocated_bytes) - + (a.persist.allocated_bytes + a.transient.allocated_bytes)); if (d > 0) opt_measured_bytes = static_cast(d); } } // 逻辑估算（RV-04：查 nn::optimizer_table()，唯一真源）： // state_bufs = 每权重全尺寸状态缓冲数；half_state（AdamWQuant）= 每缓冲 2 B - const auto *opt_meta = nn::find_optimizer(opt_name); + const auto* opt_meta = nn::find_optimizer(opt_name); const std::size_t opt_buffers = (opt_meta != nullptr) ? opt_meta->state_bufs : 0; - const std::size_t opt_state_bytes = - (opt_meta != nullptr && opt_meta->half_state) ? 2u : 4u; + const std::size_t opt_state_bytes = (opt_meta != nullptr && opt_meta->half_state) ? 2u : 4u; const std::size_t opt_logical_bytes = param_elems * opt_state_bytes * opt_buffers; // ── 训练阶段采样（镜像 text_train 顺序）──────────────────────────── @@ -511,9 +536,9 @@ int main(int argc, char* argv[]) act_bytes_fwd = static_cast(g_samples.back().transient.allocated_bytes); std::size_t num_valid = 0; - auto loss_r = ce.forward_sparse_sum(*engine, logits, - std::span(flat_targets), - std::span(flat_mask), vocab, num_valid); + auto loss_r = + ce.forward_sparse_sum(*engine, logits, std::span(flat_targets), + std::span(flat_mask), vocab, num_valid); NN_EXIT(loss_r, 1, "loss 失败: "); take_sample(tag + "loss-fwd"); @@ -554,10 +579,7 @@ int main(int argc, char* argv[]) if (do_kv) { auto* gpt = dynamic_cast(&model.layer_at(0)); - if (!gpt) - { - std::cout << "[KV] layer_at(0) 不是 GPTModel，跳过\n"; - } + if (!gpt) { std::cout << "[KV] layer_at(0) 不是 GPTModel，跳过\n"; } else { take_sample("kv/pre"); @@ -582,7 +604,11 @@ int main(int argc, char* argv[]) for (std::size_t t = 0; t < kv_steps && cur < seq; ++t, ++cur) { auto r = gpt->forward_step(*engine, uid(rng), cur, kc, vc, cur); - if (!r) { std::cerr << "forward_step 失败: " << r.error().message << "\n"; break; } + if (!r) + { + std::cerr << "forward_step 失败: " << r.error().message << "\n"; + break; + } last_logits = std::move(*r); } take_sample("kv/after-steps(batch内)"); @@ -610,16 +636,16 @@ int main(int argc, char* argv[]) peak_t_commit = std::max(peak_t_commit, static_cast(s.transient.total_bytes)); peak_p_dev = std::max(peak_p_dev, static_cast(s.persist.device_bytes)); peak_t_dev = std::max(peak_t_dev, static_cast(s.transient.device_bytes)); - peak_live = std::max(peak_live, static_cast( - s.persist.allocated_bytes + s.transient.allocated_bytes)); + peak_live = std::max(peak_live, static_cast(s.persist.allocated_bytes + + s.transient.allocated_bytes)); } // 稳态（最后一个采样点） const Sample& last = g_samples.back(); - const std::size_t slack = (last.pools_valid) - ? static_cast( - (last.persist.total_bytes - last.persist.allocated_bytes) - + (last.transient.total_bytes - last.transient.allocated_bytes)) - : 0; + const std::size_t slack = + (last.pools_valid) ? static_cast( + (last.persist.total_bytes - last.persist.allocated_bytes) + + (last.transient.total_bytes - last.transient.allocated_bytes)) + : 0; const std::size_t dev_commit = peak_p_dev + peak_t_dev; const double MB = 1.0 / (1024.0 * 1024.0); @@ -635,8 +661,9 @@ int main(int argc, char* argv[]) << "; 实测 Δlive=" << to_mb(opt_measured_bytes) << " MB)\n"; std::cout << " 激活 activations " << std::setw(9) << to_mb(act_bytes_fwd) << " MB (实测 transient.live @forward)\n"; - std::cout << " logits " << std::setw(9) << to_mb(logits_seen ? logits_seen : logits_bytes) - << " MB (vocab×seq×batch×elem" << (f16 ? ", f16" : ", f32") << ")\n"; + std::cout << " logits " << std::setw(9) + << to_mb(logits_seen ? logits_seen : logits_bytes) << " MB (vocab×seq×batch×elem" + << (f16 ? ", f16" : ", f32") << ")\n"; std::cout << " CE 梯度 (同 logits 形) " << std::setw(9) << to_mb(logits_bytes) << " MB (loss-fwd 期与 logits 短暂共存)\n"; if (do_kv) @@ -651,8 +678,7 @@ int main(int argc, char* argv[]) std::cout << " 稳态池空闲(高水位/碎片) " << std::setw(9) << to_mb(slack) << " MB\n"; if (peak_smi >= 0) { - std::cout << " nvidia-smi 整卡峰值 " << std::setw(9) << peak_smi - << " MB\n"; + std::cout << " nvidia-smi 整卡峰值 " << std::setw(9) << peak_smi << " MB\n"; std::cout << " 驱动/池外开销 " << std::setw(9) << (peak_smi - static_cast(dev_commit * MB)) << " MB (= 整卡峰值 − 池 device 提交)\n"; diff --git a/src/mnist_infer.cpp b/src/mnist_infer.cpp index eb7b923c..d220d330 100644 --- a/src/mnist_infer.cpp +++ b/src/mnist_infer.cpp @@ -32,7 +32,7 @@ namespace fs = std::filesystem; using nn::Scalar; // ==================== 帮助信息 ==================== -void print_usage(const char *prog) +void print_usage(const char* prog) { nn::cli::Help help(std::cout, prog, "MNIST 手写数字推理程序"); @@ -43,7 +43,9 @@ void print_usage(const char *prog) help.opt("--model ", "模型文件路径 (默认: pretrained/model.bin)"); help.opt("--topk ", "显示前 n 个预测结果 (默认: 3)"); help.opt("--show-pixels", "显示像素矩阵 (调试用)"); - help.opt("--gpu [索引|名称]", "启用 GPU 加速 (需要 Vulkan SDK)\n空格形式只收枚举索引；名称子串用 --gpu=<名称>，如 --gpu=NVIDIA / --gpu=40HX"); + help.opt("--gpu [索引|名称]", + "启用 GPU 加速 (需要 Vulkan SDK)\n空格形式只收枚举索引；名称子串用 --gpu=<名称>，如 " + "--gpu=NVIDIA / --gpu=40HX"); help.opt("--help, -h", "显示此帮助信息"); } @@ -58,7 +60,7 @@ struct InferConfig std::string gpu_device; // --gpu 的可选设备选择子（空 = 自动选卡） }; -nn::Result parse_args(int argc, char *argv[]) +nn::Result parse_args(int argc, char* argv[]) { InferConfig cfg; bool has_input = false; @@ -71,24 +73,17 @@ nn::Result parse_args(int argc, char *argv[]) print_usage(argv[0]); std::exit(0); } - else if (arg == "--model" && i + 1 < argc) - { - cfg.model_path = argv[++i]; - } + else if (arg == "--model" && i + 1 < argc) { cfg.model_path = argv[++i]; } else if (arg == "--topk" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); NN_TRY_CHECK(v); - if (*v <= 0) - NN_FAIL("--topk 必须为正整数"); + if (*v <= 0) NN_FAIL("--topk 必须为正整数"); if (*v > static_cast(nn::MNIST_NUM_CLASSES)) NN_FAIL("--topk 不能超过类别数 " + std::to_string(nn::MNIST_NUM_CLASSES)); cfg.topk = *v; } - else if (arg == "--show-pixels") - { - cfg.show_pixels = true; - } + else if (arg == "--show-pixels") { cfg.show_pixels = true; } else if (auto dev = nn::cli::parse_gpu_option(argc, argv, i)) { cfg.gpu_enabled = true; @@ -105,8 +100,7 @@ nn::Result parse_args(int argc, char *argv[]) } } - if (!has_input) - NN_FAIL("请指定图片文件或目录，使用 --help 查看用法"); + if (!has_input) NN_FAIL("请指定图片文件或目录，使用 --help 查看用法"); return cfg; } @@ -119,9 +113,9 @@ struct Prediction // 前向：Matrix → engine.from_matrix → model.forward → engine.to_matrix // 后续 softmax / top-k 在 CPU 上完成（单张图片，开销可忽略） -nn::Result> predict_with_confidence( - nn::Model &model, nn::ComputeEngine &engine, - const nn::Matrix &img, int topk) +nn::Result> predict_with_confidence(nn::Model& model, + nn::ComputeEngine& engine, + const nn::Matrix& img, int topk) { auto x_tensor_r = engine.from_matrix(img); NN_TRY_CHECK(x_tensor_r); @@ -129,7 +123,7 @@ nn::Result> predict_with_confidence( NN_TRY_CHECK(logits_tensor_r); auto logits_r = engine.to_matrix(*logits_tensor_r); NN_TRY_CHECK(logits_r); - const auto &logits = *logits_r; + const auto& logits = *logits_r; // Softmax 计算概率 Scalar max_val = logits.at_unchecked(0, 0); @@ -146,12 +140,12 @@ nn::Result> predict_with_confidence( // 防御：logits 全为 -inf 时 sum_exp == 0，回退为均匀分布避免 NaN if (!(sum_exp > 0.0) || !std::isfinite(sum_exp)) { - for (auto &p : probs) + for (auto& p : probs) p = Scalar{1} / static_cast(nn::MNIST_NUM_CLASSES); } else { - for (auto &p : probs) + for (auto& p : probs) p /= sum_exp; } @@ -169,7 +163,7 @@ nn::Result> predict_with_confidence( } // ==================== 显示像素矩阵 ==================== -void show_pixels(const nn::Matrix &img) +void show_pixels(const nn::Matrix& img) { for (int r = 0; r < 28; ++r) { @@ -177,25 +171,20 @@ void show_pixels(const nn::Matrix &img) for (int c = 0; c < 28; ++c) { Scalar val = img.at_unchecked(r * 28 + c, 0); - if (val > 0.5) - std::cout << "##"; - else if (val > 0.1) - std::cout << ".."; - else - std::cout << " "; + if (val > 0.5) std::cout << "##"; + else if (val > 0.1) std::cout << ".."; + else std::cout << " "; } std::cout << "\n"; } } // ==================== 推理单张图片 ==================== -nn::Result infer_single( - nn::Model &model, nn::ComputeEngine &engine, - const std::string &filepath, const InferConfig &cfg) +nn::Result infer_single(nn::Model& model, nn::ComputeEngine& engine, + const std::string& filepath, const InferConfig& cfg) { std::ifstream file(filepath); - if (!file) - NN_FAIL("无法打开文件: " + filepath); + if (!file) NN_FAIL("无法打开文件: " + filepath); std::string line; std::getline(file, line); auto img_result = nn::load_image_from_csv_line(line); @@ -215,15 +204,15 @@ nn::Result infer_single( for (std::size_t k = 0; k < results->size(); ++k) { if (k > 0) std::cout << " "; - std::cout << (*results)[k].digit - << " (" << std::fixed << std::setprecision(1) << (*results)[k].confidence * 100.0 << "%)"; + std::cout << (*results)[k].digit << " (" << std::fixed << std::setprecision(1) + << (*results)[k].confidence * 100.0 << "%)"; } std::cout << std::endl; return {}; } // ==================== 主函数 ==================== -int main(int argc, char *argv[]) +int main(int argc, char* argv[]) { auto cfg_result = parse_args(argc, argv); NN_EXIT(cfg_result, 1, "参数错误: "); @@ -237,18 +226,9 @@ int main(int argc, char *argv[]) NN_EXIT(role_r, 1, "读取模型文件失败: "); // 支持的模型类型：MLP、Transformer（ViT）、CNN - if (spec.is_mlp()) - { - std::cout << "从模型文件读取规格: MLP\n"; - } - else if (spec.is_transformer()) - { - std::cout << "从模型文件读取规格: Transformer (ViT)\n"; - } - else if (spec.is_cnn()) - { - std::cout << "从模型文件读取规格: CNN (LeNet)\n"; - } + if (spec.is_mlp()) { std::cout << "从模型文件读取规格: MLP\n"; } + else if (spec.is_transformer()) { std::cout << "从模型文件读取规格: Transformer (ViT)\n"; } + else if (spec.is_cnn()) { std::cout << "从模型文件读取规格: CNN (LeNet)\n"; } else { std::cerr << "不支持的模型类型 (type=" << spec.type_name() << ")\n"; @@ -272,8 +252,9 @@ int main(int argc, char *argv[]) { auto load_result = nn::load_any_model(cfg.model_path, model, nullptr, nullptr); NN_EXIT(load_result, 1, "加载模型失败: "); - std::cout << "模型已加载（角色: " << nn::model_role_name(*role_r) << "）: " - << cfg.model_path << "\n" << std::endl; + std::cout << "模型已加载（角色: " << nn::model_role_name(*role_r) + << "）: " << cfg.model_path << "\n" + << std::endl; } // 推理模式：BatchNorm 使用保存的 running 统计量（batch 可能为 1，不能算 batch 统计） @@ -285,10 +266,9 @@ int main(int argc, char *argv[]) { // 批量推理 std::vector csv_files; - for (auto &entry : fs::directory_iterator(input)) + for (auto& entry : fs::directory_iterator(input)) { - if (entry.path().extension() == ".csv") - csv_files.push_back(entry.path()); + if (entry.path().extension() == ".csv") csv_files.push_back(entry.path()); } std::sort(csv_files.begin(), csv_files.end()); @@ -300,10 +280,14 @@ int main(int argc, char *argv[]) std::cout << "找到 " << csv_files.size() << " 个文件\n" << std::endl; - for (auto &f : csv_files) + for (auto& f : csv_files) { auto result = infer_single(model, *engine, f.string(), cfg); - if (!result) { std::cerr << "推理失败: " << f.filename().string() << ": " << result.error().message << '\n'; } + if (!result) + { + std::cerr << "推理失败: " << f.filename().string() << ": " << result.error().message + << '\n'; + } } std::cout << "\n共推理 " << csv_files.size() << " 张图片" << std::endl; diff --git a/src/mnist_train.cpp b/src/mnist_train.cpp index 6b67be24..5e6804b9 100644 --- a/src/mnist_train.cpp +++ b/src/mnist_train.cpp @@ -45,8 +45,7 @@ using nn::Scalar; // allow_t1_58：三值（T1_58）**只适用于 param 槽**（"该层权重是三值"）， // compute/stable/optimizer 出现它一律在参数解析处拒绝，而不是拖到 init 期才报 // 能力校验错误（docs/internal/21-quantized-weights.md §4.1）。 -nn::Precision parse_precision(const std::string& name, const char* flag, - bool allow_t1_58 = false) +nn::Precision parse_precision(const std::string& name, const char* flag, bool allow_t1_58 = false) { if (name == "f16" || name == "half") return nn::Precision::F16; if (name == "f32" || name == "float") return nn::Precision::F32; @@ -57,15 +56,20 @@ nn::Precision parse_precision(const std::string& name, const char* flag, << "，可选: f16, f32\n"; std::exit(1); } - std::cerr << "无效 --" << flag << ": " << name << "，可选: " - << (allow_t1_58 ? "f16, f32, t1_58" : "f16, f32") << "\n"; + std::cerr << "无效 --" << flag << ": " << name + << "，可选: " << (allow_t1_58 ? "f16, f32, t1_58" : "f16, f32") << "\n"; std::exit(1); } -enum class ArchType { MLP, Transformer, CNN }; +enum class ArchType +{ + MLP, + Transformer, + CNN +}; // ==================== 帮助信息 ==================== -void print_usage(const char *prog) +void print_usage(const char* prog) { nn::cli::Help help(std::cout, prog, "MNIST 手写数字训练程序 (MLP / Transformer / CNN)"); @@ -73,28 +77,40 @@ void print_usage(const char *prog) help.section("选项"); help.opt("--arch ", "模型架构: mlp/transformer/cnn (默认: mlp)"); - help.opt("--resume ", "从已有模型恢复训练 (自动读取模型规格与架构；\n推理模型 .nnmodel 没有训练参数，会被明确拒绝)"); - help.opt("--save ", "推理模型保存路径 (默认: mnist_model.nnmodel；\n只含推理所需参数——三值层写打包形态，不含训练参数)"); + help.opt("--resume ", "从已有模型恢复训练 (自动读取模型规格与架构；\n推理模型 .nnmodel " + "没有训练参数，会被明确拒绝)"); + help.opt("--save ", + "推理模型保存路径 (默认: " + "mnist_model.nnmodel；\n只含推理所需参数——三值层写打包形态，不含训练参数)"); help.opt("--dataset ", "数据集目录 (默认: datasets/mnist_data)"); help.opt("--epochs ", "训练轮数 (默认: 10)"); help.opt("--lr ", "学习率 (默认: 0.001)"); help.opt("--batch-size ", "批大小 (默认: 64)"); // RV-04：名单文本由 nn::optimizer_table()（唯一真源）生成，不再手抄 - help.opt("--optimizer ", - "优化器: " + nn::optimizer_names("/") + - " (默认: adam)\nadamwquant: 量化状态的 AdamW（状态 4 B/权重 vs 8 B；见 docs/internal/21 §4.13）"); + help.opt("--optimizer ", "优化器: " + nn::optimizer_names("/") + + " (默认: adam)\nadamwquant: 量化状态的 AdamW（状态 4 B/权重 " + "vs 8 B；见 docs/internal/21 §4.13）"); help.opt("--weight-decay ", "AdamW 权重衰减系数 (默认: 0.01)"); help.opt("--max-samples ", "限制训练样本数 (默认: 全部，用于快速测试)"); help.opt("--shuffle-steps ", "每 epoch 打乱 batch 顺序 (默认: true)"); - help.opt("--gpu [索引|名称]", "启用 GPU 加速 (需要 Vulkan SDK)\n空格形式只收枚举索引；名称子串用 --gpu=<名称>，如 --gpu=NVIDIA / --gpu=40HX"); + help.opt("--gpu [索引|名称]", + "启用 GPU 加速 (需要 Vulkan SDK)\n空格形式只收枚举索引；名称子串用 --gpu=<名称>，如 " + "--gpu=NVIDIA / --gpu=40HX"); help.opt("--help, -h", "显示此帮助信息"); help.section("MLP 专用"); help.opt("--layer-dims ", "各层维度，逗号分隔 (默认: 784,512,256,128,64,10)"); help.section("归一化"); - help.opt("--norm ", "归一化层: auto/layernorm/rmsnorm/batchnorm (默认: auto)\nauto = 按架构取最合适: MLP/ViT = layernorm, CNN = batchnorm\nlayernorm: 按特征维归一化\nrmsnorm: 无均值、更轻量\nbatchnorm: 沿 batch 维归一化"); - help.opt("--norm-place ", "归一化挂载位置: auto/none/conv/head/both/final (默认: auto)\nauto = 按架构: CNN = conv, ViT = final, MLP = 结构内置(不额外加)\nconv: CNN 每个卷积(+池化)之后; head: CNN 全连接头 Linear→Norm→ReLU; both: 两者\nfinal: ViT 编码器末端 final norm (全局平均池化之前)\nMLP 的挂载位置由结构固定，只接受 none/auto"); + help.opt("--norm ", + "归一化层: auto/layernorm/rmsnorm/batchnorm (默认: auto)\nauto = 按架构取最合适: " + "MLP/ViT = layernorm, CNN = batchnorm\nlayernorm: 按特征维归一化\nrmsnorm: " + "无均值、更轻量\nbatchnorm: 沿 batch 维归一化"); + help.opt("--norm-place ", + "归一化挂载位置: auto/none/conv/head/both/final (默认: auto)\nauto = 按架构: CNN = " + "conv, ViT = final, MLP = 结构内置(不额外加)\nconv: CNN 每个卷积(+池化)之后; head: " + "CNN 全连接头 Linear→Norm→ReLU; both: 两者\nfinal: ViT 编码器末端 final norm " + "(全局平均池化之前)\nMLP 的挂载位置由结构固定，只接受 none/auto"); help.section("Transformer 专用"); help.opt("--d-model ", "模型维度 (默认: 64)"); @@ -111,14 +127,17 @@ void print_usage(const char *prog) help.opt("--cnn-fc ", "展平后的全连接头 (默认: 120,10；末位为类别数)"); help.section("学习率调度"); - help.opt("--lr-schedule ", "学习率调度: fixed/constant/cosine (默认: fixed)\ncosine: 余弦退火 (epoch 级)，lr 从初始值衰减到 min-lr"); + help.opt("--lr-schedule ", "学习率调度: fixed/constant/cosine (默认: fixed)\ncosine: " + "余弦退火 (epoch 级)，lr 从初始值衰减到 min-lr"); help.opt("--warmup-epochs ", "线性预热轮数 (默认: 0，即不预热)"); help.opt("--min-lr ", "余弦退火最低学习率 (默认: 1e-6)"); help.opt("--lr-per-epoch ", "手动指定每轮学习率 (逗号分隔，优先级最高)"); help.section("混合精度 (docs/development/05-mixed-precision.md)"); help.opt("--f16", "快捷方式: f16 存储 (param/compute=F16，stable/optimizer=F32)"); - help.opt("--t1_58", "快捷方式: 精度预设 t1_58 = f16 存储 + 三值权重\n(param=T1_58，compute=F16，stable/optimizer=F32；等价于 --f16 再把 param 换成 t1_58)"); + help.opt("--t1_58", "快捷方式: 精度预设 t1_58 = f16 存储 + " + "三值权重\n(param=T1_58，compute=F16，stable/optimizer=F32；等价于 --f16 " + "再把 param 换成 t1_58)"); help.opt("--precision-param ", "权重/参数精度 (默认: f32)；t1_58 = 三值权重 BitLinear（1.58-bit）；" "本入口仅 MLP 架构支持（transformer/CNN 由层能力校验报错），" @@ -163,8 +182,8 @@ struct TrainConfig // CNN 参数（LeNet-5 风格默认） std::size_t cnn_in_channels = 1; - std::size_t cnn_in_size = 28; - std::size_t cnn_pool = 2; + std::size_t cnn_in_size = 28; + std::size_t cnn_pool = 2; std::vector cnn_channels; // 每卷积输出通道 std::vector cnn_kernels; // 每卷积核大小 std::vector cnn_fc_dims; // 全连接头（不含展平输入） @@ -179,7 +198,7 @@ struct TrainConfig nn::PrecisionProfile precision; // 默认全 F32（零回归） }; -TrainConfig parse_args(int argc, char *argv[]) +TrainConfig parse_args(int argc, char* argv[]) { TrainConfig cfg; nn::cli::TrainCommonArgs common; @@ -198,19 +217,13 @@ TrainConfig parse_args(int argc, char *argv[]) std::exit(0); } // 通用训练参数委托给 nn::cli::parse_train_common_args - else if (nn::cli::parse_train_common_args(argc, argv, i, common)) - { - continue; - } + else if (nn::cli::parse_train_common_args(argc, argv, i, common)) { continue; } else if (arg == "--arch" && i + 1 < argc) { std::string v = argv[++i]; - if (v == "mlp") - cfg.arch = ArchType::MLP; - else if (v == "transformer" || v == "tf") - cfg.arch = ArchType::Transformer; - else if (v == "cnn") - cfg.arch = ArchType::CNN; + if (v == "mlp") cfg.arch = ArchType::MLP; + else if (v == "transformer" || v == "tf") cfg.arch = ArchType::Transformer; + else if (v == "cnn") cfg.arch = ArchType::CNN; else { std::cerr << "未知 --arch: " << v << "，可选: mlp, transformer, cnn\n"; @@ -222,14 +235,8 @@ TrainConfig parse_args(int argc, char *argv[]) cfg.resume_path = argv[++i]; cfg.load_existing = true; } - else if (arg == "--save" && i + 1 < argc) - { - cfg.save_path = argv[++i]; - } - else if (arg == "--dataset" && i + 1 < argc) - { - cfg.dataset_path = argv[++i]; - } + else if (arg == "--save" && i + 1 < argc) { cfg.save_path = argv[++i]; } + else if (arg == "--dataset" && i + 1 < argc) { cfg.dataset_path = argv[++i]; } else if (arg == "--layer-dims" && i + 1 < argc) { std::string dims_str = argv[++i]; @@ -238,8 +245,16 @@ TrainConfig parse_args(int argc, char *argv[]) while (std::getline(ss, token, ',')) { auto v = nn::parse_number(token); - if (!v) { std::cerr << "无效层维度: " << v.error().message << "\n"; std::exit(1); } - if (*v == 0) { std::cerr << "层维度不能为 0\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效层维度: " << v.error().message << "\n"; + std::exit(1); + } + if (*v == 0) + { + std::cerr << "层维度不能为 0\n"; + std::exit(1); + } cfg.layer_dims.push_back(*v); } if (cfg.layer_dims.size() < 2) @@ -251,29 +266,26 @@ TrainConfig parse_args(int argc, char *argv[]) else if (arg == "--norm" && i + 1 < argc) { std::string v = argv[++i]; - if (v == "auto") - cfg.norm_type = std::nullopt; // 回到按架构默认 - else if (v == "layernorm") - cfg.norm_type = nn::NormType::LayerNorm; - else if (v == "rmsnorm") - cfg.norm_type = nn::NormType::RMSNorm; - else if (v == "batchnorm") - cfg.norm_type = nn::NormType::BatchNorm; + if (v == "auto") cfg.norm_type = std::nullopt; // 回到按架构默认 + else if (v == "layernorm") cfg.norm_type = nn::NormType::LayerNorm; + else if (v == "rmsnorm") cfg.norm_type = nn::NormType::RMSNorm; + else if (v == "batchnorm") cfg.norm_type = nn::NormType::BatchNorm; else { - std::cerr << "无效 --norm: " << v << "，可选: auto, layernorm, rmsnorm, batchnorm\n"; + std::cerr << "无效 --norm: " << v + << "，可选: auto, layernorm, rmsnorm, batchnorm\n"; std::exit(1); } } else if (arg == "--norm-place" && i + 1 < argc) { std::string v = argv[++i]; - if (v == "auto") cfg.norm_place = std::nullopt; // 回到按架构默认 - else if (v == "none") cfg.norm_place = nn::NormPlace::None; - else if (v == "conv") cfg.norm_place = nn::NormPlace::Conv; - else if (v == "head") cfg.norm_place = nn::NormPlace::Head; - else if (v == "both") cfg.norm_place = nn::NormPlace::Both; - else if (v == "final") cfg.norm_place = nn::NormPlace::Final; + if (v == "auto") cfg.norm_place = std::nullopt; // 回到按架构默认 + else if (v == "none") cfg.norm_place = nn::NormPlace::None; + else if (v == "conv") cfg.norm_place = nn::NormPlace::Conv; + else if (v == "head") cfg.norm_place = nn::NormPlace::Head; + else if (v == "both") cfg.norm_place = nn::NormPlace::Both; + else if (v == "final") cfg.norm_place = nn::NormPlace::Final; else { std::cerr << "无效 --norm-place: " << v @@ -284,35 +296,55 @@ TrainConfig parse_args(int argc, char *argv[]) else if (arg == "--d-model" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --d-model: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --d-model: " << v.error().message << "\n"; + std::exit(1); + } cfg.d_model = *v; } else if (arg == "--num-heads" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --num-heads: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --num-heads: " << v.error().message << "\n"; + std::exit(1); + } cfg.num_heads = *v; } else if (arg == "--num-layers" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --num-layers: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --num-layers: " << v.error().message << "\n"; + std::exit(1); + } cfg.num_layers = *v; } else if (arg == "--d-ff" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --d-ff: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --d-ff: " << v.error().message << "\n"; + std::exit(1); + } cfg.d_ff = *v; } else if (arg == "--patch-size" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --patch-size: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --patch-size: " << v.error().message << "\n"; + std::exit(1); + } if (*v == 0 || nn::MNIST_IMG_SIZE % *v != 0) { - std::cerr << "--patch-size 必须能整除 " << nn::MNIST_IMG_SIZE - << "（当前: " << *v << "）\n"; + std::cerr << "--patch-size 必须能整除 " << nn::MNIST_IMG_SIZE << "（当前: " << *v + << "）\n"; std::exit(1); } cfg.patch_size = *v; @@ -320,13 +352,21 @@ TrainConfig parse_args(int argc, char *argv[]) else if (arg == "--max-samples" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --max-samples: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --max-samples: " << v.error().message << "\n"; + std::exit(1); + } cfg.max_train_samples = *v; } else if (arg == "--eval-samples" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --eval-samples: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --eval-samples: " << v.error().message << "\n"; + std::exit(1); + } cfg.eval_samples = *v; } else if (arg == "--cnn-channels" && i + 1 < argc) @@ -337,11 +377,23 @@ TrainConfig parse_args(int argc, char *argv[]) while (std::getline(ss, token, ',')) { auto v = nn::parse_number(token); - if (!v) { std::cerr << "无效 --cnn-channels: " << v.error().message << "\n"; std::exit(1); } - if (*v == 0) { std::cerr << "--cnn-channels 不能为 0\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --cnn-channels: " << v.error().message << "\n"; + std::exit(1); + } + if (*v == 0) + { + std::cerr << "--cnn-channels 不能为 0\n"; + std::exit(1); + } cfg.cnn_channels.push_back(*v); } - if (cfg.cnn_channels.empty()) { std::cerr << "--cnn-channels 至少需要一个值\n"; std::exit(1); } + if (cfg.cnn_channels.empty()) + { + std::cerr << "--cnn-channels 至少需要一个值\n"; + std::exit(1); + } } else if (arg == "--cnn-kernels" && i + 1 < argc) { @@ -351,16 +403,32 @@ TrainConfig parse_args(int argc, char *argv[]) while (std::getline(ss, token, ',')) { auto v = nn::parse_number(token); - if (!v) { std::cerr << "无效 --cnn-kernels: " << v.error().message << "\n"; std::exit(1); } - if (*v == 0) { std::cerr << "--cnn-kernels 不能为 0\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --cnn-kernels: " << v.error().message << "\n"; + std::exit(1); + } + if (*v == 0) + { + std::cerr << "--cnn-kernels 不能为 0\n"; + std::exit(1); + } cfg.cnn_kernels.push_back(*v); } - if (cfg.cnn_kernels.empty()) { std::cerr << "--cnn-kernels 至少需要一个值\n"; std::exit(1); } + if (cfg.cnn_kernels.empty()) + { + std::cerr << "--cnn-kernels 至少需要一个值\n"; + std::exit(1); + } } else if (arg == "--cnn-pool" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --cnn-pool: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --cnn-pool: " << v.error().message << "\n"; + std::exit(1); + } if (*v != 0 && *v > nn::MNIST_IMG_SIZE) { std::cerr << "--cnn-pool 不能大于输入边长 " << nn::MNIST_IMG_SIZE @@ -377,19 +445,29 @@ TrainConfig parse_args(int argc, char *argv[]) while (std::getline(ss, token, ',')) { auto v = nn::parse_number(token); - if (!v) { std::cerr << "无效 --cnn-fc: " << v.error().message << "\n"; std::exit(1); } - if (*v == 0) { std::cerr << "--cnn-fc 不能为 0\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --cnn-fc: " << v.error().message << "\n"; + std::exit(1); + } + if (*v == 0) + { + std::cerr << "--cnn-fc 不能为 0\n"; + std::exit(1); + } cfg.cnn_fc_dims.push_back(*v); } - if (cfg.cnn_fc_dims.size() < 2) { std::cerr << "--cnn-fc 至少需要 2 个维度\n"; std::exit(1); } + if (cfg.cnn_fc_dims.size() < 2) + { + std::cerr << "--cnn-fc 至少需要 2 个维度\n"; + std::exit(1); + } } else if (arg == "--shuffle-steps" && i + 1 < argc) { std::string v = argv[++i]; - if (v == "true" || v == "1") - cfg.shuffle_steps = true; - else if (v == "false" || v == "0") - cfg.shuffle_steps = false; + if (v == "true" || v == "1") cfg.shuffle_steps = true; + else if (v == "false" || v == "0") cfg.shuffle_steps = false; else { std::cerr << "无效 --shuffle-steps: " << v << "，可选: true, false\n"; @@ -447,28 +525,22 @@ TrainConfig parse_args(int argc, char *argv[]) // ==================== 构建模型规格 ==================== // 从 CLI 配置构建 CNN 配置（未指定的项用 LeNet-5 默认补齐） -nn::CnnConfig build_cnn_config(const TrainConfig &cfg) +nn::CnnConfig build_cnn_config(const TrainConfig& cfg) { nn::CnnConfig c; c.in_channels = cfg.cnn_in_channels; - c.in_size = cfg.cnn_in_size; - c.pool = cfg.cnn_pool; + c.in_size = cfg.cnn_in_size; + c.pool = cfg.cnn_pool; std::vector channels = cfg.cnn_channels; - std::vector kernels = cfg.cnn_kernels; + std::vector kernels = cfg.cnn_kernels; if (channels.empty() && kernels.empty()) { channels = {6, 16}; - kernels = {5, 5}; - } - else if (channels.empty()) - { - channels.assign(kernels.size(), 6); - } - else if (kernels.empty()) - { - kernels.assign(channels.size(), 5); + kernels = {5, 5}; } + else if (channels.empty()) { channels.assign(kernels.size(), 6); } + else if (kernels.empty()) { kernels.assign(channels.size(), 5); } else if (channels.size() != kernels.size()) { std::cerr << "--cnn-channels 与 --cnn-kernels 长度必须一致\n"; @@ -481,7 +553,7 @@ nn::CnnConfig build_cnn_config(const TrainConfig &cfg) return c; } -nn::ModelSpec build_spec(const TrainConfig &cfg) +nn::ModelSpec build_spec(const TrainConfig& cfg) { if (cfg.arch == ArchType::Transformer) { @@ -494,8 +566,7 @@ nn::ModelSpec build_spec(const TrainConfig &cfg) std::exit(1); } return nn::make_mnist_transformer_spec( - cfg.patch_size, - cfg.d_model, cfg.num_heads, cfg.d_ff, cfg.num_layers, + cfg.patch_size, cfg.d_model, cfg.num_heads, cfg.d_ff, cfg.num_layers, cfg.norm_type.value_or(nn::NormType::LayerNorm), place); } if (cfg.arch == ArchType::CNN) @@ -509,11 +580,11 @@ nn::ModelSpec build_spec(const TrainConfig &cfg) std::exit(1); } nn::CnnConfig c = build_cnn_config(cfg); - c.norm_type = cfg.norm_type.value_or(nn::NormType::BatchNorm); + c.norm_type = cfg.norm_type.value_or(nn::NormType::BatchNorm); c.norm_place = place; std::vector convs = c.convs; - return nn::make_cnn_spec(c.in_channels, c.in_size, c.pool, convs, c.fc_dims, - c.norm_type, c.norm_place); + return nn::make_cnn_spec(c.in_channels, c.in_size, c.pool, convs, c.fc_dims, c.norm_type, + c.norm_place); } if (cfg.norm_place && *cfg.norm_place != nn::NormPlace::None) { @@ -523,26 +594,25 @@ nn::ModelSpec build_spec(const TrainConfig &cfg) } // MLP（结构内置 Linear→Norm→激活；auto = layernorm） nn::ModelSpec spec; - auto &m = spec.emplace(); + auto& m = spec.emplace(); m.layer_dims = cfg.layer_dims.empty() ? nn::MNIST_LAYER_DIMS : cfg.layer_dims; m.norm_type = cfg.norm_type.value_or(nn::NormType::LayerNorm); // 三值权重：param=t1_58 → 线性层用 BitLinear（规格是加载时的权威来源， // 缺键 = None 保证旧文件零破坏）。见 docs 21 §4.7。 - m.weight_quant = (cfg.precision.param == nn::Precision::T1_58) - ? nn::WeightQuant::T1_58 : nn::WeightQuant::None; + m.weight_quant = (cfg.precision.param == nn::Precision::T1_58) ? nn::WeightQuant::T1_58 + : nn::WeightQuant::None; return spec; } // ==================== 主函数 ==================== -int main(int argc, char *argv[]) +int main(int argc, char* argv[]) { TrainConfig cfg = parse_args(argc, argv); // ── 危险精度组合启动期告警（仅警告，不拒绝；docs/internal/24 §2.3 F1）──── // optimizer=f16（含四字段全 f16）只打印一次到 stderr：不改变退出码、不阻断 // 训练、不静默降级；`--f16` 与合法组合返回空串（静默）。 - if (const std::string warn = nn::cli::dangerous_precision_warning(cfg.precision); - !warn.empty()) + if (const std::string warn = nn::cli::dangerous_precision_warning(cfg.precision); !warn.empty()) std::cerr << warn; // ── 构建规格 ───────────────────────────────────────────── @@ -579,7 +649,7 @@ int main(int argc, char *argv[]) std::cout << "从模型文件读取 Transformer 规格\n"; spec = std::move(*spec_result); cfg.arch = ArchType::Transformer; - const auto *t = spec.as(); + const auto* t = spec.as(); cfg.norm_type = t->norm_type; cfg.norm_place = t->norm_place; cfg.patch_size = t->patch_size != 0 ? t->patch_size : nn::MNIST_PATCH_SIZE; @@ -593,20 +663,19 @@ int main(int argc, char *argv[]) std::cout << "从模型文件读取 CNN 规格\n"; spec = std::move(*spec_result); cfg.arch = ArchType::CNN; - const auto *c = spec.as(); + const auto* c = spec.as(); cfg.norm_type = c->norm_type; cfg.norm_place = c->norm_place; cfg.cnn_in_channels = c->cnn_in_channels; - cfg.cnn_in_size = c->cnn_in_size; - cfg.cnn_pool = c->cnn_pool; - cfg.cnn_channels = c->cnn_channels; - cfg.cnn_kernels = c->cnn_kernels; - cfg.cnn_fc_dims = c->layer_dims; + cfg.cnn_in_size = c->cnn_in_size; + cfg.cnn_pool = c->cnn_pool; + cfg.cnn_channels = c->cnn_channels; + cfg.cnn_kernels = c->cnn_kernels; + cfg.cnn_fc_dims = c->layer_dims; } else { - std::cerr << "模型文件类型不支持 (type=" << spec_result->type_name() - << ")\n"; + std::cerr << "模型文件类型不支持 (type=" << spec_result->type_name() << ")\n"; return 1; } } @@ -621,8 +690,11 @@ int main(int argc, char *argv[]) std::cout << "========================================\n"; std::cout << " MNIST 手写数字训练\n"; std::cout << "========================================\n"; - std::cout << " 架构: " << (cfg.arch == ArchType::Transformer ? "Transformer (ViT)" - : cfg.arch == ArchType::CNN ? "CNN (LeNet)" : "MLP") << "\n"; + std::cout << " 架构: " + << (cfg.arch == ArchType::Transformer ? "Transformer (ViT)" + : cfg.arch == ArchType::CNN ? "CNN (LeNet)" + : "MLP") + << "\n"; if (cfg.arch == ArchType::CNN) { @@ -646,23 +718,21 @@ int main(int argc, char *argv[]) { const std::size_t grid = nn::MNIST_IMG_SIZE / cfg.patch_size; const std::size_t num_patches = grid * grid; - std::cout << " Patch: " << nn::MNIST_IMG_SIZE << "/" << cfg.patch_size - << " → " << grid << "×" << grid << "=" << num_patches << " patches\n"; - std::cout << " d_model: " << cfg.d_model - << " heads: " << cfg.num_heads - << " layers: " << cfg.num_layers - << " d_ff: " << cfg.d_ff << "\n"; + std::cout << " Patch: " << nn::MNIST_IMG_SIZE << "/" << cfg.patch_size << " → " << grid + << "×" << grid << "=" << num_patches << " patches\n"; + std::cout << " d_model: " << cfg.d_model << " heads: " << cfg.num_heads + << " layers: " << cfg.num_layers << " d_ff: " << cfg.d_ff << "\n"; } else { - const auto &dims = + const auto& dims = spec.is_mlp() ? spec.as()->layer_dims : nn::MNIST_LAYER_DIMS; - const auto p_norm = spec.is_mlp() ? spec.as()->norm_type + const auto p_norm = spec.is_mlp() ? spec.as()->norm_type : spec.is_transformer() ? spec.as()->norm_type : spec.as()->norm_type; - const std::string_view norm_name = - (p_norm == nn::NormType::RMSNorm) ? "RMSNorm" : - (p_norm == nn::NormType::BatchNorm) ? "BatchNorm" : "LayerNorm"; + const std::string_view norm_name = (p_norm == nn::NormType::RMSNorm) ? "RMSNorm" + : (p_norm == nn::NormType::BatchNorm) ? "BatchNorm" + : "LayerNorm"; std::cout << " 网络: "; for (std::size_t i = 0; i < dims.size(); ++i) { @@ -670,10 +740,8 @@ int main(int argc, char *argv[]) // 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 << " -> "; + if (i > 0 && i < dims.size() - 1) std::cout << "(" << norm_name << "+GeLU)"; + if (i < dims.size() - 1) std::cout << " -> "; } std::cout << "\n"; } @@ -685,17 +753,17 @@ int main(int argc, char *argv[]) : spec.is_cnn() ? spec.as()->norm_place : nn::NormPlace::None; std::cout << " 归一化: " - << ((q_norm == nn::NormType::RMSNorm) ? "RMSNorm" : - (q_norm == nn::NormType::BatchNorm) ? "BatchNorm" : "LayerNorm"); - if (q_place != nn::NormPlace::None) - std::cout << " @ " << nn::norm_place_name(q_place); + << ((q_norm == nn::NormType::RMSNorm) ? "RMSNorm" + : (q_norm == nn::NormType::BatchNorm) ? "BatchNorm" + : "LayerNorm"); + if (q_place != nn::NormPlace::None) std::cout << " @ " << nn::norm_place_name(q_place); std::cout << "\n"; std::cout << " 优化器: " << cfg.optimizer_name << " 学习率: " << cfg.lr << "\n"; std::cout << " 轮数: " << cfg.epochs << " 批大小: " << cfg.batch_size << "\n"; std::cout << " GPU: " << (cfg.gpu_enabled ? "启用" : "禁用") << " 打乱: " << (cfg.shuffle_steps ? "启用" : "禁用") << "\n"; - std::cout << " 模型: " << (cfg.load_existing ? cfg.resume_path : "(从头训练)") - << " -> " << cfg.save_path << "\n"; + std::cout << " 模型: " << (cfg.load_existing ? cfg.resume_path : "(从头训练)") << " -> " + << cfg.save_path << "\n"; std::cout << "========================================\n\n"; // ── 创建计算引擎 ───────────────────────────────────────── @@ -716,10 +784,8 @@ int main(int argc, char *argv[]) "latent 恒 f32）\n" " [范围] P1 支持 MLP 路径；docs/internal/21-quantized-weights.md\n"; } - if (cfg.precision.param == nn::Precision::F16 || - cfg.precision.compute == nn::Precision::F16 || - cfg.precision.stable == nn::Precision::F16 || - cfg.precision.optimizer == nn::Precision::F16) + if (cfg.precision.param == nn::Precision::F16 || cfg.precision.compute == nn::Precision::F16 || + cfg.precision.stable == nn::Precision::F16 || cfg.precision.optimizer == nn::Precision::F16) { std::cout << "[精度] f16 存储已启用（基类边界 cast）\n" " [注意] 边界 cast 的 transient 放大：峰值可能高于 f32，" @@ -738,7 +804,8 @@ int main(int argc, char *argv[]) auto csv_test_result = nn::cli::load_mnist_csv(cfg.dataset_path + "/test.csv"); NN_EXIT(csv_test_result, 1, "Error: "); auto [test_x, test_y] = std::move(*csv_test_result); - std::cout << "训练集: " << train_x.cols() << " 样本, 测试集: " << test_x.cols() << " 样本\n" << std::endl; + std::cout << "训练集: " << train_x.cols() << " 样本, 测试集: " << test_x.cols() << " 样本\n" + << std::endl; // ── 构建模型（绑定引擎） ───────────────────────────────── // 显式告知而不是静默失效：CNN 构建器目前**不接收**精度 profile @@ -779,14 +846,11 @@ int main(int argc, char *argv[]) // （mnist 侧的 role 语义接线是 P1-2 的遗留项，见 docs/history/history.md）。 auto load_result = nn::load_model(cfg.resume_path, model, nullptr, nullptr, nn::WeightLoadIntent::Trainable); - if (load_result) - { - std::cout << "已加载模型: " << cfg.resume_path << "\n" << std::endl; - } + if (load_result) { std::cout << "已加载模型: " << cfg.resume_path << "\n" << std::endl; } else { - std::cerr << "加载模型失败: " << load_result.error().message - << "，将从头训练。\n" << std::endl; + std::cerr << "加载模型失败: " << load_result.error().message << "，将从头训练。\n" + << std::endl; } } @@ -823,8 +887,8 @@ int main(int argc, char *argv[]) const std::size_t num_batches = train_x.cols() / cfg.batch_size; if (num_batches == 0) { - std::cerr << "训练样本数 (" << train_x.cols() << ") 小于 batch_size (" - << cfg.batch_size << ")，请减小 --batch-size 或增大 --max-samples\n"; + std::cerr << "训练样本数 (" << train_x.cols() << ") 小于 batch_size (" << cfg.batch_size + << ")，请减小 --batch-size 或增大 --max-samples\n"; return 1; } @@ -913,14 +977,16 @@ int main(int argc, char *argv[]) NN_EXIT(obb, 1, "\nbegin_batch (optimizer) failed: "); auto step_result = session.step(); - if (!step_result) { + if (!step_result) + { std::cerr << "\n优化器 step 失败: " << step_result.error().message << '\n'; (void)engine->end_batch(); return 1; } auto zero_result = session.zero_grad(); - if (!zero_result) { + if (!zero_result) + { std::cerr << "\n优化器 zero_grad 失败: " << zero_result.error().message << '\n'; (void)engine->end_batch(); return 1; @@ -933,11 +999,12 @@ int main(int argc, char *argv[]) if ((batch + 1) % 10 == 0 || batch + 1 == num_batches) { auto batch_now = std::chrono::steady_clock::now(); - Scalar batch_ms = std::chrono::duration(batch_now - ep_start).count(); - std::cout << "\r Epoch " << epoch + 1 << "/" << cfg.epochs - << " batch " << batch + 1 << "/" << num_batches - << " loss: " << std::fixed << std::setprecision(4) << loss - << " time: " << std::setprecision(0) << batch_ms << "ms" + Scalar batch_ms = + std::chrono::duration(batch_now - ep_start).count(); + std::cout << "\r Epoch " << epoch + 1 << "/" << cfg.epochs << " batch " + << batch + 1 << "/" << num_batches << " loss: " << std::fixed + << std::setprecision(4) << loss << " time: " << std::setprecision(0) + << batch_ms << "ms" << " " << std::flush; } } @@ -950,24 +1017,25 @@ int main(int argc, char *argv[]) // 评估分块大小取训练 batch：CNN 的 im2col 展开是 k²·C_in 倍， // 不分块时全量 60000 样本单次 forward 需 ~6.4 GB 显存 → OOM。 const std::size_t eval_n = (cfg.arch == ArchType::Transformer) ? cfg.eval_samples : 0; - auto train_acc_r = nn::cli::evaluate_mnist(session.model(), *engine, train_x, train_y, eval_n, cfg.batch_size); - auto test_acc_r = nn::cli::evaluate_mnist(session.model(), *engine, test_x, test_y, eval_n, cfg.batch_size); + auto train_acc_r = nn::cli::evaluate_mnist(session.model(), *engine, train_x, train_y, + eval_n, cfg.batch_size); + auto test_acc_r = nn::cli::evaluate_mnist(session.model(), *engine, test_x, test_y, eval_n, + cfg.batch_size); if (!train_acc_r || !test_acc_r) { - const auto &err = !train_acc_r ? train_acc_r.error() : test_acc_r.error(); + const auto& err = !train_acc_r ? train_acc_r.error() : test_acc_r.error(); std::cerr << "\n 评估失败: " << err.message << std::endl; return 1; } Scalar train_acc = *train_acc_r; - Scalar test_acc = *test_acc_r; + Scalar test_acc = *test_acc_r; - std::cout << "\r Epoch " << epoch + 1 << "/" << cfg.epochs - << " lr=" << std::scientific << std::setprecision(4) << optimizer_current_lr - << " loss=" << std::fixed << std::setprecision(4) << avg_loss - << " train_acc=" << std::setprecision(2) << train_acc * 100.0 << "%" + std::cout << "\r Epoch " << epoch + 1 << "/" << cfg.epochs << " lr=" << std::scientific + << std::setprecision(4) << optimizer_current_lr << " loss=" << std::fixed + << std::setprecision(4) << avg_loss << " train_acc=" << std::setprecision(2) + << train_acc * 100.0 << "%" << " test_acc=" << test_acc * 100.0 << "%" - << " time=" << std::setprecision(1) << ep_sec << "s" - << std::endl; + << " time=" << std::setprecision(1) << ep_sec << "s" << std::endl; } auto t_end = std::chrono::steady_clock::now(); @@ -975,10 +1043,11 @@ int main(int argc, char *argv[]) // ── 保存模型 ───────────────────────────────────────────── // 推理模型（Model 角色）：三值层写打包形态 + τ（不含训练参数） - auto save_result = nn::save_model(cfg.save_path, session.model(), spec, {}, - nn::ModelRole::Model); + auto save_result = + nn::save_model(cfg.save_path, session.model(), spec, {}, nn::ModelRole::Model); NN_EXIT(save_result, 1, "保存模型失败: "); - std::cout << "\n训练完成! 总耗时: " << std::fixed << std::setprecision(1) << total_sec << "s" << std::endl; + std::cout << "\n训练完成! 总耗时: " << std::fixed << std::setprecision(1) << total_sec << "s" + << std::endl; return 0; } diff --git a/src/model_roles_test.cpp b/src/model_roles_test.cpp index 874594d2..f838a770 100644 --- a/src/model_roles_test.cpp +++ b/src/model_roles_test.cpp @@ -45,7 +45,7 @@ constexpr std::size_t B = 4; nn::ModelSpec ternary_spec() { nn::ModelSpec spec; - auto &sp = spec.emplace(); + auto& sp = spec.emplace(); sp.layer_dims = {IN, H, OUT}; sp.weight_quant = nn::WeightQuant::T1_58; return spec; @@ -56,7 +56,8 @@ nn::ModelSpec ternary_spec() nn::Matrix m(rows, cols); std::mt19937 rng(seed); std::uniform_real_distribution d(-1.0f, 1.0f); - for (auto& v : m.span()) v = d(rng); + for (auto& v : m.span()) + v = d(rng); return m; } @@ -142,10 +143,8 @@ void test_model_and_checkpoint_roles() const std::uint64_t sm = file_size(f_model); const std::uint64_t scp = file_size(f_ckpt); std::printf(" .nnmodel = %llu B, .nnckpt = %llu B\n", - static_cast(sm), - static_cast(scp)); - CHECK(sm > 0 && scp > 0 && sm < scp, - "推理文件小于断点文件（base-3 打包 + 无 latent）"); + static_cast(sm), static_cast(scp)); + CHECK(sm > 0 && scp > 0 && sm < scp, "推理文件小于断点文件（base-3 打包 + 无 latent）"); } // ── [1] Model 角色载入：只读态 + 前向逐位一致 ── @@ -165,16 +164,14 @@ void test_model_and_checkpoint_roles() CHECK(bl != nullptr, "重建模型里有 BitLinear"); if (bl != nullptr) { - CHECK(!bl->latent_weights().valid(), - "推理文件载入后没有 latent（只读推理态）"); + CHECK(!bl->latent_weights().valid(), "推理文件载入后没有 latent（只读推理态）"); auto g = bl->backward(bl->latent_weights()); // 只读态 → 错误（不静默） CHECK(!g.has_value(), "只读态 backward() 明确报错"); } auto y2_r = m2.forward(x); CHECK(y2_r.has_value(), "载入模型前向成功"); if (y2_r.has_value()) - CHECK(matrices_equal(y_ref, download(eng, *y2_r)), - "前向输出与保存前**逐位一致**"); + CHECK(matrices_equal(y_ref, download(eng, *y2_r)), "前向输出与保存前**逐位一致**"); } // ── [2] Checkpoint 角色载入：latent 逐位一致 + 可继续训练 ── @@ -226,8 +223,7 @@ void test_model_and_checkpoint_roles() if (bl != nullptr) { auto e = bl->export_param(0, nn::ModelRole::Checkpoint); - CHECK(!e.has_value(), - "只读推理态按断点角色导出 → 明确错误（不静默）"); + CHECK(!e.has_value(), "只读推理态按断点角色导出 → 明确错误（不静默）"); } } } @@ -325,8 +321,7 @@ void test_ternary_warm_start() { for (std::size_t i = 0; i < wpr; ++i) words[i] = std::bit_cast(src_packed.span()[o * wpr + i]); - const float s = - nn::trit_pack::latent_row_scale(words.data(), IN, tau.span()[o]); + const float s = nn::trit_pack::latent_row_scale(words.data(), IN, tau.span()[o]); for (std::size_t c = 0; c < IN; ++c) if (lat.span()[o * IN + c] != nn::trit_pack::at(words.data(), c) * s) formula_ok = false; @@ -334,18 +329,16 @@ void test_ternary_warm_start() for (std::size_t c = 0; c < IN; ++c) acc += std::fabs(static_cast(lat.span()[o * IN + c])); const double mean = acc / static_cast(IN); - const double tol = - 2e-6 * std::max(1.0, std::fabs(static_cast(tau.span()[o]))); - if (std::fabs(mean - static_cast(tau.span()[o])) > tol) - fixed_point_ok = false; + const double tol = 2e-6 * std::max(1.0, std::fabs(static_cast(tau.span()[o]))); + if (std::fabs(mean - static_cast(tau.span()[o])) > tol) fixed_point_ok = false; } CHECK(formula_ok, "反量化 latent == trit·(τ·K/n)（逐位）"); CHECK(fixed_point_ok, "逐行 absmean(latent) == τ（量化器不动点）"); } // ⑤ 8 步 Adam：backward 可跑、latent（f32 主权重）确实移动 { - auto opt = nn::create_optimizer("adam", eng, m2.parameters(), m2.param_gradients(), - 1e-2f, 0.0f, t1); + auto opt = nn::create_optimizer("adam", eng, m2.parameters(), m2.param_gradients(), 1e-2f, + 0.0f, t1); CHECK(opt != nullptr, "warm start 后能建 Adam 优化器"); if (opt != nullptr) { @@ -361,21 +354,40 @@ void test_ternary_warm_start() { steps_ok = m2.zero_grad().has_value(); auto logits = m2.forward(x); - if (!logits) { steps_ok = false; break; } + if (!logits) + { + steps_ok = false; + break; + } auto loss = ce.forward(eng, *logits, tgt_t); - if (!loss) { steps_ok = false; break; } + if (!loss) + { + steps_ok = false; + break; + } auto grad = ce.backward(); - if (!grad) { steps_ok = false; break; } - if (!m2.backward(*grad)) { steps_ok = false; break; } - if (!opt->step()) { steps_ok = false; break; } + if (!grad) + { + steps_ok = false; + break; + } + if (!m2.backward(*grad)) + { + steps_ok = false; + break; + } + if (!opt->step()) + { + steps_ok = false; + break; + } } CHECK(steps_ok, "warm start 后 8 步 Adam（forward/loss/backward/step）全部成功"); const nn::Matrix lat_after = download(eng, bl->latent_weights()); std::size_t moved = 0; for (std::size_t i = 0; i < lat_before.size() && i < lat_after.size(); ++i) if (lat_before.span()[i] != lat_after.span()[i]) ++moved; - std::printf(" 三值 latent 8 步 Adam：移动 %zu/%zu\n", moved, - lat_before.size()); + std::printf(" 三值 latent 8 步 Adam：移动 %zu/%zu\n", moved, lat_before.size()); CHECK(moved > 0, "8 步 Adam 后 latent 确实移动（三值续训可用）"); CHECK(bl->quantize_calls() > 1, "latent 被写过 ⇒ 量化缓存按写版本重算"); } @@ -405,8 +417,8 @@ void test_ternary_warm_start() nn::Model& a = *a_r; nn::Model& b = *b_r; auto la = nn::load_model(df, a); // 默认 Inference - auto lb = nn::load_model(df, b, nullptr, nullptr, - nn::WeightLoadIntent::Trainable); + auto lb = + nn::load_model(df, b, nullptr, nullptr, nn::WeightLoadIntent::Trainable); CHECK(la.has_value() && lb.has_value(), "两种意图载入 dense 均成功"); if (la.has_value() && lb.has_value()) { diff --git a/src/model_spec_validation_test.cpp b/src/model_spec_validation_test.cpp index 83cd5a8f..f7a51d14 100644 --- a/src/model_spec_validation_test.cpp +++ b/src/model_spec_validation_test.cpp @@ -22,15 +22,15 @@ #include "neuralnet.cpp/nn.hpp" using nn::ActivationType; -using nn::ModelSpec; -using nn::ModelType; using nn::CnnSpec; using nn::GptSpec; using nn::MlpSpec; -using nn::TransformerSpec; -using nn::NormType; +using nn::ModelSpec; +using nn::ModelType; using nn::NormPlace; +using nn::NormType; using nn::PosEncodingType; +using nn::TransformerSpec; namespace { @@ -38,8 +38,7 @@ namespace // 记录断言结果，全部通过时 run_test 返回 0。 void expect(bool cond, const std::string& name, bool& all) { - if (cond) - std::cout << " ✓ " << name << "\n"; + if (cond) std::cout << " ✓ " << name << "\n"; else { std::cout << " ✗ " << name << "\n"; @@ -61,13 +60,15 @@ int run_test() const ModelSpec g1 = nn::make_gpt_spec(100, 64, 32, 4, 128, 2); const ModelSpec g2 = nn::make_gpt_spec(100, 64, 32, 4, 128, 2); const ModelSpec g3 = nn::make_gpt_spec(101, 64, 32, 4, 128, 2); - expect(nn::spec_matches(g1, g2), "GPT 相同 spec 匹配", all); - expect(!nn::spec_matches(g1, g3), "GPT 不同 vocab 不匹配", all); + expect(nn::spec_matches(g1, g2), "GPT 相同 spec 匹配", all); + expect(!nn::spec_matches(g1, g3), "GPT 不同 vocab 不匹配", all); // GPT 家族兼容：ALiBi/RoPE 变体同归 GptSpec（判别符随多态对象走） - ModelSpec g_fam_a = g1; g_fam_a.as()->pos_encoding = PosEncodingType::ALiBi; - ModelSpec g_fam_b = g1; g_fam_b.as()->pos_encoding = PosEncodingType::ALiBi; - expect(nn::spec_matches(g_fam_a, g_fam_b), "GPT 与 ALiBi_GPT 家族兼容（同字段）", all); + ModelSpec g_fam_a = g1; + g_fam_a.as()->pos_encoding = PosEncodingType::ALiBi; + ModelSpec g_fam_b = g1; + g_fam_b.as()->pos_encoding = PosEncodingType::ALiBi; + expect(nn::spec_matches(g_fam_a, g_fam_b), "GPT 与 ALiBi_GPT 家族兼容（同字段）", all); // GPT 跨类型不匹配 ModelSpec g_other = g1; @@ -75,47 +76,64 @@ int run_test() expect(!nn::spec_matches(g1, g_other), "GPT 与 MLP 跨类型不匹配", all); // MLP - ModelSpec m1; m1.emplace().layer_dims = {128, 10}; - ModelSpec m2; m2.emplace().layer_dims = {128, 10}; - ModelSpec m3; m3.emplace().layer_dims = {64, 10}; - expect(nn::spec_matches(m1, m2), "MLP 相同 layer_dims 匹配", all); - expect(!nn::spec_matches(m1, m3), "MLP 不同 layer_dims 不匹配", all); - ModelSpec m4 = m1; m4.as()->norm_type = NormType::RMSNorm; - expect(!nn::spec_matches(m1, m4), "MLP 不同 norm_type 不匹配", all); + ModelSpec m1; + m1.emplace().layer_dims = {128, 10}; + ModelSpec m2; + m2.emplace().layer_dims = {128, 10}; + ModelSpec m3; + m3.emplace().layer_dims = {64, 10}; + expect(nn::spec_matches(m1, m2), "MLP 相同 layer_dims 匹配", all); + expect(!nn::spec_matches(m1, m3), "MLP 不同 layer_dims 不匹配", all); + ModelSpec m4 = m1; + m4.as()->norm_type = NormType::RMSNorm; + expect(!nn::spec_matches(m1, m4), "MLP 不同 norm_type 不匹配", all); // Transformer ModelSpec t1; { - auto &tt = t1.emplace(); - tt.d_model = 64; tt.num_heads = 4; tt.d_ff = 128; tt.num_layers = 2; + auto& tt = t1.emplace(); + tt.d_model = 64; + tt.num_heads = 4; + tt.d_ff = 128; + tt.num_layers = 2; tt.patch_size = 7; } ModelSpec t2 = t1; - ModelSpec t3 = t1; t3.as()->num_heads = 8; - expect(nn::spec_matches(t1, t2), "Transformer 相同 spec 匹配", all); - expect(!nn::spec_matches(t1, t3), "Transformer 不同 heads 不匹配", all); - ModelSpec t4 = t1; t4.as()->norm_place = NormPlace::Final; - expect(!nn::spec_matches(t1, t4), "Transformer 不同 norm_place 不匹配", all); - ModelSpec t5 = t1; t5.as()->norm_type = NormType::BatchNorm; - expect(!nn::spec_matches(t1, t5), "Transformer 不同 norm_type 不匹配", all); + ModelSpec t3 = t1; + t3.as()->num_heads = 8; + expect(nn::spec_matches(t1, t2), "Transformer 相同 spec 匹配", all); + expect(!nn::spec_matches(t1, t3), "Transformer 不同 heads 不匹配", all); + ModelSpec t4 = t1; + t4.as()->norm_place = NormPlace::Final; + expect(!nn::spec_matches(t1, t4), "Transformer 不同 norm_place 不匹配", all); + ModelSpec t5 = t1; + t5.as()->norm_type = NormType::BatchNorm; + expect(!nn::spec_matches(t1, t5), "Transformer 不同 norm_type 不匹配", all); // CNN ModelSpec c1; { - auto &cc = c1.emplace(); - cc.cnn_in_channels = 1; cc.cnn_in_size = 28; cc.cnn_pool = 2; - cc.cnn_channels = {16, 32}; cc.cnn_kernels = {5, 5}; - cc.cnn_strides = {1, 1}; cc.cnn_paddings = {0, 0}; + auto& cc = c1.emplace(); + cc.cnn_in_channels = 1; + cc.cnn_in_size = 28; + cc.cnn_pool = 2; + cc.cnn_channels = {16, 32}; + cc.cnn_kernels = {5, 5}; + cc.cnn_strides = {1, 1}; + cc.cnn_paddings = {0, 0}; cc.layer_dims = {10}; } ModelSpec c2 = c1; - ModelSpec c3 = c1; c3.as()->cnn_channels = {16, 64}; - expect(nn::spec_matches(c1, c2), "CNN 相同 spec 匹配", all); - expect(!nn::spec_matches(c1, c3), "CNN 不同 channels 不匹配", all); - ModelSpec c4 = c1; c4.as()->norm_place = NormPlace::Both; - expect(!nn::spec_matches(c1, c4), "CNN 不同 norm_place 不匹配", all); - ModelSpec c5 = c1; c5.as()->norm_type = NormType::RMSNorm; - expect(!nn::spec_matches(c1, c5), "CNN 不同 norm_type 不匹配", all); + ModelSpec c3 = c1; + c3.as()->cnn_channels = {16, 64}; + expect(nn::spec_matches(c1, c2), "CNN 相同 spec 匹配", all); + expect(!nn::spec_matches(c1, c3), "CNN 不同 channels 不匹配", all); + ModelSpec c4 = c1; + c4.as()->norm_place = NormPlace::Both; + expect(!nn::spec_matches(c1, c4), "CNN 不同 norm_place 不匹配", all); + ModelSpec c5 = c1; + c5.as()->norm_type = NormType::RMSNorm; + expect(!nn::spec_matches(c1, c5), "CNN 不同 norm_type 不匹配", all); // ── 2. round-trip：匹配 spec 加载成功 ───────────────────────────── std::cout << "\n── round-trip（匹配 spec 加载成功） ──\n"; @@ -126,18 +144,23 @@ int run_test() const ModelSpec spec = nn::make_gpt_spec(100, 64, 32, 4, 128, 2); auto m = nn::build_gpt_model_from_spec(eng, spec); - if (!m) { expect(false, "build_gpt_model_from_spec: " + m.error().message, all); return all ? 0 : 1; } + if (!m) + { + expect(false, "build_gpt_model_from_spec: " + m.error().message, all); + return all ? 0 : 1; + } expect(m->spec().has_value(), "build_gpt_model_from_spec 记录 spec", all); if (auto r = nn::save_model(file, *m, spec); !r) - { expect(false, "save_model: " + r.error().message, all); return all ? 0 : 1; } + { + expect(false, "save_model: " + r.error().message, all); + return all ? 0 : 1; + } auto m2 = nn::build_gpt_model_from_spec(eng, spec); auto lr = nn::load_model(file, *m2); - if (lr) - expect(true, "匹配 spec 的 load_model 成功", all); - else - expect(false, "匹配 spec 应加载成功: " + lr.error().message, all); + if (lr) expect(true, "匹配 spec 的 load_model 成功", all); + else expect(false, "匹配 spec 应加载成功: " + lr.error().message, all); std::remove(file.c_str()); } @@ -150,18 +173,22 @@ int run_test() const ModelSpec spec = nn::make_gpt_spec(100, 64, 32, 4, 128, 2); auto m = nn::build_gpt_model_from_spec(eng, spec); - if (!m) { expect(false, "build_gpt_model_from_spec: " + m.error().message, all); return all ? 0 : 1; } + if (!m) + { + expect(false, "build_gpt_model_from_spec: " + m.error().message, all); + return all ? 0 : 1; + } if (auto r = nn::save_model(file, *m, spec); !r) - { expect(false, "save_model: " + r.error().message, all); return all ? 0 : 1; } + { + expect(false, "save_model: " + r.error().message, all); + return all ? 0 : 1; + } // 用不同 vocab 的架构构建，加载同一文件应被校验拦截 const ModelSpec bad = nn::make_gpt_spec(999, 64, 32, 4, 128, 2); auto m2 = nn::build_gpt_model_from_spec(eng, bad); auto lr = nn::load_model(file, *m2); - if (lr) - { - expect(false, "不匹配 spec 应加载失败（但成功了）", all); - } + if (lr) { expect(false, "不匹配 spec 应加载失败（但成功了）", all); } else { expect(true, "不匹配 spec 的 load_model 正确报错", all); @@ -179,19 +206,24 @@ int run_test() const ModelSpec spec = nn::make_gpt_spec(100, 64, 32, 4, 128, 2); auto m = nn::build_gpt_model_from_spec(eng, spec); - if (!m) { expect(false, "build_gpt_model_from_spec: " + m.error().message, all); return all ? 0 : 1; } + if (!m) + { + expect(false, "build_gpt_model_from_spec: " + m.error().message, all); + return all ? 0 : 1; + } if (auto r = nn::save_model(file, *m, spec); !r) - { expect(false, "save_model: " + r.error().message, all); return all ? 0 : 1; } + { + expect(false, "save_model: " + r.error().message, all); + return all ? 0 : 1; + } // 用与文件相同架构直接构建（无 spec，跳过架构校验），加载应成功 // （build_gpt_model 不记录 spec，load_model 不会触发 spec_matches 拦截）。 auto m2 = nn::build_gpt_model(eng, 100, 64, 32, 4, 128, 2); expect(!m2->spec().has_value(), "build_gpt_model 不记录 spec", all); auto lr = nn::load_model(file, *m2); - if (lr) - expect(true, "无 spec 模型 load_model 成功（跳过校验）", all); - else - expect(false, "无 spec 模型应加载成功: " + lr.error().message, all); + if (lr) expect(true, "无 spec 模型 load_model 成功（跳过校验）", all); + else expect(false, "无 spec 模型应加载成功: " + lr.error().message, all); std::remove(file.c_str()); } @@ -264,10 +296,7 @@ int run_test() nn::KeyValueRecord kv; kv.set("type", static_cast(6)); auto r = nn::detail::spec_from_kv(kv); - if (r) - { - expect(false, "type=6（ZiPT）应被拒绝（但解析成功了）", all); - } + if (r) { expect(false, "type=6（ZiPT）应被拒绝（但解析成功了）", all); } else { const bool readable = r.error().message.find("ZiPT") != std::string::npos; @@ -293,7 +322,7 @@ int run_test() kv.set("type", static_cast(nn::ModelType::CNN)); auto r = nn::detail::spec_from_kv(kv); expect(bool(r) && r->as()->norm_place == nn::NormPlace::None && - r->as()->norm_type == nn::NormType::LayerNorm, + r->as()->norm_type == nn::NormType::LayerNorm, "缺 norm_place/norm_type 键的旧文件 → None/LayerNorm（旧布局）", all); } @@ -312,15 +341,17 @@ int run_test() // 7.1 CNN：conv+head 挂 RMSNorm { nn::CnnConfig ccfg; - ccfg.in_channels = 1; ccfg.in_size = 12; ccfg.pool = 2; + ccfg.in_channels = 1; + ccfg.in_size = 12; + ccfg.pool = 2; ccfg.convs = {{4, 3, 1, 0}, {8, 3, 1, 0}}; ccfg.fc_dims = {16, 10}; ccfg.norm_type = NormType::RMSNorm; ccfg.norm_place = NormPlace::Both; - const nn::ModelSpec spec = nn::make_cnn_spec( - ccfg.in_channels, ccfg.in_size, ccfg.pool, ccfg.convs, - ccfg.fc_dims, ccfg.norm_type, ccfg.norm_place); + const nn::ModelSpec spec = + nn::make_cnn_spec(ccfg.in_channels, ccfg.in_size, ccfg.pool, ccfg.convs, + ccfg.fc_dims, ccfg.norm_type, ccfg.norm_place); auto m = nn::build_cnn_model_from_spec(eng, spec); if (!m) { expect(false, "build_cnn_model(norm=both): " + m.error().message, all); } else @@ -340,8 +371,7 @@ int run_test() else { auto m2 = nn::build_cnn_model_from_spec(eng, spec); - if (!m2) - expect(false, "CNN build_from_spec: " + m2.error().message, all); + if (!m2) expect(false, "CNN build_from_spec: " + m2.error().message, all); else { auto lr = nn::load_model(file, *m2); @@ -367,8 +397,7 @@ int run_test() { const nn::ModelSpec spec = nn::make_mnist_transformer_spec( nn::MNIST_PATCH_SIZE, nn::MNIST_TF_D_MODEL, nn::MNIST_TF_NUM_HEADS, - nn::MNIST_TF_D_FF, nn::MNIST_TF_NUM_LAYERS, - NormType::LayerNorm, NormPlace::Final); + nn::MNIST_TF_D_FF, nn::MNIST_TF_NUM_LAYERS, NormType::LayerNorm, NormPlace::Final); auto m = nn::build_mnist_model_from_spec(eng, spec); if (!m) { expect(false, "build ViT(final norm): " + m.error().message, all); } else @@ -390,7 +419,11 @@ int run_test() ok = static_cast(gm); if (ok) for (auto v : gm->span()) - if (!std::isfinite(v)) { ok = false; break; } + if (!std::isfinite(v)) + { + ok = false; + break; + } } } expect(ok, "ViT(final norm) forward/backward 梯度全有限", all); @@ -401,8 +434,7 @@ int run_test() else { auto m2 = nn::build_mnist_model_from_spec(eng, spec); - if (!m2) - expect(false, "ViT build_from_spec: " + m2.error().message, all); + if (!m2) expect(false, "ViT build_from_spec: " + m2.error().message, all); else { auto lr = nn::load_model(file, *m2); @@ -430,28 +462,19 @@ int run_test() struct MyNetSpec : nn::ModelSpecBase { std::size_t hidden = 0; - [[nodiscard]] const char *type_name() const noexcept override - { - return "MyNet"; - } + [[nodiscard]] const char* type_name() const noexcept override { return "MyNet"; } [[nodiscard]] std::unique_ptr clone() const override + { return std::make_unique(*this); } + [[nodiscard]] bool arch_compatible(const nn::ModelSpecBase& o) const override { - return std::make_unique(*this); - } - [[nodiscard]] bool arch_compatible(const nn::ModelSpecBase &o) const override - { - const auto *p = dynamic_cast(&o); + const auto* p = dynamic_cast(&o); return p && hidden == p->hidden; } - [[nodiscard]] std::string summary(const nn::ModelSpec &) const override - { - return "MyNet(hidden=" + std::to_string(hidden) + ")"; - } - void write_kv(nn::KeyValueRecord &kv) const override - { - kv.set("hidden", static_cast(hidden)); - } - [[nodiscard]] nn::Result read_kv(const nn::KeyValueRecord &kv) override + [[nodiscard]] std::string summary(const nn::ModelSpec&) const override + { return "MyNet(hidden=" + std::to_string(hidden) + ")"; } + void write_kv(nn::KeyValueRecord& kv) const override + { kv.set("hidden", static_cast(hidden)); } + [[nodiscard]] nn::Result read_kv(const nn::KeyValueRecord& kv) override { std::uint64_t v = 0; if (kv.get("hidden", v)) hidden = static_cast(v); @@ -459,8 +482,8 @@ int run_test() } }; - nn::ModelRegistry::instance().register_type( - "MyNet", [] { return std::make_unique(); }); + nn::ModelRegistry::instance().register_type("MyNet", + [] { return std::make_unique(); }); nn::ModelSpec mine; mine.emplace().hidden = 42; @@ -473,8 +496,7 @@ int run_test() "自定义类型 save → peek 往返（经 ModelRegistry 构造）", all); expect(back && nn::spec_matches(mine, *back), "自定义类型 spec_matches 走 arch_compatible 虚分发", all); - expect(back && nn::spec_summary(*back).find("MyNet(hidden=42)") != - std::string::npos, + expect(back && nn::spec_summary(*back).find("MyNet(hidden=42)") != std::string::npos, "自定义类型 spec_summary 走 summary 虚分发", all); // 未知类型名（未注册）→ 明确报错 nn::KeyValueRecord bad_kv; @@ -492,6 +514,4 @@ int run_test() } // namespace int main() -{ - return run_test(); -} +{ return run_test(); } diff --git a/src/optimizer_state_test.cpp b/src/optimizer_state_test.cpp index 462518fd..cf84025a 100644 --- a/src/optimizer_state_test.cpp +++ b/src/optimizer_state_test.cpp @@ -51,8 +51,10 @@ struct Fixture params.push_back(engine.create_tensor(3, 1, nn::Precision::F32, nn::InitSpec::zero())); grads.push_back(engine.create_tensor(4, 3, nn::Precision::F32, nn::InitSpec::zero())); grads.push_back(engine.create_tensor(3, 1, nn::Precision::F32, nn::InitSpec::zero())); - for (auto& p : params) pref.push_back(std::ref(p)); - for (auto& g : grads) gref.push_back(std::ref(g)); + for (auto& p : params) + pref.push_back(std::ref(p)); + for (auto& g : grads) + gref.push_back(std::ref(g)); set_params_from_seed(/*seed=*/1); set_grads(0); } @@ -65,8 +67,8 @@ struct Fixture nn::Matrix m(params[pi].rows(), params[pi].cols()); auto s = m.span(); for (std::size_t k = 0; k < s.size(); ++k) - s[k] = 0.05f * static_cast((static_cast(k) * 7 + seed * 3) % 11 - 5) - + 0.001f * static_cast(pi); + s[k] = 0.05f * static_cast((static_cast(k) * 7 + seed * 3) % 11 - 5) + + 0.001f * static_cast(pi); NN_ASSERT_ALWAYS(engine.copy_from(params[pi], m), "upload params"); } } @@ -90,7 +92,8 @@ struct Fixture [[nodiscard]] std::vector snapshot() { std::vector out; - for (const auto& p : params) out.push_back(download(engine, p)); + for (const auto& p : params) + out.push_back(download(engine, p)); return out; } @@ -102,8 +105,7 @@ struct Fixture } }; -[[nodiscard]] bool same_params(const std::vector& a, - const std::vector& b) +[[nodiscard]] bool same_params(const std::vector& a, const std::vector& b) { if (a.size() != b.size()) return false; for (std::size_t i = 0; i < a.size(); ++i) @@ -120,8 +122,8 @@ struct Fixture [[nodiscard]] std::unique_ptr make(Fixture& fx, const std::string& name) { return nn::create_optimizer(name, fx.engine, fx.pref, fx.gref, - /*lr=*/0.01f, /*weight_decay=*/0.01f, - nn::PrecisionProfile{}, 0.9f, 0.999f); + /*lr=*/0.01f, /*weight_decay=*/0.01f, nn::PrecisionProfile{}, 0.9f, + 0.999f); } // 表里声明的"逻辑状态缓冲数" → 期望的张量数（每参数） @@ -145,7 +147,8 @@ void test_state_presence() CHECK(opt != nullptr, ("create_optimizer: " + std::string(meta.name)).c_str()); if (!opt) continue; - for (std::size_t s = 0; s < 3; ++s) fx.step(*opt, s); + for (std::size_t s = 0; s < 3; ++s) + fx.step(*opt, s); std::vector state; auto sr = opt->save_state(state); @@ -153,15 +156,16 @@ void test_state_presence() const std::size_t want = expected_tensors_per_param(meta.name) * kParams; CHECK(state.size() == want, ("状态张量数 = " + std::to_string(want) + "（" + std::string(meta.name) + - "），实际 " + std::to_string(state.size())).c_str()); + "），实际 " + std::to_string(state.size())) + .c_str()); if (meta.state_bufs == 0 && meta.name != "adamwquant") CHECK(state.empty(), "无状态优化器（sgd）导出的状态为空"); else if (state.empty()) CHECK(false, "有状态优化器不得导出空状态（新优化器必须实现 save_state）"); // 步计数器：Adam 家族非零；无状态优化器恒 0 - const bool adam_family = meta.name == "adam" || meta.name == "adamw" || - meta.name == "adamwquant"; + const bool adam_family = + meta.name == "adam" || meta.name == "adamw" || meta.name == "adamwquant"; CHECK(opt->state_step() == (adam_family ? 3u : 0u), ("state_step: " + std::string(meta.name)).c_str()); @@ -173,7 +177,11 @@ void test_state_presence() { const auto m = download(fx.engine, state[i]); for (float v : m.span()) - if (v != 0.0f) { nonzero = true; break; } + if (v != 0.0f) + { + nonzero = true; + break; + } } CHECK(nonzero, ("重启后状态非零（" + std::string(meta.name) + "）").c_str()); } @@ -195,8 +203,13 @@ void test_resume_continuity() Fixture a; a.build(); auto opt_a = make(a, name); - if (!opt_a) { CHECK(false, "create_optimizer (A)"); continue; } - for (std::size_t s = 0; s < kTotal; ++s) a.step(*opt_a, s); + if (!opt_a) + { + CHECK(false, "create_optimizer (A)"); + continue; + } + for (std::size_t s = 0; s < kTotal; ++s) + a.step(*opt_a, s); const auto params_a = a.snapshot(); // B：跑到中途 → 存状态 → 新建优化器（模拟进程重启）→ 载入续跑 @@ -204,8 +217,13 @@ void test_resume_continuity() b.build(); b.set_params_from_seed(1); // 与 A 同一初值 auto opt_b = make(b, name); - if (!opt_b) { CHECK(false, "create_optimizer (B)"); continue; } - for (std::size_t s = 0; s < kSplit; ++s) b.step(*opt_b, s); + if (!opt_b) + { + CHECK(false, "create_optimizer (B)"); + continue; + } + for (std::size_t s = 0; s < kSplit; ++s) + b.step(*opt_b, s); std::vector saved; auto sr = opt_b->save_state(saved); CHECK(sr.has_value(), ("save_state（中断点）: " + name).c_str()); @@ -213,11 +231,16 @@ void test_resume_continuity() opt_b.reset(); // "进程死了"：优化器连同状态一起消失 auto opt_b2 = make(b, name); // 重启后重新构造（状态为零初始化） - if (!opt_b2) { CHECK(false, "create_optimizer (B2)"); continue; } + if (!opt_b2) + { + CHECK(false, "create_optimizer (B2)"); + continue; + } opt_b2->set_state_step(saved_step); auto lr2 = opt_b2->load_state(saved); CHECK(lr2.has_value(), ("load_state（重启）: " + name).c_str()); - for (std::size_t s = kSplit; s < kTotal; ++s) b.step(*opt_b2, s); + for (std::size_t s = kSplit; s < kTotal; ++s) + b.step(*opt_b2, s); CHECK(same_params(params_a, b.snapshot()), ("重启续跑与直跑参数逐位一致（" + name + "）").c_str()); @@ -228,12 +251,22 @@ void test_resume_continuity() c.build(); c.set_params_from_seed(1); auto opt_c = make(c, name); - if (!opt_c) { CHECK(false, "create_optimizer (C)"); continue; } - for (std::size_t s = 0; s < kSplit; ++s) c.step(*opt_c, s); + if (!opt_c) + { + CHECK(false, "create_optimizer (C)"); + continue; + } + for (std::size_t s = 0; s < kSplit; ++s) + c.step(*opt_c, s); opt_c.reset(); auto opt_c2 = make(c, name); // 不载入状态、不恢复 t - if (!opt_c2) { CHECK(false, "create_optimizer (C2)"); continue; } - for (std::size_t s = kSplit; s < kTotal; ++s) c.step(*opt_c2, s); + if (!opt_c2) + { + CHECK(false, "create_optimizer (C2)"); + continue; + } + for (std::size_t s = kSplit; s < kTotal; ++s) + c.step(*opt_c2, s); if (stateful) CHECK(!same_params(params_a, c.snapshot()), ("冷启动续跑与直跑不同（负例具判别力：" + name + "）").c_str()); @@ -250,8 +283,13 @@ void test_state_mismatch() Fixture fx; fx.build(); auto opt = make(fx, "adam"); - if (!opt) { CHECK(false, "create_optimizer"); return; } - for (std::size_t s = 0; s < 2; ++s) fx.step(*opt, s); + if (!opt) + { + CHECK(false, "create_optimizer"); + return; + } + for (std::size_t s = 0; s < 2; ++s) + fx.step(*opt, s); std::vector state; NN_ASSERT_ALWAYS(opt->save_state(state), "save_state"); @@ -266,8 +304,7 @@ void test_state_mismatch() CHECK(!bad_shape.has_value(), "状态张量形状不符 → 报错"); // sgd 无状态：给非空状态 → 报错 auto sgd = make(fx, "sgd"); - if (sgd) - CHECK(!sgd->load_state(state).has_value(), "无状态优化器收到状态 → 报错"); + if (sgd) CHECK(!sgd->load_state(state).has_value(), "无状态优化器收到状态 → 报错"); } } // namespace diff --git a/src/optimizer_test.cpp b/src/optimizer_test.cpp index e015eca1..98370e0b 100644 --- a/src/optimizer_test.cpp +++ b/src/optimizer_test.cpp @@ -37,17 +37,19 @@ namespace { int g_fail = 0; -#define CHECK(cond, msg) \ - do { \ - if (!(cond)) { \ - std::printf("[FAIL] %s:%d: %s\n", __FILE__, __LINE__, (msg)); \ - ++g_fail; \ - } \ +#define CHECK(cond, msg) \ + do \ + { \ + if (!(cond)) \ + { \ + std::printf("[FAIL] %s:%d: %s\n", __FILE__, __LINE__, (msg)); \ + ++g_fail; \ + } \ } while (0) // ── 引擎侧工具（M1：Tensor 一律经引擎创建/读写）──────────────────────────── -nn::Tensor make_tensor(nn::CpuEngine& eng, std::size_t rows, std::size_t cols, - float base, float step) +nn::Tensor make_tensor(nn::CpuEngine& eng, std::size_t rows, std::size_t cols, float base, + float step) { nn::Tensor t = eng.create_tensor(rows, cols); std::vector buf(rows * cols); @@ -123,8 +125,7 @@ void test_optimizer_table() // ① 表内每个 name 都能被工厂构造（名单与构造器不得漂移） nn::Tensor p = make_tensor(eng, 2, 3, 0.1f, 0.01f); nn::Tensor g = make_tensor(eng, 2, 3, 0.2f, 0.01f); - auto opt = nn::create_optimizer(m.name, eng, - std::vector{p}, + auto opt = nn::create_optimizer(m.name, eng, std::vector{p}, std::vector{g}, /*lr=*/Scalar{1e-3f}, /*wd=*/Scalar{0.01f}); CHECK(opt != nullptr, (std::string("工厂可构造: ") + std::string(m.name)).c_str()); @@ -150,9 +151,8 @@ void test_optimizer_table() // ── AdamW 单步参考（与 compute_optimizer.hpp::AdamW::step 同一公式）──────── // decay_first = true → 先衰减后更新（AdamW 语义，正确） // decay_first = false → 先更新后衰减（错误顺序，仅用于证明顺序敏感） -void adamw_reference(std::vector& p, const std::vector& g, - std::vector& m, std::vector& v, - Scalar lr, Scalar b1, Scalar b2, Scalar eps, Scalar wd, +void adamw_reference(std::vector& p, const std::vector& g, std::vector& m, + std::vector& v, Scalar lr, Scalar b1, Scalar b2, Scalar eps, Scalar wd, std::size_t t, bool decay_first) { const Scalar bc1 = Scalar{1} - std::pow(b1, static_cast(t)); @@ -202,8 +202,8 @@ void test_adamw_order() const Scalar denom = std::max(std::fabs(ref[i]), Scalar{1e-6f}); max_rel = std::max(max_rel, d / denom); } - std::printf(" AdamW 单步 vs 参考: max_abs=%.3e max_rel=%.3e\n", - static_cast(max_abs), static_cast(max_rel)); + std::printf(" AdamW 单步 vs 参考: max_abs=%.3e max_rel=%.3e\n", static_cast(max_abs), + static_cast(max_rel)); CHECK(max_rel < 1e-5f, "AdamW 单步与参考实现一致（wd≠0，先衰减后更新）"); // 负例：同样的超参、只把两次写的顺序对调 → 结果必须明显不同 @@ -214,8 +214,7 @@ void test_adamw_order() Scalar swap_diff = 0; for (std::size_t i = 0; i < ref.size(); ++i) swap_diff = std::max(swap_diff, std::fabs(ref[i] - ref_swapped[i])); - std::printf(" 顺序对调后的差值: %.3e（应当显著非零）\n", - static_cast(swap_diff)); + std::printf(" 顺序对调后的差值: %.3e（应当显著非零）\n", static_cast(swap_diff)); CHECK(swap_diff > 1e-6f, "decay→update 的顺序确实影响数值（顺序边非形式主义）"); std::printf("[OK] test_adamw_order\n"); } @@ -231,13 +230,12 @@ void test_dst_order_hazard() nn::ExprGraph graph; // 节点 0：p *= 2（写目标 = p） - const int n0 = graph.add_node(make_unary_const(nn::ExprOp::Mul, Scalar{2}), - std::vector{p}, rows, cols, false, - nn::Precision::F32, p.storage_id()); + const int n0 = + graph.add_node(make_unary_const(nn::ExprOp::Mul, Scalar{2}), std::vector{p}, + rows, cols, false, nn::Precision::F32, p.storage_id()); // 节点 1：p = p + x（读 p、写目标同样 = p） - const int n1 = graph.add_node(make_binary(nn::ExprOp::Add), - std::vector{p, x}, rows, cols, false, - nn::Precision::F32, p.storage_id()); + const int n1 = graph.add_node(make_binary(nn::ExprOp::Add), std::vector{p, x}, rows, + cols, false, nn::Precision::F32, p.storage_id()); // ① 顺序隐患边 (n0, n1) const auto edges = nn::dst_order_constraints(graph); @@ -253,7 +251,11 @@ void test_dst_order_hazard() for (const auto& k : kernels) { bool a = false, b = false; - for (const int m : k.members) { a = a || (m == n0); b = b || (m == n1); } + for (const int m : k.members) + { + a = a || (m == n0); + b = b || (m == n1); + } if (a && b) same_kernel = true; } CHECK(!same_kernel, "两节点未出现在同一 kernel 的成员表里"); diff --git a/src/precision_profile_test.cpp b/src/precision_profile_test.cpp index 0ed43107..8e001221 100644 --- a/src/precision_profile_test.cpp +++ b/src/precision_profile_test.cpp @@ -28,9 +28,7 @@ int g_failures = 0; bool is_all(nn::PrecisionProfile p, nn::Precision v) -{ - return p.param == v && p.compute == v && p.stable == v && p.optimizer == v; -} +{ return p.param == v && p.compute == v && p.stable == v && p.optimizer == v; } } // namespace int main() @@ -55,7 +53,7 @@ int main() const PrecisionProfile mw = nn::profile_master_weights(); CHECK(mw.param == Precision::F32 && mw.compute == Precision::F16 && - mw.stable == Precision::F32 && mw.optimizer == Precision::F32, + mw.stable == Precision::F32 && mw.optimizer == Precision::F32, "profile_master_weights 应为 F32/F16/F32/F32（混合，参数仍 f32）"); // ── CLI `--t1_58` 语义：param + 其余三槽照抄 --f16 ─────────────────── @@ -73,8 +71,8 @@ int main() // 各配方两两可区分：锁定 --f16 = f16 存储配方，不被改成混合/全 f16 const auto same = [](const PrecisionProfile& a, const PrecisionProfile& b) { - return a.param == b.param && a.compute == b.compute && - a.stable == b.stable && a.optimizer == b.optimizer; + return a.param == b.param && a.compute == b.compute && a.stable == b.stable && + a.optimizer == b.optimizer; }; CHECK(!same(f16, mw), "profile_f16 与 profile_master_weights 必须可区分"); CHECK(!same(f16, all16), "profile_f16 与 profile_all_f16 必须可区分"); diff --git a/src/precision_test.cpp b/src/precision_test.cpp index 868d3566..e0b67400 100644 --- a/src/precision_test.cpp +++ b/src/precision_test.cpp @@ -39,12 +39,14 @@ // 上溢边界、±inf、NaN（含负载）、以及一半值的往返。 namespace { -constexpr std::uint16_t soft_h(float v) { return nn::detail::float_to_half_bits(v); } -constexpr float soft_f(std::uint16_t h) { return nn::detail::half_bits_to_float(h); } +constexpr std::uint16_t soft_h(float v) +{ return nn::detail::float_to_half_bits(v); } +constexpr float soft_f(std::uint16_t h) +{ return nn::detail::half_bits_to_float(h); } struct HalfCase { - float v; + float v; std::uint16_t soft; // 编译期（软件路径）结果 }; constexpr HalfCase kHalfCases[] = { @@ -64,10 +66,8 @@ constexpr HalfCase kHalfCases[] = { {65536.0f, soft_h(65536.0f)}, // 2^16 → inf {1e30f, soft_h(1e30f)}, {-1e30f, soft_h(-1e30f)}, - {std::numeric_limits::infinity(), - soft_h(std::numeric_limits::infinity())}, - {-std::numeric_limits::infinity(), - soft_h(-std::numeric_limits::infinity())}, + {std::numeric_limits::infinity(), soft_h(std::numeric_limits::infinity())}, + {-std::numeric_limits::infinity(), soft_h(-std::numeric_limits::infinity())}, {std::bit_cast(0x7FC00000u), soft_h(std::bit_cast(0x7FC00000u))}, {std::bit_cast(0xFFC00000u), soft_h(std::bit_cast(0xFFC00000u))}, {std::bit_cast(0x7F800001u), soft_h(std::bit_cast(0x7F800001u))}, @@ -76,9 +76,9 @@ constexpr HalfCase kHalfCases[] = { }; // f16 → f32 侧：软件路径恒精确，硬件同；用一组位型对比 -constexpr std::uint16_t kHalfBits[] = { - 0x0000u, 0x8000u, 0x0001u, 0x03FFu, 0x0400u, 0x3555u, 0x3C00u, 0x3C01u, - 0x7BFFu, 0x7C00u, 0xFC00u, 0x7E00u, 0xFE00u, 0x7C01u, 0xFFFFu}; +constexpr std::uint16_t kHalfBits[] = {0x0000u, 0x8000u, 0x0001u, 0x03FFu, 0x0400u, + 0x3555u, 0x3C00u, 0x3C01u, 0x7BFFu, 0x7C00u, + 0xFC00u, 0x7E00u, 0xFE00u, 0x7C01u, 0xFFFFu}; } // namespace namespace @@ -100,16 +100,14 @@ int test_f16_fastpath_parity() for (const std::uint16_t h : kHalfBits) { const float hw = nn::detail::half_bits_to_float(h); // 运行期 - if (std::bit_cast(hw) != - std::bit_cast(soft_f(h))) + if (std::bit_cast(hw) != std::bit_cast(soft_f(h))) { std::printf(" [FAIL] f16→f32 位型不一致: h=0x%04x\n", h); ++bad; } } std::printf(" %s（%zu 条 f32→f16 + %zu 条 f16→f32，逐位一致）\n", - bad == 0 ? "[PASS]" : "[FAIL]", - sizeof(kHalfCases) / sizeof(kHalfCases[0]), + bad == 0 ? "[PASS]" : "[FAIL]", sizeof(kHalfCases) / sizeof(kHalfCases[0]), sizeof(kHalfBits) / sizeof(kHalfBits[0])); return bad; } diff --git a/src/rapt_checkpoint_test.cpp b/src/rapt_checkpoint_test.cpp index 32522e14..93b926a0 100644 --- a/src/rapt_checkpoint_test.cpp +++ b/src/rapt_checkpoint_test.cpp @@ -40,12 +40,12 @@ namespace int ckpt_run_test() { - const std::size_t vocab = 257; - const std::size_t d_model = 64; - const std::size_t seq_len = 16; - const std::size_t batch = 2; + const std::size_t vocab = 257; + const std::size_t d_model = 64; + const std::size_t seq_len = 16; + const std::size_t batch = 2; const std::size_t num_heads = 4; - const std::size_t d_ff = 128; + const std::size_t d_ff = 128; const std::size_t num_layers = 6; auto eng_owner = nn::CpuEngine::create(); @@ -54,16 +54,13 @@ int ckpt_run_test() std::cout << "========================================\n"; std::cout << " RAPT 梯度检查点一致性测试 (L1)\n"; std::cout << "========================================\n"; - std::cout << " vocab=" << vocab << " d_model=" << d_model - << " seq=" << seq_len << " batch=" << batch - << " heads=" << num_heads << " d_ff=" << d_ff + std::cout << " vocab=" << vocab << " d_model=" << d_model << " seq=" << seq_len + << " batch=" << batch << " heads=" << num_heads << " d_ff=" << d_ff << " layers=" << num_layers << "\n"; // RLA 约束：位置编码必须 RoPE - RAPTModel model(vocab, d_model, seq_len, num_heads, d_ff, num_layers, - PosEncodingType::RoPE, - ActivationType::SwiGLU, - NormType::RMSNorm); + RAPTModel model(vocab, d_model, seq_len, num_heads, d_ff, num_layers, PosEncodingType::RoPE, + ActivationType::SwiGLU, NormType::RMSNorm); NN_EXIT(model.init(eng), 1, "RAPTModel init 失败: "); std::mt19937_64 rng(7); @@ -76,7 +73,8 @@ int ckpt_run_test() // 固定输出梯度 (vocab, seq*batch) nn::Matrix go_m(vocab, seq_len * batch); - for (std::size_t i = 0; i < go_m.size(); ++i) go_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); NN_EXIT(x, 1, "from_matrix(x) failed\n"); @@ -154,6 +152,4 @@ int ckpt_run_test() } // namespace int main() -{ - return ckpt_run_test(); -} +{ return ckpt_run_test(); } diff --git a/src/rapt_gradcheck.cpp b/src/rapt_gradcheck.cpp index f7019d50..b890e86b 100644 --- a/src/rapt_gradcheck.cpp +++ b/src/rapt_gradcheck.cpp @@ -22,21 +22,20 @@ #include #include "test_common.hpp" -using nn::Scalar; +using nn::ComputeEngine; using nn::Matrix; +using nn::ReLULinearAttention; +using nn::Scalar; using nn::Tensor; using nn::TensorRef; -using nn::ComputeEngine; -using nn::ReLULinearAttention; - -namespace { +namespace +{ -bool check_param(ComputeEngine& engine, - const std::function& eval_loss, - Tensor& param, const Matrix& base, const Matrix& ana, - const std::string& name, Scalar eps, Scalar tol) +bool check_param(ComputeEngine& engine, const std::function& eval_loss, Tensor& param, + const Matrix& base, const Matrix& ana, const std::string& name, Scalar eps, + Scalar tol) { bool ok = true; std::size_t bad = 0, skipped = 0; @@ -51,10 +50,12 @@ bool check_param(ComputeEngine& engine, const Scalar orig = base.at_unchecked(r, c); if (!engine.write(param, base.span())) return false; const Scalar f0 = eval_loss(); - Matrix pp = base; pp.set_value_unchecked(r, c, orig + eps); + Matrix pp = base; + pp.set_value_unchecked(r, c, orig + eps); if (!engine.write(param, pp.span())) return false; const Scalar fp = eval_loss(); - Matrix pm = base; pm.set_value_unchecked(r, c, orig - eps); + Matrix pm = base; + pm.set_value_unchecked(r, c, orig - eps); if (!engine.write(param, pm.span())) return false; const Scalar fm = eval_loss(); if (!engine.write(param, base.span())) return false; @@ -74,32 +75,36 @@ bool check_param(ComputeEngine& engine, { ok = false; if (bad < 5) - std::cout << " [FAIL] " << name << "(" << r << "," << c - << ") num=" << num << " ana=" << ana_v - << " err=" << err << "\n"; + std::cout << " [FAIL] " << name << "(" << r << "," << c << ") num=" << num + << " ana=" << ana_v << " err=" << err << "\n"; ++bad; } } } - std::cout << " " << name << ": " << (ok ? "OK" : "FAIL") - << " max_err=" << max_err << " bad=" << bad - << " skipped(kink)=" << skipped << "\n"; + std::cout << " " << name << ": " << (ok ? "OK" : "FAIL") << " max_err=" << max_err + << " bad=" << bad << " skipped(kink)=" << skipped << "\n"; return ok; } -bool check_all_params(ComputeEngine& engine, - const std::function& eval_loss, - std::vector& params, - std::vector& grads, - Scalar eps, Scalar tol) +bool check_all_params(ComputeEngine& engine, const std::function& eval_loss, + std::vector& params, std::vector& grads, Scalar eps, + Scalar tol) { bool ok = true; for (std::size_t i = 0; i < params.size(); ++i) { auto base = engine.to_matrix(params[i]); - if (!base) { std::cerr << " to_matrix failed\n"; return false; } + if (!base) + { + std::cerr << " to_matrix failed\n"; + return false; + } auto ana = engine.to_matrix(grads[i]); - if (!ana) { std::cerr << " to_matrix(grad) failed\n"; return false; } + if (!ana) + { + std::cerr << " to_matrix(grad) failed\n"; + return false; + } ok &= check_param(engine, eval_loss, params[i], *base, *ana, "param[" + std::to_string(i) + "]", eps, tol); } @@ -108,8 +113,7 @@ bool check_all_params(ComputeEngine& engine, // 用固定种子重设全部参数，使 gradcheck 跨运行可复现 // （ReLU 层拐点附近的有限差分梯度不稳定，固定权重避免随机初始化抖动） -bool reseed_params(ComputeEngine& engine, std::vector& params, - unsigned seed) +bool reseed_params(ComputeEngine& engine, std::vector& params, unsigned seed) { std::mt19937_64 rng{seed}; std::normal_distribution dist(0.0, 0.5); @@ -118,7 +122,8 @@ bool reseed_params(ComputeEngine& engine, std::vector& params, auto m = engine.to_matrix(p); if (!m) return false; auto sp = m->span(); - for (auto& v : sp) v = dist(rng); + for (auto& v : sp) + v = dist(rng); if (!engine.write(p, m->span())) return false; } return true; @@ -148,30 +153,48 @@ int main(int argc, char* argv[]) { std::cout << "=== ReLULinearAttention (causal) gradcheck ===\n"; const std::size_t d_model = 8, heads = 2, seq = 4, batch = 2; - ReLULinearAttention attn(d_model, heads, seq, /*causal=*/true, - nn::PosEncodingType::RoPE); - if (!attn.init(eng)) { std::cerr << " init failed\n"; return 1; } - { auto prm = attn.parameters(); if (!reseed_params(eng, prm, 101)) return 1; } + ReLULinearAttention attn(d_model, heads, seq, /*causal=*/true, nn::PosEncodingType::RoPE); + if (!attn.init(eng)) + { + std::cerr << " init failed\n"; + return 1; + } + { + auto prm = attn.parameters(); + if (!reseed_params(eng, prm, 101)) return 1; + } std::mt19937_64 rng{7}; std::normal_distribution dist(0.0, 1.0); Matrix x_m(d_model, batch * seq); - { auto sp = x_m.span(); for (auto& v : sp) v = dist(rng); } + { + auto sp = x_m.span(); + for (auto& v : sp) + v = dist(rng); + } auto x_t = eng.from_matrix(x_m); NN_EXIT(x_t, 1); Matrix go_m(d_model, batch * seq); - { auto sp = go_m.span(); for (auto& v : sp) v = dist(rng); } + { + auto sp = go_m.span(); + for (auto& v : sp) + v = dist(rng); + } auto go_t = eng.from_matrix(go_m); NN_EXIT(go_t, 1); - auto eval_loss = [&]() -> Scalar { + auto eval_loss = [&]() -> Scalar + { auto y = attn.forward(*x_t); auto ym = eng.to_matrix(*y); return dot(*ym, go_m); }; auto grads = attn.param_gradients(); - for (auto& g : grads) { if (!eng.zero(g)) return 1; } + for (auto& g : grads) + { + if (!eng.zero(g)) return 1; + } auto y0 = attn.forward(*x_t); NN_EXIT(y0, 1, " fwd failed\n"); auto gx = attn.backward(*go_t); @@ -185,30 +208,48 @@ int main(int argc, char* argv[]) { std::cout << "=== ReLULinearAttention (bidirectional) gradcheck ===\n"; const std::size_t d_model = 8, heads = 2, seq = 4, batch = 1; - ReLULinearAttention attn(d_model, heads, seq, /*causal=*/false, - nn::PosEncodingType::RoPE); - if (!attn.init(eng)) { std::cerr << " init failed\n"; return 1; } - { auto prm = attn.parameters(); if (!reseed_params(eng, prm, 202)) return 1; } + ReLULinearAttention attn(d_model, heads, seq, /*causal=*/false, nn::PosEncodingType::RoPE); + if (!attn.init(eng)) + { + std::cerr << " init failed\n"; + return 1; + } + { + auto prm = attn.parameters(); + if (!reseed_params(eng, prm, 202)) return 1; + } std::mt19937_64 rng{9}; std::normal_distribution dist(0.0, 1.0); Matrix x_m(d_model, batch * seq); - { auto sp = x_m.span(); for (auto& v : sp) v = dist(rng); } + { + auto sp = x_m.span(); + for (auto& v : sp) + v = dist(rng); + } auto x_t = eng.from_matrix(x_m); NN_EXIT(x_t, 1); Matrix go_m(d_model, batch * seq); - { auto sp = go_m.span(); for (auto& v : sp) v = dist(rng); } + { + auto sp = go_m.span(); + for (auto& v : sp) + v = dist(rng); + } auto go_t = eng.from_matrix(go_m); NN_EXIT(go_t, 1); - auto eval_loss = [&]() -> Scalar { + auto eval_loss = [&]() -> Scalar + { auto y = attn.forward(*x_t); auto ym = eng.to_matrix(*y); return dot(*ym, go_m); }; auto grads = attn.param_gradients(); - for (auto& g : grads) { if (!eng.zero(g)) return 1; } + for (auto& g : grads) + { + if (!eng.zero(g)) return 1; + } auto y0 = attn.forward(*x_t); NN_EXIT(y0, 1, " fwd failed\n"); auto gx = attn.backward(*go_t); @@ -223,37 +264,61 @@ int main(int argc, char* argv[]) { std::cout << "=== ReLULinearAttention (causal + doc-aware) gradcheck ===\n"; const std::size_t d_model = 8, heads = 2, seq = 4, batch = 2; - ReLULinearAttention attn(d_model, heads, seq, /*causal=*/true, - nn::PosEncodingType::RoPE); - if (!attn.init(eng)) { std::cerr << " init failed\n"; return 1; } - { auto prm = attn.parameters(); if (!reseed_params(eng, prm, 404)) return 1; } + ReLULinearAttention attn(d_model, heads, seq, /*causal=*/true, nn::PosEncodingType::RoPE); + if (!attn.init(eng)) + { + std::cerr << " init failed\n"; + return 1; + } + { + auto prm = attn.parameters(); + if (!reseed_params(eng, prm, 404)) return 1; + } // 文档边界（batch-major）：batch0=[0,0,1,1]，batch1=[2,2,2,3] // → 边界位置：batch0 t0,t2；batch1 t0,t3 std::vector doc_ids(batch * seq); - doc_ids[0]=0; doc_ids[1]=0; doc_ids[2]=1; doc_ids[3]=1; - doc_ids[4]=2; doc_ids[5]=2; doc_ids[6]=2; doc_ids[7]=3; + doc_ids[0] = 0; + doc_ids[1] = 0; + doc_ids[2] = 1; + doc_ids[3] = 1; + doc_ids[4] = 2; + doc_ids[5] = 2; + doc_ids[6] = 2; + doc_ids[7] = 3; attn.set_doc_ids(doc_ids); std::mt19937_64 rng{13}; std::normal_distribution dist(0.0, 1.0); Matrix x_m(d_model, batch * seq); - { auto sp = x_m.span(); for (auto& v : sp) v = dist(rng); } + { + auto sp = x_m.span(); + for (auto& v : sp) + v = dist(rng); + } auto x_t = eng.from_matrix(x_m); NN_EXIT(x_t, 1); Matrix go_m(d_model, batch * seq); - { auto sp = go_m.span(); for (auto& v : sp) v = dist(rng); } + { + auto sp = go_m.span(); + for (auto& v : sp) + v = dist(rng); + } auto go_t = eng.from_matrix(go_m); NN_EXIT(go_t, 1); - auto eval_loss = [&]() -> Scalar { + auto eval_loss = [&]() -> Scalar + { auto y = attn.forward(*x_t); auto ym = eng.to_matrix(*y); return dot(*ym, go_m); }; auto grads = attn.param_gradients(); - for (auto& g : grads) { if (!eng.zero(g)) return 1; } + for (auto& g : grads) + { + if (!eng.zero(g)) return 1; + } auto y0 = attn.forward(*x_t); NN_EXIT(y0, 1, " fwd failed\n"); auto gx = attn.backward(*go_t); @@ -263,7 +328,6 @@ int main(int argc, char* argv[]) all_ok &= check_all_params(eng, eval_loss, params, grads, eps, tol); } - std::cout << (all_ok ? "\nALL RAPT GRADCHECKS PASSED\n" - : "\nSOME RAPT GRADCHECKS FAILED\n"); + std::cout << (all_ok ? "\nALL RAPT GRADCHECKS PASSED\n" : "\nSOME RAPT GRADCHECKS FAILED\n"); return all_ok ? 0 : 1; } \ No newline at end of file diff --git a/src/rapt_offload_test.cpp b/src/rapt_offload_test.cpp index 8853fc5d..0341d089 100644 --- a/src/rapt_offload_test.cpp +++ b/src/rapt_offload_test.cpp @@ -52,12 +52,12 @@ namespace int off_run_test() { - const std::size_t vocab = 257; - const std::size_t d_model = 64; - const std::size_t seq_len = 16; - const std::size_t batch = 2; + const std::size_t vocab = 257; + const std::size_t d_model = 64; + const std::size_t seq_len = 16; + const std::size_t batch = 2; const std::size_t num_heads = 4; - const std::size_t d_ff = 128; + const std::size_t d_ff = 128; const std::size_t num_layers = 6; auto backend_owner = GpuBackend::shared(); @@ -75,15 +75,12 @@ int off_run_test() std::cout << "========================================\n" << " RAPT activation offload 一致性测试\n" << "========================================\n"; - std::cout << " vocab=" << vocab << " d_model=" << d_model - << " seq=" << seq_len << " batch=" << batch - << " heads=" << num_heads << " d_ff=" << d_ff + std::cout << " vocab=" << vocab << " d_model=" << d_model << " seq=" << seq_len + << " batch=" << batch << " heads=" << num_heads << " d_ff=" << d_ff << " layers=" << num_layers << "\n"; - RAPTModel model(vocab, d_model, seq_len, num_heads, d_ff, num_layers, - PosEncodingType::RoPE, - ActivationType::SwiGLU, - NormType::RMSNorm); + RAPTModel model(vocab, d_model, seq_len, num_heads, d_ff, num_layers, PosEncodingType::RoPE, + ActivationType::SwiGLU, NormType::RMSNorm); NN_EXIT(model.init(eng), 1, "RAPTModel init 失败: "); std::mt19937_64 rng(7); @@ -93,7 +90,8 @@ int off_run_test() for (std::size_t i = 0; i < x_m.size(); ++i) x_m.span()[i] = static_cast(rng() % vocab); Matrix go_m(vocab, seq_len * batch); - for (std::size_t i = 0; i < go_m.size(); ++i) go_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); NN_EXIT(x, 1, "from_matrix(x) failed\n"); @@ -163,7 +161,5 @@ int off_run_test() } // namespace int main() -{ - return off_run_test(); -} +{ return off_run_test(); } #endif \ No newline at end of file diff --git a/src/rapt_smoke_test.cpp b/src/rapt_smoke_test.cpp index 3dbcabd9..b34bbd05 100644 --- a/src/rapt_smoke_test.cpp +++ b/src/rapt_smoke_test.cpp @@ -38,10 +38,10 @@ int main(int argc, char* argv[]) // ── 1. 规格往返校验 ───────────────────────────────────────────── nn::RAPTConfig cfg; cfg.vocab_size = 32; - cfg.d_model = 16; // d_k = d_model/heads = 8（偶数，RoPE 约束） - cfg.seq_len = 8; - cfg.num_heads = 2; - cfg.d_ff = 32; + cfg.d_model = 16; // d_k = d_model/heads = 8（偶数，RoPE 约束） + cfg.seq_len = 8; + cfg.num_heads = 2; + cfg.d_ff = 32; cfg.num_layers = 2; // P2-3：`--f16` 走 profile_f16()（param/compute=F16，stable/optimizer=F32） if (use_f16) @@ -50,8 +50,8 @@ int main(int argc, char* argv[]) std::cout << "精度: --f16（param/compute=f16）\n"; } - auto spec = nn::make_rapt_spec(cfg.vocab_size, cfg.d_model, cfg.seq_len, - cfg.num_heads, cfg.d_ff, cfg.num_layers); + auto spec = nn::make_rapt_spec(cfg.vocab_size, cfg.d_model, cfg.seq_len, cfg.num_heads, + cfg.d_ff, cfg.num_layers); auto model_from_spec = nn::build_rapt_model_from_spec(eng, spec); NN_EXIT(model_from_spec, 1, "build_rapt_model_from_spec failed\n"); if (!nn::spec_matches(spec, *model_from_spec->spec())) @@ -65,14 +65,13 @@ int main(int argc, char* argv[]) NN_EXIT(model, 1, "build_rapt_model failed\n"); // ── 2. 训练冒烟：合成数据（causal 下一位置预测），验证 loss 下降 ── - auto optimizer = nn::create_optimizer( - "adamw", eng, model->parameters(), model->param_gradients(), - nn::Scalar{1e-3}, nn::Scalar{0.01}); + auto optimizer = + nn::create_optimizer("adamw", eng, model->parameters(), model->param_gradients(), + nn::Scalar{1e-3}, nn::Scalar{0.01}); NN_EXIT(optimizer, 1, "optimizer creation failed\n"); nn::CrossEntropyLoss ce; - if (use_f16) - ce.set_precision_profile(nn::profile_f16()); + if (use_f16) ce.set_precision_profile(nn::profile_f16()); std::mt19937_64 rng{123}; std::uniform_int_distribution tok(0, cfg.vocab_size - 1); @@ -80,7 +79,8 @@ int main(int argc, char* argv[]) nn::Matrix in(cfg.seq_len, 1); std::vector labels(cfg.seq_len); std::vector toks(cfg.seq_len); - for (std::size_t t = 0; t < cfg.seq_len; ++t) toks[t] = tok(rng); + for (std::size_t t = 0; t < cfg.seq_len; ++t) + toks[t] = tok(rng); for (std::size_t t = 0; t < cfg.seq_len; ++t) { in.set_value_unchecked(t, 0, static_cast(toks[t])); @@ -90,7 +90,7 @@ int main(int argc, char* argv[]) NN_EXIT(x, 1); nn::Scalar first_loss = 0; - nn::Scalar last_loss = 0; + nn::Scalar last_loss = 0; for (std::size_t step = 0; step < steps; ++step) { auto zero = optimizer->zero_grad(); @@ -126,12 +126,12 @@ int main(int argc, char* argv[]) for (std::size_t b = 0; b < batch; ++b) { std::vector bt(cfg.seq_len); - for (std::size_t t = 0; t < cfg.seq_len; ++t) bt[t] = tok(rng); + for (std::size_t t = 0; t < cfg.seq_len; ++t) + bt[t] = tok(rng); for (std::size_t t = 0; t < cfg.seq_len; ++t) { bin.set_value_unchecked(t, b, static_cast(bt[t])); - blabels[b * cfg.seq_len + t] = - (t + 1 < cfg.seq_len) ? bt[t + 1] : tok(rng); + blabels[b * cfg.seq_len + t] = (t + 1 < cfg.seq_len) ? bt[t + 1] : tok(rng); } } auto bx = eng.from_matrix(bin); @@ -145,7 +145,11 @@ int main(int argc, char* argv[]) NN_EXIT(logits, 1, "batch>1 forward failed: "); auto loss = ce.forward_sparse(eng, *logits, blabels, {}, cfg.vocab_size); NN_EXIT(loss, 1, "batch>1 loss failed\n"); - if (!(*loss == *loss)) { std::cerr << "batch>1 loss NaN\n"; return 1; } + if (!(*loss == *loss)) + { + std::cerr << "batch>1 loss NaN\n"; + return 1; + } auto grad = ce.backward(); NN_EXIT(grad, 1, "batch>1 loss backward failed\n"); auto b = model->backward(*grad); @@ -162,11 +166,19 @@ int main(int argc, char* argv[]) const std::size_t d_model = 16, heads = 2, L = 6; // d_k = 8 nn::ReLULinearAttention attn(d_model, heads, /*seq_len=*/0, /*causal=*/true, nn::PosEncodingType::RoPE); - if (!attn.init(eng)) { std::cerr << "kv attn init failed\n"; return 1; } + if (!attn.init(eng)) + { + std::cerr << "kv attn init failed\n"; + return 1; + } std::mt19937_64 rng{77}; std::normal_distribution dist(0.0, 1.0); nn::Matrix x_m(d_model, L); - { auto sp = x_m.span(); for (auto& v : sp) v = dist(rng); } + { + auto sp = x_m.span(); + for (auto& v : sp) + v = dist(rng); + } auto x_t = eng.from_matrix(x_m); NN_EXIT(x_t, 1); auto full = attn.forward(*x_t); diff --git a/src/rl_experience_test.cpp b/src/rl_experience_test.cpp index f6c80e8f..9059a434 100644 --- a/src/rl_experience_test.cpp +++ b/src/rl_experience_test.cpp @@ -22,80 +22,70 @@ int g_fail = 0; namespace { -nn::Tensor matrix_tensor(nn::ComputeEngine& engine, std::size_t rows, - std::size_t cols, +nn::Tensor matrix_tensor(nn::ComputeEngine& engine, std::size_t rows, std::size_t cols, std::initializer_list values) { nn::Matrix m(rows, cols); NN_ASSERT_ALWAYS(values.size() == m.size(), "matrix_tensor size"); auto span = m.span(); std::size_t i = 0; - for (float v : values) span[i++] = v; + for (float v : values) + span[i++] = v; return upload(engine, m); } -std::vector values_of(nn::ComputeEngine& engine, - const nn::Tensor& tensor) +std::vector values_of(nn::ComputeEngine& engine, const nn::Tensor& tensor) { auto m = download(engine, tensor); std::vector out; out.reserve(m.size()); - for (float v : m.span()) out.push_back(v); + for (float v : m.span()) + out.push_back(v); return out; } -bool near_values(const std::vector& got, - const std::vector& want, float tol = 1e-5f) +bool near_values(const std::vector& got, const std::vector& want, float tol = 1e-5f) { if (got.size() != want.size()) return false; for (std::size_t i = 0; i < got.size(); ++i) - if (std::fabs(got[i] - want[i]) > - tol * (1.0f + std::fabs(want[i]))) - return false; + if (std::fabs(got[i] - want[i]) > tol * (1.0f + std::fabs(want[i]))) return false; return true; } void test_experience_schema_and_view() { - auto schema_r = nn::ExperienceSchema::create( - {"action", "action_mask"}, - {nn::tensor_field("custom_score")}); + auto schema_r = + nn::ExperienceSchema::create({"action", "action_mask"}, {nn::tensor_field("custom_score")}); CHECK(schema_r.has_value(), "custom experience schema builds"); if (!schema_r) return; nn::ExperienceSchema schema = std::move(*schema_r); CHECK(schema.action().has_value(), "action handle resolved"); CHECK(schema.action_mask().has_value(), "action mask handle resolved"); CHECK(schema.group_id().has_value(), "group id standard field registered"); - CHECK(schema.schema().fields().back().name == "custom_score", - "custom field appended"); + CHECK(schema.schema().fields().back().name == "custom_score", "custom field appended"); auto owner = nn::CpuEngine::create(); auto action = make_tensor(*owner, 2, 2, 0.0f, 1.0f); auto mask = matrix_tensor(*owner, 2, 2, {1, 1, 1, 0}); auto custom = make_tensor(*owner, 2, 2, 0.5f, 0.1f); auto record_r = nn::TensorRecord::create( - schema.schema(), - {nn::NamedTensor{.name = "action", .value = action}, - nn::NamedTensor{.name = "action_mask", .value = mask}, - nn::NamedTensor{.name = "custom_score", .value = custom}}); + schema.schema(), {nn::NamedTensor{.name = "action", .value = action}, + nn::NamedTensor{.name = "action_mask", .value = mask}, + nn::NamedTensor{.name = "custom_score", .value = custom}}); CHECK(record_r.has_value(), "experience record builds"); - if (!record_r) - std::printf("[info] TensorRecord error: %s\n", - record_r.error().message.c_str()); + if (!record_r) std::printf("[info] TensorRecord error: %s\n", record_r.error().message.c_str()); if (!record_r) return; auto view = nn::PolicyBatchView::create(*record_r); CHECK(view.has_value(), "policy batch view builds"); if (!view) return; CHECK(view->action().size() == 4, "view exposes action"); CHECK(view->action_mask().size() == 4, "view exposes explicit mask"); - CHECK(view->old_log_prob() == nullptr, - "view distinguishes absent optional policy fields"); + CHECK(view->old_log_prob() == nullptr, "view distinguishes absent optional policy fields"); auto custom_field = record_r->require("custom_score"); CHECK(custom_field.has_value(), "custom field readable by consumer"); auto missing = nn::TensorRecord::create( - schema.schema(), - {nn::NamedTensor{.name = "action_mask", .value = mask}}); + schema.schema(), {nn::NamedTensor{.name = "action_mask", .value = mask}}); CHECK(!missing, "required action rejects incomplete record"); CHECK(!nn::ExperienceSchema::create({"not_a_standard_field"}), "unknown required standard field rejected"); @@ -108,18 +98,14 @@ void test_gae_hand_calculation() auto value = matrix_tensor(*owner, 3, 1, {0.5f, 0.5f, 0.5f}); auto done = matrix_tensor(*owner, 3, 1, {0, 0, 1}); auto bootstrap = matrix_tensor(*owner, 1, 1, {0.7f}); - nn::GaeAdvantage gae(nn::AdvantageConfig{.gamma = 0.5f, - .lambda = 0.8f, - .precision = nn::Precision::F32}); + nn::GaeAdvantage gae( + nn::AdvantageConfig{.gamma = 0.5f, .lambda = 0.8f, .precision = nn::Precision::F32}); auto out = gae.apply(*owner, reward, value, &done, nullptr, &bootstrap); CHECK(out.has_value(), "GAE forward"); if (!out) return; - CHECK(near_values(values_of(*owner, out->advantage), - {1.85f, 2.75f, 2.5f}), + CHECK(near_values(values_of(*owner, out->advantage), {1.85f, 2.75f, 2.5f}), "GAE hand-computed advantages"); - CHECK(near_values(values_of(*owner, out->returns), - {2.35f, 3.25f, 3.0f}), - "GAE returns"); + CHECK(near_values(values_of(*owner, out->returns), {2.35f, 3.25f, 3.0f}), "GAE returns"); } void test_gae_done_mask_and_bootstrap() @@ -129,28 +115,23 @@ void test_gae_done_mask_and_bootstrap() auto value = matrix_tensor(*owner, 3, 1, {0, 0, 0}); auto done = matrix_tensor(*owner, 3, 1, {0, 1, 0}); auto mask = matrix_tensor(*owner, 3, 1, {1, 0, 1}); - nn::GaeAdvantage gae(nn::AdvantageConfig{.gamma = 1.0f, - .lambda = 1.0f, - .precision = nn::Precision::F32}); + nn::GaeAdvantage gae( + nn::AdvantageConfig{.gamma = 1.0f, .lambda = 1.0f, .precision = nn::Precision::F32}); auto out = gae.apply(*owner, reward, value, &done, &mask); CHECK(out.has_value(), "GAE done/mask forward"); if (!out) return; - CHECK(near_values(values_of(*owner, out->advantage), - {1.0f, 0.0f, 3.0f}), + CHECK(near_values(values_of(*owner, out->advantage), {1.0f, 0.0f, 3.0f}), "done terminates GAE and mask outside steps contributes zero"); auto reward2 = matrix_tensor(*owner, 2, 1, {1, 1}); auto value2 = matrix_tensor(*owner, 2, 1, {0, 0}); auto bootstrap = matrix_tensor(*owner, 1, 1, {2}); nn::GaeAdvantage gae_lambda1( - nn::AdvantageConfig{.gamma = 0.5f, .lambda = 1.0f, - .precision = nn::Precision::F32}); - auto boot = gae_lambda1.apply(*owner, reward2, value2, nullptr, nullptr, - &bootstrap); + nn::AdvantageConfig{.gamma = 0.5f, .lambda = 1.0f, .precision = nn::Precision::F32}); + auto boot = gae_lambda1.apply(*owner, reward2, value2, nullptr, nullptr, &bootstrap); CHECK(boot.has_value(), "GAE bootstrap forward"); if (boot) - CHECK(near_values(values_of(*owner, boot->advantage), - {2.0f, 2.0f}), + CHECK(near_values(values_of(*owner, boot->advantage), {2.0f, 2.0f}), "terminal bootstrap enters GAE"); auto reward3 = matrix_tensor(*owner, 3, 1, {1, 2, 3}); @@ -171,8 +152,7 @@ void test_returns_and_normalization() auto returns = mc.apply(*owner, reward); CHECK(returns.has_value(), "Monte Carlo return forward"); if (returns) - CHECK(near_values(values_of(*owner, *returns), - {2.75f, 3.5f, 3.0f}), + CHECK(near_values(values_of(*owner, *returns), {2.75f, 3.5f, 3.0f}), "Monte Carlo discounted returns"); auto values = matrix_tensor(*owner, 3, 1, {1, 2, 3}); @@ -183,8 +163,7 @@ void test_returns_and_normalization() if (norm) { const float s = std::sqrt(2.0f / 3.0f); - CHECK(near_values(values_of(*owner, *norm), - {-1.0f / s, 0.0f, 1.0f / s}), + CHECK(near_values(values_of(*owner, *norm), {-1.0f / s, 0.0f, 1.0f / s}), "population standard deviation normalization"); } @@ -192,16 +171,14 @@ void test_returns_and_normalization() auto zero_var = normalizer.apply(*owner, constant, mask); CHECK(zero_var.has_value(), "zero-variance normalization returns"); if (zero_var) - CHECK(near_values(values_of(*owner, *zero_var), - {0, 0, 0}), + CHECK(near_values(values_of(*owner, *zero_var), {0, 0, 0}), "zero variance yields zero, not NaN"); } void test_group_advantage() { auto owner = nn::CpuEngine::create(); - auto returns = matrix_tensor(*owner, 6, 1, - {1, 2, 3, 10, 10, 10}); + auto returns = matrix_tensor(*owner, 6, 1, {1, 2, 3, 10, 10, 10}); const std::vector spans{ nn::GroupSpan{.first_row = 0, .row_count = 3}, nn::GroupSpan{.first_row = 3, .row_count = 3}, @@ -211,16 +188,14 @@ void test_group_advantage() CHECK(out.has_value(), "GRPO group normalization"); if (!out) return; const float s = std::sqrt(2.0f / 3.0f); - CHECK(near_values(values_of(*owner, *out), - {-1.0f / s, 0.0f, 1.0f / s, 0, 0, 0}), + CHECK(near_values(values_of(*owner, *out), {-1.0f / s, 0.0f, 1.0f / s, 0, 0, 0}), "multiple groups normalize independently; zero-variance group -> 0"); const std::vector bad{ nn::GroupSpan{.first_row = 0, .row_count = 2}, nn::GroupSpan{.first_row = 3, .row_count = 1}, }; - CHECK(!group.apply(*owner, returns, bad).has_value(), - "non-contiguous group spans rejected"); + CHECK(!group.apply(*owner, returns, bad).has_value(), "non-contiguous group spans rejected"); } } // namespace diff --git a/src/rl_gpu_smoke_test.cpp b/src/rl_gpu_smoke_test.cpp index 55f20976..fc402978 100644 --- a/src/rl_gpu_smoke_test.cpp +++ b/src/rl_gpu_smoke_test.cpp @@ -21,15 +21,15 @@ int g_fail = 0; namespace { -nn::Tensor matrix_tensor(nn::ComputeEngine& engine, std::size_t rows, - std::size_t cols, +nn::Tensor matrix_tensor(nn::ComputeEngine& engine, std::size_t rows, std::size_t cols, std::initializer_list values) { nn::Matrix m(rows, cols); NN_ASSERT_ALWAYS(values.size() == m.size(), "matrix_tensor size"); auto span = m.span(); std::size_t i = 0; - for (float v : values) span[i++] = v; + for (float v : values) + span[i++] = v; return upload(engine, m); } @@ -51,14 +51,10 @@ bool finite(nn::ComputeEngine& engine, const nn::Tensor& tensor) auto gae_out = gae.apply(engine, reward, reward, &done, &mask_rc); CHECK(gae_out.has_value(), "GPU GAE"); if (gae_out) - CHECK(finite(engine, gae_out->advantage) && - finite(engine, gae_out->returns), + CHECK(finite(engine, gae_out->advantage) && finite(engine, gae_out->returns), "GPU GAE finite"); - auto logits = matrix_tensor(engine, 3, 4, - {1, 0, 0, 0, - 0, 1, 0, 0, - -1, 0, 1, 0}); + auto logits = matrix_tensor(engine, 3, 4, {1, 0, 0, 0, 0, 1, 0, 0, -1, 0, 1, 0}); auto actions = matrix_tensor(engine, 1, 4, {0, 1, 2, 0}); auto mask = matrix_tensor(engine, 1, 4, {1, 1, 1, 1}); auto advantage = matrix_tensor(engine, 1, 4, {1, -0.5f, 0.2f, -0.1f}); @@ -67,50 +63,39 @@ bool finite(nn::ComputeEngine& engine, const nn::Tensor& tensor) auto old_logp = matrix_tensor(engine, 1, 4, {0, 0, 0, 0}); auto value = matrix_tensor(engine, 1, 4, {0.7f, 1.2f, 2.5f, 3.4f}); nn::ExperienceSchema exp = NN_UNWRAP(nn::ExperienceSchema::create( - {"action", "action_mask", "advantage", "return", "old_value", - "old_log_prob"})); + {"action", "action_mask", "advantage", "return", "old_value", "old_log_prob"})); nn::TensorRecord record = NN_UNWRAP(nn::TensorRecord::create( - exp.schema(), - {nn::NamedTensor{.name = "action", .value = actions}, - nn::NamedTensor{.name = "action_mask", .value = mask}, - nn::NamedTensor{.name = "value_mask", .value = mask}, - nn::NamedTensor{.name = "advantage", .value = advantage}, - nn::NamedTensor{.name = "return", .value = returns}, - nn::NamedTensor{.name = "old_value", .value = old_value}, - nn::NamedTensor{.name = "old_log_prob", .value = old_logp}})); - nn::PolicyBatchView view = NN_UNWRAP( - nn::PolicyBatchView::create(record)); + exp.schema(), {nn::NamedTensor{.name = "action", .value = actions}, + nn::NamedTensor{.name = "action_mask", .value = mask}, + nn::NamedTensor{.name = "value_mask", .value = mask}, + nn::NamedTensor{.name = "advantage", .value = advantage}, + nn::NamedTensor{.name = "return", .value = returns}, + nn::NamedTensor{.name = "old_value", .value = old_value}, + nn::NamedTensor{.name = "old_log_prob", .value = old_logp}})); + nn::PolicyBatchView view = NN_UNWRAP(nn::PolicyBatchView::create(record)); auto ppo_schema = NN_UNWRAP(nn::TensorSchema::create( - "ppo_output", - {nn::tensor_field("policy", - nn::TensorShapeSpec::exact(3, 4)), - nn::tensor_field("value", - nn::TensorShapeSpec::exact(1, 4))})); + "ppo_output", {nn::tensor_field("policy", nn::TensorShapeSpec::exact(3, 4)), + nn::tensor_field("value", nn::TensorShapeSpec::exact(1, 4))})); nn::ClippedPolicyObjective ppo_policy; auto policy_out = ppo_policy.evaluate_context(engine, logits, view); CHECK(policy_out.has_value(), "GPU PPO policy component"); if (!policy_out) - std::printf("[info] GPU PPO policy error: %s\n", - policy_out.error().message.c_str()); + std::printf("[info] GPU PPO policy error: %s\n", policy_out.error().message.c_str()); nn::EntropyBonusObjective entropy; auto entropy_out = entropy.evaluate_context(engine, logits, view); CHECK(entropy_out.has_value(), "GPU entropy component"); if (!entropy_out) - std::printf("[info] GPU entropy error: %s\n", - entropy_out.error().message.c_str()); + std::printf("[info] GPU entropy error: %s\n", entropy_out.error().message.c_str()); nn::ClippedValueObjective ppo_value; auto value_out = ppo_value.evaluate_context(engine, value, view); CHECK(value_out.has_value(), "GPU PPO value component"); if (!value_out) - std::printf("[info] GPU PPO value error: %s\n", - value_out.error().message.c_str()); + std::printf("[info] GPU PPO value error: %s\n", value_out.error().message.c_str()); nn::PpoObjective ppo; auto ppo_out = ppo.evaluate(engine, ppo_schema, logits, value, view); CHECK(ppo_out.has_value(), "GPU PPO preset"); - if (!ppo_out) - std::printf("[info] GPU PPO error: %s\n", - ppo_out.error().message.c_str()); + if (!ppo_out) std::printf("[info] GPU PPO error: %s\n", ppo_out.error().message.c_str()); if (ppo_out) { auto ph = ppo_out->gradients.field("policy"); @@ -121,15 +106,12 @@ bool finite(nn::ComputeEngine& engine, const nn::Tensor& tensor) auto vg = ppo_out->gradients.require(*vh); CHECK(pg && vg, "GPU PPO gradients"); if (pg && vg) - CHECK(finite(engine, **pg) && finite(engine, **vg), - "GPU PPO gradients finite"); + CHECK(finite(engine, **pg) && finite(engine, **vg), "GPU PPO gradients finite"); } } auto grpo_schema = NN_UNWRAP(nn::TensorSchema::create( - "grpo_output", - {nn::tensor_field("policy", - nn::TensorShapeSpec::exact(3, 4))})); + "grpo_output", {nn::tensor_field("policy", nn::TensorShapeSpec::exact(3, 4))})); const std::vector spans{ nn::GroupSpan{.first_row = 0, .row_count = 2}, nn::GroupSpan{.first_row = 2, .row_count = 2}, @@ -137,17 +119,14 @@ bool finite(nn::ComputeEngine& engine, const nn::Tensor& tensor) nn::GrpoObjective grpo; auto grpo_out = grpo.evaluate(engine, grpo_schema, logits, view, spans); CHECK(grpo_out.has_value(), "GPU GRPO preset"); - if (!grpo_out) - std::printf("[info] GPU GRPO error: %s\n", - grpo_out.error().message.c_str()); + if (!grpo_out) std::printf("[info] GPU GRPO error: %s\n", grpo_out.error().message.c_str()); if (grpo_out) { auto gh = grpo_out->gradients.field("policy"); if (gh) { auto grad = grpo_out->gradients.require(*gh); - CHECK(grad && finite(engine, **grad), - "GPU GRPO gradient finite"); + CHECK(grad && finite(engine, **grad), "GPU GRPO gradient finite"); } } } diff --git a/src/rl_grpo_test.cpp b/src/rl_grpo_test.cpp index 089b9e70..f11d9e03 100644 --- a/src/rl_grpo_test.cpp +++ b/src/rl_grpo_test.cpp @@ -23,46 +23,40 @@ int g_fail = 0; namespace { -nn::Tensor matrix_tensor(nn::ComputeEngine& engine, std::size_t rows, - std::size_t cols, +nn::Tensor matrix_tensor(nn::ComputeEngine& engine, std::size_t rows, std::size_t cols, std::initializer_list values) { nn::Matrix m(rows, cols); NN_ASSERT_ALWAYS(values.size() == m.size(), "matrix_tensor size"); auto span = m.span(); std::size_t i = 0; - for (float v : values) span[i++] = v; + for (float v : values) + span[i++] = v; return upload(engine, m); } -std::vector values_of(nn::ComputeEngine& engine, - const nn::Tensor& tensor) +std::vector values_of(nn::ComputeEngine& engine, const nn::Tensor& tensor) { auto m = download(engine, tensor); std::vector out; out.reserve(m.size()); - for (float v : m.span()) out.push_back(v); + for (float v : m.span()) + out.push_back(v); return out; } -bool near_values(const std::vector& got, - const std::vector& want, - float tol = 2e-5f) +bool near_values(const std::vector& got, const std::vector& want, float tol = 2e-5f) { if (got.size() != want.size()) return false; for (std::size_t i = 0; i < got.size(); ++i) - if (std::fabs(got[i] - want[i]) > - tol * (1.0f + std::fabs(want[i]))) - return false; + if (std::fabs(got[i] - want[i]) > tol * (1.0f + std::fabs(want[i]))) return false; return true; } -nn::Result grpo_record( - nn::Tensor actions, nn::Tensor mask, nn::Tensor returns, - nn::Tensor old_log_prob, nn::Tensor advantage = {}) +nn::Result grpo_record(nn::Tensor actions, nn::Tensor mask, nn::Tensor returns, + nn::Tensor old_log_prob, nn::Tensor advantage = {}) { - auto schema = nn::ExperienceSchema::create( - {"action", "action_mask", "return", "old_log_prob"}); + auto schema = nn::ExperienceSchema::create({"action", "action_mask", "return", "old_log_prob"}); NN_TRY_CHECK(schema); std::vector values{ nn::NamedTensor{.name = "action", .value = actions}, @@ -71,18 +65,15 @@ nn::Result grpo_record( nn::NamedTensor{.name = "old_log_prob", .value = old_log_prob}, }; if (advantage.valid()) - values.push_back(nn::NamedTensor{.name = "advantage", - .value = advantage}); + values.push_back(nn::NamedTensor{.name = "advantage", .value = advantage}); return nn::TensorRecord::create(schema->schema(), std::move(values)); } void test_grpo_multiple_groups() { auto owner = nn::CpuEngine::create(); - auto logits = matrix_tensor(*owner, 3, 6, - {1, 0, 0, 0, 0, 0, - 0, 1, 0, 0, 0, 0, - -1, 0, 1, 0, 0, 0}); + auto logits = + matrix_tensor(*owner, 3, 6, {1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, -1, 0, 1, 0, 0, 0}); auto actions = matrix_tensor(*owner, 1, 6, {0, 1, 2, 0, 1, 2}); auto mask = matrix_tensor(*owner, 1, 6, {1, 1, 1, 1, 1, 1}); auto returns = matrix_tensor(*owner, 6, 1, {1, 2, 3, 10, 10, 10}); @@ -99,15 +90,12 @@ void test_grpo_multiple_groups() nn::GroupSpan{.first_row = 3, .row_count = 3}, }; auto output_schema = nn::TensorSchema::create( - "grpo_output", - {nn::tensor_field("policy", - nn::TensorShapeSpec::exact(3, 6))}); + "grpo_output", {nn::tensor_field("policy", nn::TensorShapeSpec::exact(3, 6))}); CHECK(output_schema.has_value(), "GRPO output schema"); if (!output_schema) return; nn::GrpoObjective preset; - auto report = preset.evaluate(*owner, *output_schema, logits, *view, - spans); + auto report = preset.evaluate(*owner, *output_schema, logits, *view, spans); CHECK(report.has_value(), "GRPO preset"); if (!report) return; @@ -116,8 +104,7 @@ void test_grpo_multiple_groups() CHECK(grouped.has_value(), "GRPO group advantage"); if (!grouped) return; auto grouped_row = NN_UNWRAP(owner->reshape(*grouped, 1, 6)); - auto manual_record = grpo_record(actions, mask, returns, old_logp, - grouped_row); + auto manual_record = grpo_record(actions, mask, returns, old_logp, grouped_row); CHECK(manual_record.has_value(), "manual GRPO record with advantage"); if (!manual_record) return; auto manual_view = nn::PolicyBatchView::create(*manual_record); @@ -127,8 +114,7 @@ void test_grpo_multiple_groups() const float s = std::sqrt(2.0f / 3.0f); auto manual_adv = manual_view->advantage(); CHECK(manual_adv && - near_values(values_of(*owner, *manual_adv), - {-1.0f / s, 0, 1.0f / s, 0, 0, 0}), + near_values(values_of(*owner, *manual_adv), {-1.0f / s, 0, 1.0f / s, 0, 0, 0}), "GRPO group advantages"); nn::ClippedPolicyObjective policy; @@ -147,8 +133,8 @@ void test_grpo_multiple_groups() auto got = report->gradients.require(*gh); CHECK(got.has_value(), "GRPO gradient present"); if (got) - check_close(*owner, **got, expected_grad, - "GRPO preset gradient == manual composition", g_fail); + check_close(*owner, **got, expected_grad, "GRPO preset gradient == manual composition", + g_fail); } void test_grpo_zero_variance_and_missing_old() @@ -165,15 +151,12 @@ void test_grpo_zero_variance_and_missing_old() auto view = nn::PolicyBatchView::create(*record); CHECK(view.has_value(), "zero-variance GRPO view"); if (!view) return; - const std::vector one_group{ - nn::GroupSpan{.first_row = 0, .row_count = 3}}; - auto schema = nn::TensorSchema::create( - "grpo_output", {nn::tensor_field("policy")}); + const std::vector one_group{nn::GroupSpan{.first_row = 0, .row_count = 3}}; + auto schema = nn::TensorSchema::create("grpo_output", {nn::tensor_field("policy")}); CHECK(schema.has_value(), "zero-variance GRPO schema"); if (!schema) return; nn::GrpoObjective preset; - auto report = preset.evaluate(*owner, *schema, logits, *view, - one_group); + auto report = preset.evaluate(*owner, *schema, logits, *view, one_group); CHECK(report.has_value(), "zero-variance GRPO evaluates"); if (report) { @@ -189,22 +172,19 @@ void test_grpo_zero_variance_and_missing_old() } } - auto bad_schema = nn::ExperienceSchema::create( - {"action", "action_mask", "return"}); + auto bad_schema = nn::ExperienceSchema::create({"action", "action_mask", "return"}); CHECK(bad_schema.has_value(), "missing-old schema"); if (!bad_schema) return; auto missing_record = nn::TensorRecord::create( - bad_schema->schema(), - {nn::NamedTensor{.name = "action", .value = actions}, - nn::NamedTensor{.name = "action_mask", .value = mask}, - nn::NamedTensor{.name = "return", .value = returns}}); + bad_schema->schema(), {nn::NamedTensor{.name = "action", .value = actions}, + nn::NamedTensor{.name = "action_mask", .value = mask}, + nn::NamedTensor{.name = "return", .value = returns}}); CHECK(missing_record.has_value(), "missing-old record"); if (!missing_record) return; auto missing_view = nn::PolicyBatchView::create(*missing_record); CHECK(missing_view.has_value(), "missing-old view"); if (missing_view) - CHECK(!preset.evaluate(*owner, *schema, logits, *missing_view, - one_group), + CHECK(!preset.evaluate(*owner, *schema, logits, *missing_view, one_group), "GRPO rejects missing old_log_prob"); } diff --git a/src/rl_output_test.cpp b/src/rl_output_test.cpp index 0ebfd496..eb0cb052 100644 --- a/src/rl_output_test.cpp +++ b/src/rl_output_test.cpp @@ -28,23 +28,16 @@ nn::TensorSchema policy_value_schema() { auto r = nn::TensorSchema::create( "policy_value", - {nn::tensor_field( - "policy", - nn::TensorShapeSpec{.rows = 2, .cols = std::nullopt}, - true, nn::Precision::F32, nn::TensorPlacement::CPU), - nn::tensor_field( - "value", - nn::TensorShapeSpec{.rows = 1, .cols = std::nullopt}, - true, nn::Precision::F32, nn::TensorPlacement::CPU)}); + {nn::tensor_field("policy", nn::TensorShapeSpec{.rows = 2, .cols = std::nullopt}, true, + nn::Precision::F32, nn::TensorPlacement::CPU), + nn::tensor_field("value", nn::TensorShapeSpec{.rows = 1, .cols = std::nullopt}, true, + nn::Precision::F32, nn::TensorPlacement::CPU)}); NN_ASSERT_ALWAYS(r, "policy_value_schema"); return std::move(*r); } -void set_matrix(nn::ComputeEngine& engine, nn::Tensor& tensor, - const nn::Matrix& value) -{ - NN_UNWRAP(engine.copy_from(tensor, value)); -} +void set_matrix(nn::ComputeEngine& engine, nn::Tensor& tensor, const nn::Matrix& value) +{ NN_UNWRAP(engine.copy_from(tensor, value)); } auto make_bundle_model(nn::CpuEngine& engine, const nn::TensorSchema& schema) -> nn::Result @@ -53,22 +46,20 @@ auto make_bundle_model(nn::CpuEngine& engine, const nn::TensorSchema& schema) NN_TRY_CHECK(model.add(2, 2)); auto bundle = std::make_unique(schema); - NN_TRY_CHECK(bundle->add_port( - "policy", std::make_unique(2, 2))); - NN_TRY_CHECK(bundle->add_port( - "value", std::make_unique(2, 1))); + NN_TRY_CHECK(bundle->add_port("policy", std::make_unique(2, 2))); + NN_TRY_CHECK(bundle->add_port("value", std::make_unique(2, 1))); NN_TRY_CHECK(bundle->set_row_band_packing({"policy", "value"})); NN_TRY_CHECK(model.add_layer(std::move(bundle))); // 固定参数，使 named 与 packed 两条路径比较的是同一函数。 - auto matrix = [](std::size_t rows, std::size_t cols, - std::initializer_list items) { + auto matrix = [](std::size_t rows, std::size_t cols, std::initializer_list items) + { nn::Matrix m(rows, cols); - NN_ASSERT_ALWAYS(items.size() == m.size(), - "matrix initializer size"); + NN_ASSERT_ALWAYS(items.size() == m.size(), "matrix initializer size"); auto dst = m.span(); std::size_t i = 0; - for (float v : items) dst[i++] = v; + for (float v : items) + dst[i++] = v; return m; }; const std::vector values{ @@ -80,15 +71,13 @@ auto make_bundle_model(nn::CpuEngine& engine, const nn::TensorSchema& schema) matrix(1, 1, {-0.03f}), }; auto params = model.parameters(); - NN_ASSERT_ALWAYS(params.size() == values.size(), - "bundle model parameter inventory"); + NN_ASSERT_ALWAYS(params.size() == values.size(), "bundle model parameter inventory"); for (std::size_t i = 0; i < values.size(); ++i) set_matrix(engine, params[i].get(), values[i]); return model; } -std::vector snapshot(nn::CpuEngine& engine, - std::vector refs) +std::vector snapshot(nn::CpuEngine& engine, std::vector refs) { std::vector out; out.reserve(refs.size()); @@ -97,21 +86,17 @@ std::vector snapshot(nn::CpuEngine& engine, return out; } -bool matrices_equal(const std::vector& a, - const std::vector& b) +bool matrices_equal(const std::vector& a, const std::vector& b) { if (a.size() != b.size()) return false; for (std::size_t i = 0; i < a.size(); ++i) { - if (a[i].rows() != b[i].rows() || - a[i].cols() != b[i].cols()) - return false; + if (a[i].rows() != b[i].rows() || a[i].cols() != b[i].cols()) return false; const auto as = a[i].span(); const auto bs = b[i].span(); for (std::size_t j = 0; j < as.size(); ++j) { - if (std::fabs(as[j] - bs[j]) > - 1e-5f * (1.0f + std::fabs(as[j]) + std::fabs(bs[j]))) + if (std::fabs(as[j] - bs[j]) > 1e-5f * (1.0f + std::fabs(as[j]) + std::fabs(bs[j]))) return false; } } @@ -147,9 +132,9 @@ class FailSecondForwardLayer final : public nn::Layer [[nodiscard]] nn::Result init_impl(nn::ComputeEngine&) override { return {}; } }; -nn::Result make_objective_grads( - nn::CpuEngine& engine, const nn::TensorSchema& schema, - const nn::NamedTensorSet& output) +nn::Result make_objective_grads(nn::CpuEngine& engine, + const nn::TensorSchema& schema, + const nn::NamedTensorSet& output) { NN_TRY(policy_h, output.field("policy")); NN_TRY(value_h, output.field("value")); @@ -158,12 +143,8 @@ nn::Result make_objective_grads( nn::MSELoss policy_loss; nn::MSELoss value_loss; - auto zero_policy = make_tensor(engine, (*policy)->rows(), - (*policy)->cols(), - 0.0f, 0.0f); - auto zero_value = make_tensor(engine, (*value)->rows(), - (*value)->cols(), - 0.0f, 0.0f); + auto zero_policy = make_tensor(engine, (*policy)->rows(), (*policy)->cols(), 0.0f, 0.0f); + auto zero_value = make_tensor(engine, (*value)->rows(), (*value)->cols(), 0.0f, 0.0f); NN_TRY_CHECK(policy_loss.forward(engine, **policy, zero_policy)); NN_TRY_CHECK(value_loss.forward(engine, **value, zero_value)); NN_TRY(policy_grad, policy_loss.backward()); @@ -171,20 +152,15 @@ nn::Result make_objective_grads( NN_TRY_CHECK(engine.scale_inplace(*value_grad, 0.5f)); auto named = nn::NamedTensorSet::create( - schema, - {nn::NamedTensor{.name = "policy", .value = *policy_grad}, - nn::NamedTensor{.name = "value", .value = *value_grad}}); + schema, {nn::NamedTensor{.name = "policy", .value = *policy_grad}, + nn::NamedTensor{.name = "value", .value = *value_grad}}); NN_TRY_CHECK(named); NN_TRY(packing, nn::RowBandPacking::create(schema, {"policy", "value"})); NN_TRY(packed, packing->pack(engine, *named)); NN_TRY(named_grads, nn::GradientBundle::create( - schema, - {nn::NamedTensor{.name = "policy", - .value = *policy_grad}, - nn::NamedTensor{.name = "value", - .value = *value_grad}})); - return ObjectiveGrads{.named = std::move(*named_grads), - .packed = std::move(*packed)}; + schema, {nn::NamedTensor{.name = "policy", .value = *policy_grad}, + nn::NamedTensor{.name = "value", .value = *value_grad}})); + return ObjectiveGrads{.named = std::move(*named_grads), .packed = std::move(*packed)}; } void test_named_and_packed() @@ -200,9 +176,7 @@ void test_named_and_packed() nn::Model& named_model = *named_model_r; nn::Model& packed_model = *packed_model_r; CHECK(named_model.output_schema().has_value(), "bundle exposes schema"); - CHECK(named_model.forward_named( - make_tensor(engine, 2, 3, 0.2f, 0.1f)) - .has_value(), + CHECK(named_model.forward_named(make_tensor(engine, 2, 3, 0.2f, 0.1f)).has_value(), "named forward"); const nn::Tensor input = make_tensor(engine, 2, 3, 0.2f, 0.1f); @@ -220,34 +194,27 @@ void test_named_and_packed() auto value = named->require(*value_h); CHECK(policy && value, "named output values"); if (!policy || !value) return; - CHECK((*policy)->rows() == 2 && (*policy)->cols() == 3, - "policy shape (A,T)"); - CHECK((*value)->rows() == 1 && (*value)->cols() == 3, - "value shape (1,T)"); + CHECK((*policy)->rows() == 2 && (*policy)->cols() == 3, "policy shape (A,T)"); + CHECK((*value)->rows() == 1 && (*value)->cols() == 3, "value shape (1,T)"); - auto packing = NN_UNWRAP(nn::RowBandPacking::create( - schema, {"policy", "value"})); + auto packing = NN_UNWRAP(nn::RowBandPacking::create(schema, {"policy", "value"})); auto named_packed = packing.pack(engine, *named); CHECK(named_packed.has_value(), "named outputs pack"); auto packed = packed_model.forward(input); CHECK(packed.has_value(), "legacy packed forward"); if (named_packed && packed) - check_close(engine, *packed, *named_packed, - "packed forward equals named ports", g_fail); + check_close(engine, *packed, *named_packed, "packed forward equals named ports", g_fail); auto objective = make_objective_grads(engine, schema, *named); CHECK(objective.has_value(), "build policy/value objective gradients"); if (!objective) return; - CHECK(named_model.backward_named(objective->named).has_value(), - "named backward"); - CHECK(packed_model.backward(objective->packed).has_value(), - "packed backward"); + CHECK(named_model.backward_named(objective->named).has_value(), "named backward"); + CHECK(packed_model.backward(objective->packed).has_value(), "packed backward"); auto named_grads = snapshot(engine, named_model.param_gradients()); auto packed_grads = snapshot(engine, packed_model.param_gradients()); - CHECK(matrices_equal(named_grads, packed_grads), - "named and packed parameter gradients agree"); + CHECK(matrices_equal(named_grads, packed_grads), "named and packed parameter gradients agree"); CHECK(named_grads.size() == 6, "all branch/trunk gradients materialized"); if (named_grads.size() == 6) { @@ -265,8 +232,7 @@ void test_negative_paths() nn::CpuEngine& engine = *owner; const nn::TensorSchema schema = policy_value_schema(); - auto bad_order = nn::RowBandPacking::create( - schema, {"policy", "policy"}); + auto bad_order = nn::RowBandPacking::create(schema, {"policy", "policy"}); CHECK(!bad_order, "packing rejects duplicate field"); auto missing_field = nn::RowBandPacking::create(schema, {"policy"}); CHECK(!missing_field, "packing requires complete coverage"); @@ -282,13 +248,11 @@ void test_negative_paths() if (!policy_h) return; auto policy = output->require(*policy_h); if (!policy) return; - auto wrong_schema = nn::TensorSchema::create( - "wrong", {nn::tensor_field("policy")}); + auto wrong_schema = nn::TensorSchema::create("wrong", {nn::tensor_field("policy")}); CHECK(wrong_schema.has_value(), "wrong schema builds"); if (!wrong_schema) return; auto wrong_grads = nn::GradientBundle::create( - *wrong_schema, - {nn::NamedTensor{.name = "policy", .value = **policy}}); + *wrong_schema, {nn::NamedTensor{.name = "policy", .value = **policy}}); CHECK(wrong_grads.has_value(), "wrong-schema gradient bundle builds"); if (wrong_grads) CHECK(!model_r->backward_named(*wrong_grads).has_value(), @@ -329,12 +293,9 @@ void test_model_spec_roundtrip() CHECK(loaded.has_value(), "output schema spec decodes"); if (!loaded) return; CHECK(loaded->output_schema.has_value(), "loaded schema present"); - CHECK(loaded->output_packing == spec.output_packing, - "loaded packing order"); - CHECK(*loaded->output_schema == *spec.output_schema, - "schema semantic roundtrip"); - CHECK(nn::spec_matches(spec, *loaded), - "spec_matches includes output schema"); + CHECK(loaded->output_packing == spec.output_packing, "loaded packing order"); + CHECK(*loaded->output_schema == *spec.output_schema, "schema semantic roundtrip"); + CHECK(nn::spec_matches(spec, *loaded), "spec_matches includes output schema"); nn::ModelSpec duplicate_packing = spec; duplicate_packing.output_packing = {"policy", "policy"}; @@ -345,10 +306,8 @@ void test_model_spec_roundtrip() nn::ModelSpec old; old.emplace().layer_dims = {8, 4}; auto old_loaded = nn::detail::spec_from_kv( - *nn::KeyValueRecord::parse( - nn::detail::spec_to_kv(old).serialize())); - CHECK(old_loaded && !old_loaded->output_schema && - old_loaded->output_packing.empty(), + *nn::KeyValueRecord::parse(nn::detail::spec_to_kv(old).serialize())); + CHECK(old_loaded && !old_loaded->output_schema && old_loaded->output_packing.empty(), "legacy single-output spec stays schema-free"); } diff --git a/src/rl_pipeline_test.cpp b/src/rl_pipeline_test.cpp index b3cfb0b4..7ac4e51c 100644 --- a/src/rl_pipeline_test.cpp +++ b/src/rl_pipeline_test.cpp @@ -25,11 +25,8 @@ namespace nn::TensorSchema rollout_schema() { auto r = nn::TensorSchema::create( - "rollout", - {nn::tensor_field("x", - nn::TensorShapeSpec::exact(1, 1), - true, nn::Precision::F32, - nn::TensorPlacement::CPU)}); + "rollout", {nn::tensor_field("x", nn::TensorShapeSpec::exact(1, 1), true, + nn::Precision::F32, nn::TensorPlacement::CPU)}); NN_ASSERT_ALWAYS(r, "rollout schema"); return std::move(*r); } @@ -43,42 +40,31 @@ class CountingSource final : public nn::RolloutSource public: std::size_t calls = 0; - explicit CountingSource(nn::CpuEngine& engine) - : engine_(engine), schema_(rollout_schema()) - { - } + explicit CountingSource(nn::CpuEngine& engine) : engine_(engine), schema_(rollout_schema()) {} - [[nodiscard]] nn::Result rollout( - nn::DeterministicRng& rng) override + [[nodiscard]] nn::Result rollout(nn::DeterministicRng& rng) override { ++calls; - const nn::Scalar value = - static_cast(rng.next_u64() % 100000u); - NN_TRY(tensor, nn::detail::upload_span( - engine_, 1, 1, nn::Precision::F32, - std::span(&value, 1))); - return nn::TensorRecord::create( - schema_, - {nn::NamedTensor{.name = "x", .value = *tensor}}); + const nn::Scalar value = static_cast(rng.next_u64() % 100000u); + NN_TRY(tensor, nn::detail::upload_span(engine_, 1, 1, nn::Precision::F32, + std::span(&value, 1))); + return nn::TensorRecord::create(schema_, {nn::NamedTensor{.name = "x", .value = *tensor}}); } }; class BadSizeSampler final : public nn::BatchSampler { public: - [[nodiscard]] nn::Result> sample( - std::size_t, std::size_t, nn::DeterministicRng&) override - { - return std::vector{}; - } + [[nodiscard]] nn::Result> sample(std::size_t, std::size_t, + nn::DeterministicRng&) override + { return std::vector{}; } }; class OutOfRangeSampler final : public nn::BatchSampler { public: - [[nodiscard]] nn::Result> sample( - std::size_t, std::size_t batch_size, - nn::DeterministicRng&) override + [[nodiscard]] nn::Result> sample(std::size_t, std::size_t batch_size, + nn::DeterministicRng&) override { std::vector indices(batch_size, 9999); return indices; @@ -92,32 +78,26 @@ struct RunTrace std::size_t after_hooks = 0; }; -nn::Result run_pipeline( - nn::RLTrainingPipeline& pipeline, std::size_t updates, - nn::ComputeEngine& engine, RunTrace& trace) +nn::Result run_pipeline(nn::RLTrainingPipeline& pipeline, std::size_t updates, + nn::ComputeEngine& engine, RunTrace& trace) { nn::TrainerHooks hooks; - hooks.before_update = - [&](std::span batch) - -> nn::Result - { - ++trace.before_hooks; - if (batch.empty()) NN_FAIL("empty hook batch"); - return {}; - }; - hooks.after_update = - [&](std::span batch) - -> nn::Result - { - ++trace.after_hooks; - if (batch.empty()) NN_FAIL("empty hook batch"); - return {}; - }; + hooks.before_update = [&](std::span batch) -> nn::Result + { + ++trace.before_hooks; + if (batch.empty()) NN_FAIL("empty hook batch"); + return {}; + }; + hooks.after_update = [&](std::span batch) -> nn::Result + { + ++trace.after_hooks; + if (batch.empty()) NN_FAIL("empty hook batch"); + return {}; + }; return pipeline.run( updates, - [&](std::span batch, - nn::DeterministicRng&, nn::MetricSink& metrics) - -> nn::Result + [&](std::span batch, nn::DeterministicRng&, + nn::MetricSink& metrics) -> nn::Result { double sum = 0; for (const nn::TensorRecord* record : batch) @@ -142,11 +122,10 @@ struct PipelineBundle nn::RLTrainingPipeline pipeline; }; -nn::Result make_pipeline( - nn::CpuEngine& engine, std::shared_ptr metrics, - std::shared_ptr checkpoints = {}, - std::uint64_t seed = 177, - bool clear_buffer = false) +nn::Result make_pipeline(nn::CpuEngine& engine, + std::shared_ptr metrics, + std::shared_ptr checkpoints = {}, + std::uint64_t seed = 177, bool clear_buffer = false) { nn::RLPipelineConfig cfg; cfg.batch_size = 2; @@ -158,10 +137,9 @@ nn::Result make_pipeline( auto source = std::make_unique(engine); CountingSource* source_ptr = source.get(); NN_TRY(pipeline, - nn::RLTrainingPipeline::create( - std::move(cfg), std::move(source), - std::make_unique(), - std::move(metrics), std::move(checkpoints), seed)); + nn::RLTrainingPipeline::create(std::move(cfg), std::move(source), + std::make_unique(), + std::move(metrics), std::move(checkpoints), seed)); return PipelineBundle{ .source = source_ptr, .pipeline = std::move(*pipeline), @@ -176,20 +154,13 @@ void test_pipeline_and_metrics() CHECK(bundle.has_value(), "pipeline composition"); if (!bundle) return; RunTrace trace; - CHECK(run_pipeline(bundle->pipeline, 4, *engine, trace).has_value(), - "pipeline run"); - CHECK(bundle->source->calls == 8, - "collect_per_update drives source call count"); - CHECK(trace.batches.size() == 8, - "four updates x batch_size records reach task update"); - CHECK(trace.before_hooks == 4 && trace.after_hooks == 4, - "hooks run once around each update"); - CHECK(metrics->records().size() == 4, - "metrics sink receives one record per update"); - CHECK(bundle->pipeline.buffer().empty(), - "clear_buffer_after_update policy"); - CHECK(bundle->pipeline.update_count() == 4 && - bundle->pipeline.rollout_cursor() == 8, + CHECK(run_pipeline(bundle->pipeline, 4, *engine, trace).has_value(), "pipeline run"); + CHECK(bundle->source->calls == 8, "collect_per_update drives source call count"); + CHECK(trace.batches.size() == 8, "four updates x batch_size records reach task update"); + CHECK(trace.before_hooks == 4 && trace.after_hooks == 4, "hooks run once around each update"); + CHECK(metrics->records().size() == 4, "metrics sink receives one record per update"); + CHECK(bundle->pipeline.buffer().empty(), "clear_buffer_after_update policy"); + CHECK(bundle->pipeline.update_count() == 4 && bundle->pipeline.rollout_cursor() == 8, "pipeline counters"); } @@ -203,29 +174,24 @@ void test_checkpoint_reproducibility() RunTrace first; CHECK(run_pipeline(base->pipeline, 3, *engine, first).has_value(), "base pipeline first segment"); - CHECK(base->pipeline.save_checkpoint().has_value(), - "pipeline checkpoint save"); + CHECK(base->pipeline.save_checkpoint().has_value(), "pipeline checkpoint save"); auto resumed = make_pipeline(*engine, {}, checkpoints, 4242, false); CHECK(resumed.has_value(), "resumed pipeline composition"); if (!resumed) return; RunTrace resume_prefix; - CHECK(run_pipeline(resumed->pipeline, 3, *engine, - resume_prefix).has_value(), + CHECK(run_pipeline(resumed->pipeline, 3, *engine, resume_prefix).has_value(), "resumed pipeline prefix"); - CHECK(resumed->pipeline.restore_checkpoint().has_value(), - "pipeline checkpoint restore"); + CHECK(resumed->pipeline.restore_checkpoint().has_value(), "pipeline checkpoint restore"); RunTrace resumed_tail; - CHECK(run_pipeline(resumed->pipeline, 2, *engine, - resumed_tail).has_value(), + CHECK(run_pipeline(resumed->pipeline, 2, *engine, resumed_tail).has_value(), "resumed pipeline continuation"); auto uninterrupted = make_pipeline(*engine, {}, {}, 4242, false); CHECK(uninterrupted.has_value(), "uninterrupted pipeline"); if (!uninterrupted) return; RunTrace reference; - CHECK(run_pipeline(uninterrupted->pipeline, 5, *engine, - reference).has_value(), + CHECK(run_pipeline(uninterrupted->pipeline, 5, *engine, reference).has_value(), "uninterrupted pipeline run"); const std::size_t offset = reference.batches.size() - 4; @@ -233,14 +199,12 @@ void test_checkpoint_reproducibility() for (std::size_t i = 0; same && i < 4; ++i) same = resumed_tail.batches[i] == reference.batches[offset + i]; CHECK(same, "checkpoint restores RNG/rollout/sampler sequence"); - CHECK(resumed->pipeline.update_count() == 5 && - resumed->pipeline.rollout_cursor() == 10, + CHECK(resumed->pipeline.update_count() == 5 && resumed->pipeline.rollout_cursor() == 10, "restored pipeline counters continue"); NN_UNWRAP(checkpoints->save_section("custom", "opaque")); auto custom = checkpoints->load_section("custom"); - CHECK(custom && *custom == "opaque", - "checkpoint provider supports custom sections"); + CHECK(custom && *custom == "opaque", "checkpoint provider supports custom sections"); } void test_after_update_failure_commits_state() @@ -253,11 +217,11 @@ void test_after_update_failure_commits_state() nn::TrainerHooks hooks; hooks.after_update = [](std::span) -> nn::Result { NN_FAIL("intentional after_update failure"); }; - auto run = bundle->pipeline.run(1, - [](std::span, - nn::DeterministicRng&, nn::MetricSink&) -> nn::Result - { return {}; }, - std::move(hooks)); + auto run = bundle->pipeline.run( + 1, + [](std::span, nn::DeterministicRng&, + nn::MetricSink&) -> nn::Result { return {}; }, + std::move(hooks)); CHECK(!run, "after_update failure propagates"); CHECK(bundle->pipeline.update_count() == 1, "successful model update is committed before after hook"); @@ -273,12 +237,10 @@ void test_invalid_components() cfg.min_buffer_size = 2; cfg.collect_per_update = 2; auto source = std::make_unique(*engine); - CHECK(!nn::RLTrainingPipeline::create( - cfg, nullptr, - std::make_unique()), - "null source rejected"); - CHECK(!nn::RLTrainingPipeline::create( - cfg, std::move(source), nullptr), + CHECK( + !nn::RLTrainingPipeline::create(cfg, nullptr, std::make_unique()), + "null source rejected"); + CHECK(!nn::RLTrainingPipeline::create(cfg, std::move(source), nullptr), "null sampler rejected"); nn::RLPipelineConfig too_small = cfg; @@ -293,9 +255,8 @@ void test_invalid_components() std::make_unique()), "empty checkpoint section rejected"); - auto bad_size = nn::RLTrainingPipeline::create( - cfg, std::make_unique(*engine), - std::make_unique()); + auto bad_size = nn::RLTrainingPipeline::create(cfg, std::make_unique(*engine), + std::make_unique()); CHECK(bad_size.has_value(), "bad-size sampler pipeline builds"); if (bad_size) { @@ -305,8 +266,7 @@ void test_invalid_components() } auto out_of_range = nn::RLTrainingPipeline::create( - cfg, std::make_unique(*engine), - std::make_unique()); + cfg, std::make_unique(*engine), std::make_unique()); CHECK(out_of_range.has_value(), "out-of-range sampler pipeline builds"); if (out_of_range) { diff --git a/src/rl_policy_test.cpp b/src/rl_policy_test.cpp index a3531335..776ec1aa 100644 --- a/src/rl_policy_test.cpp +++ b/src/rl_policy_test.cpp @@ -25,25 +25,25 @@ int g_fail = 0; namespace { -nn::Tensor matrix_tensor(nn::ComputeEngine& engine, std::size_t rows, - std::size_t cols, +nn::Tensor matrix_tensor(nn::ComputeEngine& engine, std::size_t rows, std::size_t cols, std::initializer_list values) { nn::Matrix m(rows, cols); NN_ASSERT_ALWAYS(values.size() == m.size(), "matrix_tensor size"); auto span = m.span(); std::size_t i = 0; - for (float v : values) span[i++] = v; + for (float v : values) + span[i++] = v; return upload(engine, m); } -std::vector values_of(nn::ComputeEngine& engine, - const nn::Tensor& tensor) +std::vector values_of(nn::ComputeEngine& engine, const nn::Tensor& tensor) { auto m = download(engine, tensor); std::vector out; out.reserve(m.size()); - for (float v : m.span()) out.push_back(v); + for (float v : m.span()) + out.push_back(v); return out; } @@ -53,57 +53,41 @@ class ForwardingDistribution final : public nn::PolicyDistribution nn::CategoricalDistribution inner_; public: - [[nodiscard]] nn::Result log_prob( - nn::ComputeEngine& engine, const nn::Tensor& parameters, - const nn::Tensor& actions) const override - { - return inner_.log_prob(engine, parameters, actions); - } - - [[nodiscard]] nn::Result entropy( - nn::ComputeEngine& engine, const nn::Tensor& parameters) const override - { - return inner_.entropy(engine, parameters); - } - - [[nodiscard]] nn::Result log_prob_with_gradient( - nn::ComputeEngine& engine, const nn::Tensor& parameters, - const nn::Tensor& actions, - const nn::Tensor& upstream) const override - { - return inner_.log_prob_with_gradient(engine, parameters, actions, - upstream); - } - - [[nodiscard]] nn::Result entropy_with_gradient( - nn::ComputeEngine& engine, const nn::Tensor& parameters, - const nn::Tensor& upstream) const override - { - return inner_.entropy_with_gradient(engine, parameters, upstream); - } - - [[nodiscard]] nn::Result sample( - nn::ComputeEngine& engine, const nn::Tensor& parameter_mesh, - std::uint64_t seed) const override - { - return inner_.sample(engine, parameter_mesh, seed); - } + [[nodiscard]] nn::Result log_prob(nn::ComputeEngine& engine, + const nn::Tensor& parameters, + const nn::Tensor& actions) const override + { return inner_.log_prob(engine, parameters, actions); } + + [[nodiscard]] nn::Result entropy(nn::ComputeEngine& engine, + const nn::Tensor& parameters) const override + { return inner_.entropy(engine, parameters); } + + [[nodiscard]] nn::Result + log_prob_with_gradient(nn::ComputeEngine& engine, const nn::Tensor& parameters, + const nn::Tensor& actions, const nn::Tensor& upstream) const override + { return inner_.log_prob_with_gradient(engine, parameters, actions, upstream); } + + [[nodiscard]] nn::Result + entropy_with_gradient(nn::ComputeEngine& engine, const nn::Tensor& parameters, + const nn::Tensor& upstream) const override + { return inner_.entropy_with_gradient(engine, parameters, upstream); } + + [[nodiscard]] nn::Result sample(nn::ComputeEngine& engine, + const nn::Tensor& parameter_mesh, + std::uint64_t seed) const override + { return inner_.sample(engine, parameter_mesh, seed); } }; -bool near_values(const std::vector& got, - const std::vector& want, - float tol = 2e-5f) +bool near_values(const std::vector& got, const std::vector& want, float tol = 2e-5f) { if (got.size() != want.size()) return false; for (std::size_t i = 0; i < got.size(); ++i) - if (std::fabs(got[i] - want[i]) > - tol * (1.0f + std::fabs(want[i]))) - return false; + if (std::fabs(got[i] - want[i]) > tol * (1.0f + std::fabs(want[i]))) return false; return true; } -std::vector softmax_column( - const std::vector& logits, std::size_t rows, std::size_t col) +std::vector softmax_column(const std::vector& logits, std::size_t rows, + std::size_t col) { const std::size_t cols = logits.size() / rows; float mx = -std::numeric_limits::infinity(); @@ -116,15 +100,15 @@ std::vector softmax_column( out[r] = std::exp(logits[r * cols + col] - mx); sum += out[r]; } - for (float& v : out) v /= sum; + for (float& v : out) + v /= sum; return out; } nn::TensorSchema fixture_schema() { auto r = nn::ExperienceSchema::create( - {"action", "action_mask", "advantage", "return", "old_value", - "old_log_prob"}); + {"action", "action_mask", "advantage", "return", "old_value", "old_log_prob"}); NN_ASSERT_ALWAYS(r, "fixture schema"); return r->schema(); } @@ -136,24 +120,20 @@ struct Fixture nn::PolicyBatchView view; }; -nn::Result make_fixture( - nn::CpuEngine& engine, nn::Tensor logits, nn::Tensor actions, - nn::Tensor mask, nn::Tensor advantage, nn::Tensor returns, - nn::Tensor old_value, nn::Tensor old_log_prob) +nn::Result make_fixture(nn::CpuEngine& engine, nn::Tensor logits, nn::Tensor actions, + nn::Tensor mask, nn::Tensor advantage, nn::Tensor returns, + nn::Tensor old_value, nn::Tensor old_log_prob) { (void)engine; - NN_TRY(record, nn::TensorRecord::create( - fixture_schema(), - {nn::NamedTensor{.name = "action", .value = actions}, - nn::NamedTensor{.name = "action_mask", .value = mask}, - nn::NamedTensor{.name = "value_mask", .value = mask}, - nn::NamedTensor{.name = "advantage", - .value = advantage}, - nn::NamedTensor{.name = "return", .value = returns}, - nn::NamedTensor{.name = "old_value", - .value = old_value}, - nn::NamedTensor{.name = "old_log_prob", - .value = old_log_prob}})); + NN_TRY(record, + nn::TensorRecord::create( + fixture_schema(), {nn::NamedTensor{.name = "action", .value = actions}, + nn::NamedTensor{.name = "action_mask", .value = mask}, + nn::NamedTensor{.name = "value_mask", .value = mask}, + nn::NamedTensor{.name = "advantage", .value = advantage}, + nn::NamedTensor{.name = "return", .value = returns}, + nn::NamedTensor{.name = "old_value", .value = old_value}, + nn::NamedTensor{.name = "old_log_prob", .value = old_log_prob}})); Fixture out; out.logits = std::move(logits); out.record = std::make_unique(std::move(*record)); @@ -162,8 +142,7 @@ nn::Result make_fixture( return out; } -bool reinforce_loss_unsupported(nn::ComputeEngine& engine, - const nn::Tensor& logits); +bool reinforce_loss_unsupported(nn::ComputeEngine& engine, const nn::Tensor& logits); void test_distribution() { @@ -178,8 +157,7 @@ void test_distribution() const std::vector logits_v{1, 0, 0, 2, -1, 0}; const auto p0 = softmax_column(logits_v, 3, 0); const auto p1 = softmax_column(logits_v, 3, 1); - CHECK(near_values(values_of(*owner, *logp), - {std::log(p0[0]), std::log(p1[1])}), + CHECK(near_values(values_of(*owner, *logp), {std::log(p0[0]), std::log(p1[1])}), "categorical log_prob hand reference"); auto entropy = dist.entropy(*owner, logits); @@ -187,15 +165,16 @@ void test_distribution() if (entropy) { float h0 = 0, h1 = 0; - for (float p : p0) h0 -= p * std::log(p); - for (float p : p1) h1 -= p * std::log(p); + for (float p : p0) + h0 -= p * std::log(p); + for (float p : p1) + h1 -= p * std::log(p); CHECK(near_values(values_of(*owner, *entropy), {h0, h1}), "categorical entropy hand reference"); } auto upstream = matrix_tensor(*owner, 1, 2, {1, 1}); - auto eval = dist.log_prob_with_gradient(*owner, logits, actions, - upstream); + auto eval = dist.log_prob_with_gradient(*owner, logits, actions, upstream); CHECK(eval.has_value(), "log_prob gradient"); if (eval) { @@ -208,9 +187,8 @@ void test_distribution() want_col1[1] = 1 - want_col1[1]; want_col1[2] = -want_col1[2]; CHECK(near_values(values_of(*owner, eval->gradient), - {want_col0[0], want_col1[0], - want_col0[1], want_col1[1], - want_col0[2], want_col1[2]}), + {want_col0[0], want_col1[0], want_col0[1], want_col1[1], want_col0[2], + want_col1[2]}), "log_prob gradient is onehot - softmax"); } @@ -219,17 +197,14 @@ void test_distribution() CHECK(s1 && s2, "categorical sampling"); if (s1 && s2) { - CHECK(values_of(*owner, s1->action) == - values_of(*owner, s2->action), + CHECK(values_of(*owner, s1->action) == values_of(*owner, s2->action), "same categorical sample seed is reproducible"); - CHECK(near_values(values_of(*owner, s1->log_prob), - values_of(*owner, *logp)), + CHECK(near_values(values_of(*owner, s1->log_prob), values_of(*owner, *logp)), "sample log_prob matches distribution"); } auto bad_action = matrix_tensor(*owner, 1, 2, {3, 0}); - CHECK(!dist.log_prob(*owner, logits, bad_action), - "out-of-range categorical action rejected"); + CHECK(!dist.log_prob(*owner, logits, bad_action), "out-of-range categorical action rejected"); } void test_reinforce() @@ -242,8 +217,8 @@ void test_reinforce() auto returns = matrix_tensor(*owner, 1, 2, {2, 3}); auto old_value = matrix_tensor(*owner, 1, 2, {1, 2.5f}); auto old_logp = matrix_tensor(*owner, 1, 2, {0, 0}); - auto fixture = make_fixture(*owner, logits, actions, mask, advantage, - returns, old_value, old_logp); + auto fixture = + make_fixture(*owner, logits, actions, mask, advantage, returns, old_value, old_logp); CHECK(fixture.has_value(), "REINFORCE fixture"); if (!fixture) return; @@ -256,42 +231,35 @@ void test_reinforce() nn::ReinforceObjective reinforce; auto loss = reinforce.forward(*owner, fixture->logits, fixture->view); - CHECK(loss && std::fabs(*loss - want_loss) < 1e-5f, - "REINFORCE loss"); + CHECK(loss && std::fabs(*loss - want_loss) < 1e-5f, "REINFORCE loss"); auto grad = reinforce.backward(); CHECK(grad.has_value(), "REINFORCE cached backward"); nn::ReinforceObjective baseline_reinforce(true); - auto baseline_loss = baseline_reinforce.forward( - *owner, fixture->logits, fixture->view); + auto baseline_loss = baseline_reinforce.forward(*owner, fixture->logits, fixture->view); CHECK(baseline_loss && std::fabs(*baseline_loss - want_loss) < 1e-5f, "REINFORCE old_value baseline equals explicit advantage"); auto baseline_grad = baseline_reinforce.backward(); if (grad && baseline_grad) - check_close(*owner, *baseline_grad, *grad, - "REINFORCE baseline gradient", g_fail); + check_close(*owner, *baseline_grad, *grad, "REINFORCE baseline gradient", g_fail); - auto polymorphic = - std::make_shared(); + auto polymorphic = std::make_shared(); nn::ReinforceObjective injected(false, nn::Precision::F32, polymorphic); - auto injected_loss = injected.forward(*owner, fixture->logits, - fixture->view); + auto injected_loss = injected.forward(*owner, fixture->logits, fixture->view); auto injected_grad = injected.backward(); CHECK(injected_loss && std::fabs(*injected_loss - want_loss) < 1e-5f, "objective accepts custom PolicyDistribution interface"); if (grad && injected_grad) - check_close(*owner, *injected_grad, *grad, - "custom distribution gradient composition", g_fail); + check_close(*owner, *injected_grad, *grad, "custom distribution gradient composition", + g_fail); auto zero_mask = matrix_tensor(*owner, 1, 2, {0, 0}); - auto all_masked = make_fixture(*owner, fixture->logits, actions, - zero_mask, advantage, returns, old_value, - old_logp); + auto all_masked = make_fixture(*owner, fixture->logits, actions, zero_mask, advantage, returns, + old_value, old_logp); CHECK(all_masked.has_value(), "all-mask fixture"); if (!all_masked) return; nn::ReinforceObjective masked_reinforce; - auto masked_loss = masked_reinforce.forward( - *owner, all_masked->logits, all_masked->view); + auto masked_loss = masked_reinforce.forward(*owner, all_masked->logits, all_masked->view); auto masked_grad = masked_reinforce.backward(); CHECK(masked_loss && *masked_loss == 0, "all-mask REINFORCE loss = 0"); if (masked_grad) @@ -313,8 +281,8 @@ void test_value_and_kl() auto returns = matrix_tensor(*owner, 1, 2, {1, 2}); auto old_value = matrix_tensor(*owner, 1, 2, {0, 0}); auto old_logp = matrix_tensor(*owner, 1, 2, {0, 0}); - auto fixture = make_fixture(*owner, logits, actions, mask, advantage, - returns, old_value, old_logp); + auto fixture = + make_fixture(*owner, logits, actions, mask, advantage, returns, old_value, old_logp); CHECK(fixture.has_value(), "value/KL fixture"); if (!fixture) return; @@ -322,24 +290,19 @@ void test_value_and_kl() nn::ValueRegressionObjective value_obj(2.0f); auto value_loss = value_obj.forward(*owner, value, fixture->view); auto value_grad = value_obj.backward(); - CHECK(value_loss && std::fabs(*value_loss - 1.25f) < 1e-5f, - "value regression loss"); - CHECK(value_grad && - near_values(values_of(*owner, *value_grad), {-1, -2}), + CHECK(value_loss && std::fabs(*value_loss - 1.25f) < 1e-5f, "value regression loss"); + CHECK(value_grad && near_values(values_of(*owner, *value_grad), {-1, -2}), "value regression gradient"); nn::SampledKlObjective kl; auto kl_loss = kl.forward(*owner, fixture->logits, fixture->view); CHECK(kl_loss, "sampled KL objective runs"); - const bool unsupported = reinforce_loss_unsupported(*owner, - fixture->logits); - CHECK(unsupported, - "contextual objective rejects target-only Loss entry"); + const bool unsupported = reinforce_loss_unsupported(*owner, fixture->logits); + CHECK(unsupported, "contextual objective rejects target-only Loss entry"); } -bool reinforce_loss_unsupported(nn::ComputeEngine& engine, - const nn::Tensor& logits) +bool reinforce_loss_unsupported(nn::ComputeEngine& engine, const nn::Tensor& logits) { nn::ReinforceObjective objective; auto target = logits; @@ -351,8 +314,7 @@ void test_entropy_finite_difference() auto owner = nn::CpuEngine::create(); const std::vector base_values{1, 0, 0, 2, -1, 0}; auto logits = matrix_tensor(*owner, 3, 2, - {base_values[0], base_values[1], - base_values[2], base_values[3], + {base_values[0], base_values[1], base_values[2], base_values[3], base_values[4], base_values[5]}); auto actions = matrix_tensor(*owner, 1, 2, {0, 1}); auto mask = matrix_tensor(*owner, 1, 2, {1, 1}); @@ -360,8 +322,8 @@ void test_entropy_finite_difference() auto returns = matrix_tensor(*owner, 1, 2, {1, 1}); auto old_value = matrix_tensor(*owner, 1, 2, {0, 0}); auto old_logp = matrix_tensor(*owner, 1, 2, {0, 0}); - auto fixture = make_fixture(*owner, logits, actions, mask, advantage, - returns, old_value, old_logp); + auto fixture = + make_fixture(*owner, logits, actions, mask, advantage, returns, old_value, old_logp); CHECK(fixture.has_value(), "entropy finite-difference fixture"); if (!fixture) return; @@ -378,13 +340,13 @@ void test_entropy_finite_difference() { auto values = base_values; values[i] += eps; - auto plus = matrix_tensor(*owner, 3, 2, - {values[0], values[1], values[2], values[3], values[4], values[5]}); + auto plus = matrix_tensor( + *owner, 3, 2, {values[0], values[1], values[2], values[3], values[4], values[5]}); auto plus_loss = objective.forward(*owner, plus, fixture->view); values = base_values; values[i] -= eps; - auto minus = matrix_tensor(*owner, 3, 2, - {values[0], values[1], values[2], values[3], values[4], values[5]}); + auto minus = matrix_tensor( + *owner, 3, 2, {values[0], values[1], values[2], values[3], values[4], values[5]}); auto minus_loss = objective.forward(*owner, minus, fixture->view); if (!plus_loss || !minus_loss) { diff --git a/src/rl_ppo_test.cpp b/src/rl_ppo_test.cpp index cf3ff807..9bc8f6ae 100644 --- a/src/rl_ppo_test.cpp +++ b/src/rl_ppo_test.cpp @@ -24,45 +24,40 @@ int g_fail = 0; namespace { -nn::Tensor matrix_tensor(nn::ComputeEngine& engine, std::size_t rows, - std::size_t cols, +nn::Tensor matrix_tensor(nn::ComputeEngine& engine, std::size_t rows, std::size_t cols, std::initializer_list values) { nn::Matrix m(rows, cols); NN_ASSERT_ALWAYS(values.size() == m.size(), "matrix_tensor size"); auto span = m.span(); std::size_t i = 0; - for (float v : values) span[i++] = v; + for (float v : values) + span[i++] = v; return upload(engine, m); } -std::vector values_of(nn::ComputeEngine& engine, - const nn::Tensor& tensor) +std::vector values_of(nn::ComputeEngine& engine, const nn::Tensor& tensor) { auto m = download(engine, tensor); std::vector out; out.reserve(m.size()); - for (float v : m.span()) out.push_back(v); + for (float v : m.span()) + out.push_back(v); return out; } -bool near_values(const std::vector& got, - const std::vector& want, - float tol = 2e-5f) +bool near_values(const std::vector& got, const std::vector& want, float tol = 2e-5f) { if (got.size() != want.size()) return false; for (std::size_t i = 0; i < got.size(); ++i) - if (std::fabs(got[i] - want[i]) > - tol * (1.0f + std::fabs(want[i]))) - return false; + if (std::fabs(got[i] - want[i]) > tol * (1.0f + std::fabs(want[i]))) return false; return true; } nn::TensorSchema fixture_schema() { auto r = nn::ExperienceSchema::create( - {"action", "action_mask", "advantage", "return", "old_value", - "old_log_prob"}); + {"action", "action_mask", "advantage", "return", "old_value", "old_log_prob"}); NN_ASSERT_ALWAYS(r, "fixture schema"); return r->schema(); } @@ -73,22 +68,18 @@ struct Fixture nn::PolicyBatchView view; }; -nn::Result make_fixture( - nn::Tensor actions, nn::Tensor mask, nn::Tensor advantage, - nn::Tensor returns, nn::Tensor old_value, nn::Tensor old_log_prob) +nn::Result make_fixture(nn::Tensor actions, nn::Tensor mask, nn::Tensor advantage, + nn::Tensor returns, nn::Tensor old_value, nn::Tensor old_log_prob) { - NN_TRY(record, nn::TensorRecord::create( - fixture_schema(), - {nn::NamedTensor{.name = "action", .value = actions}, - nn::NamedTensor{.name = "action_mask", .value = mask}, - nn::NamedTensor{.name = "value_mask", .value = mask}, - nn::NamedTensor{.name = "advantage", - .value = advantage}, - nn::NamedTensor{.name = "return", .value = returns}, - nn::NamedTensor{.name = "old_value", - .value = old_value}, - nn::NamedTensor{.name = "old_log_prob", - .value = old_log_prob}})); + NN_TRY(record, + nn::TensorRecord::create( + fixture_schema(), {nn::NamedTensor{.name = "action", .value = actions}, + nn::NamedTensor{.name = "action_mask", .value = mask}, + nn::NamedTensor{.name = "value_mask", .value = mask}, + nn::NamedTensor{.name = "advantage", .value = advantage}, + nn::NamedTensor{.name = "return", .value = returns}, + nn::NamedTensor{.name = "old_value", .value = old_value}, + nn::NamedTensor{.name = "old_log_prob", .value = old_log_prob}})); Fixture out; out.record = std::make_unique(std::move(*record)); NN_TRY(view, nn::PolicyBatchView::create(*out.record)); @@ -96,8 +87,7 @@ nn::Result make_fixture( return out; } -nn::Tensor make_old_logp(nn::ComputeEngine& engine, - const nn::Tensor& new_logp, +nn::Tensor make_old_logp(nn::ComputeEngine& engine, const nn::Tensor& new_logp, std::initializer_list ratios) { auto lp = download(engine, new_logp); @@ -127,8 +117,7 @@ void test_policy_clip_branches() CHECK(new_logp.has_value(), "PPO new log-prob"); if (!new_logp) return; auto old_logp = make_old_logp(*owner, *new_logp, {0.5f, 1.5f}); - auto fixture = make_fixture(actions, mask, advantage, returns, - old_value, old_logp); + auto fixture = make_fixture(actions, mask, advantage, returns, old_value, old_logp); CHECK(fixture.has_value(), "positive-advantage PPO fixture"); if (!fixture) return; @@ -136,12 +125,9 @@ void test_policy_clip_branches() auto out = objective.evaluate_context(*owner, logits, fixture->view); CHECK(out.has_value(), "clipped policy evaluates"); if (!out) return; - CHECK(std::fabs(out->loss - (-0.85f)) < 1e-5f, - "PPO low/high clip branches: min(rA, clip(r)A)"); - CHECK(std::fabs(objective.metrics().ratio_mean - 1.0f) < 1e-5f, - "PPO ratio metric"); - CHECK(std::fabs(objective.metrics().clip_fraction - 1.0f) < 1e-5f, - "PPO clip fraction"); + CHECK(std::fabs(out->loss - (-0.85f)) < 1e-5f, "PPO low/high clip branches: min(rA, clip(r)A)"); + CHECK(std::fabs(objective.metrics().ratio_mean - 1.0f) < 1e-5f, "PPO ratio metric"); + CHECK(std::fabs(objective.metrics().clip_fraction - 1.0f) < 1e-5f, "PPO clip fraction"); constexpr float eps = 1e-3f; const auto analytic = values_of(*owner, out->gradient); @@ -151,46 +137,47 @@ void test_policy_clip_branches() { auto plus_v = base; plus_v[i] += eps; - auto plus_t = matrix_tensor(*owner, 3, 2, - {plus_v[0], plus_v[1], plus_v[2], - plus_v[3], plus_v[4], plus_v[5]}); - auto plus_fix = make_fixture(actions, mask, advantage, returns, - old_value, old_logp); - if (!plus_fix) { finite_ok = false; break; } - auto plus = objective.evaluate_context(*owner, plus_t, - plus_fix->view); + auto plus_t = matrix_tensor( + *owner, 3, 2, {plus_v[0], plus_v[1], plus_v[2], plus_v[3], plus_v[4], plus_v[5]}); + auto plus_fix = make_fixture(actions, mask, advantage, returns, old_value, old_logp); + if (!plus_fix) + { + finite_ok = false; + break; + } + auto plus = objective.evaluate_context(*owner, plus_t, plus_fix->view); auto minus_v = base; minus_v[i] -= eps; - auto minus_t = matrix_tensor(*owner, 3, 2, - {minus_v[0], minus_v[1], minus_v[2], - minus_v[3], minus_v[4], minus_v[5]}); - auto minus_fix = make_fixture(actions, mask, advantage, returns, - old_value, old_logp); - if (!minus_fix) { finite_ok = false; break; } - auto minus = objective.evaluate_context(*owner, minus_t, - minus_fix->view); - if (!plus || !minus) { finite_ok = false; break; } + auto minus_t = matrix_tensor( + *owner, 3, 2, {minus_v[0], minus_v[1], minus_v[2], minus_v[3], minus_v[4], minus_v[5]}); + auto minus_fix = make_fixture(actions, mask, advantage, returns, old_value, old_logp); + if (!minus_fix) + { + finite_ok = false; + break; + } + auto minus = objective.evaluate_context(*owner, minus_t, minus_fix->view); + if (!plus || !minus) + { + finite_ok = false; + break; + } const float numeric = (plus->loss - minus->loss) / (2 * eps); if (std::fabs(numeric - analytic[i]) >= 3e-3f) - std::printf("[info] PPO finite-diff [%zu] numeric=%g analytic=%g\n", - i, static_cast(numeric), - static_cast(analytic[i])); - finite_ok = finite_ok && - std::fabs(numeric - analytic[i]) < 3e-3f; + std::printf("[info] PPO finite-diff [%zu] numeric=%g analytic=%g\n", i, + static_cast(numeric), static_cast(analytic[i])); + finite_ok = finite_ok && std::fabs(numeric - analytic[i]) < 3e-3f; } CHECK(finite_ok, "PPO policy clip finite-difference gradient"); auto neg_adv = matrix_tensor(*owner, 1, 2, {-1, -1}); - auto neg_fixture = make_fixture(actions, mask, neg_adv, returns, - old_value, old_logp); + auto neg_fixture = make_fixture(actions, mask, neg_adv, returns, old_value, old_logp); CHECK(neg_fixture.has_value(), "negative-advantage PPO fixture"); if (neg_fixture) { - auto neg = objective.evaluate_context(*owner, logits, - neg_fixture->view); - CHECK(neg && std::fabs(neg->loss - 1.15f) < 1e-5f, - "PPO negative advantage clip branches"); + auto neg = objective.evaluate_context(*owner, logits, neg_fixture->view); + CHECK(neg && std::fabs(neg->loss - 1.15f) < 1e-5f, "PPO negative advantage clip branches"); } } @@ -268,20 +255,16 @@ void test_value_clip() auto returns = matrix_tensor(*owner, 1, 2, {0, 1}); auto old_value = matrix_tensor(*owner, 1, 2, {0, 1}); auto old_logp = matrix_tensor(*owner, 1, 2, {0, 0}); - auto fixture = make_fixture(actions, mask, advantage, returns, - old_value, old_logp); + auto fixture = make_fixture(actions, mask, advantage, returns, old_value, old_logp); CHECK(fixture.has_value(), "value-clip fixture"); if (!fixture) return; auto value = matrix_tensor(*owner, 1, 2, {1, 1.8f}); nn::ClippedValueObjective objective(0.2f, 1.0f); auto out = objective.evaluate_context(*owner, value, fixture->view); - CHECK(out && std::fabs(out->loss - 0.82f) < 1e-5f, - "value clip uses both sides"); - CHECK(out && - near_values(values_of(*owner, out->gradient), {1, 0.8f}), + CHECK(out && std::fabs(out->loss - 0.82f) < 1e-5f, "value clip uses both sides"); + CHECK(out && near_values(values_of(*owner, out->gradient), {1, 0.8f}), "value clip selects max(unclipped, clipped) gradient"); - CHECK(std::fabs(objective.metrics().clip_fraction - 1.0f) < 1e-5f, - "value clip fraction"); + CHECK(std::fabs(objective.metrics().clip_fraction - 1.0f) < 1e-5f, "value clip fraction"); } void test_preset_matches_manual() @@ -294,8 +277,7 @@ void test_preset_matches_manual() auto returns = matrix_tensor(*owner, 1, 2, {1, 2}); auto old_value = matrix_tensor(*owner, 1, 2, {0.5f, 1.5f}); auto old_logp = matrix_tensor(*owner, 1, 2, {-0.2f, 0.1f}); - auto fixture = make_fixture(actions, mask, advantage, returns, - old_value, old_logp); + auto fixture = make_fixture(actions, mask, advantage, returns, old_value, old_logp); CHECK(fixture.has_value(), "PPO preset fixture"); if (!fixture) return; auto value = matrix_tensor(*owner, 1, 2, {0.7f, 1.2f}); @@ -306,31 +288,25 @@ void test_preset_matches_manual() cfg.value_coefficient = 0.5f; cfg.entropy_coefficient = 0.01f; auto output_schema = nn::TensorSchema::create( - "policy_value", - {nn::tensor_field("policy", nn::TensorShapeSpec{.rows = 3, - .cols = 2}), - nn::tensor_field("value", nn::TensorShapeSpec{.rows = 1, - .cols = 2})}); + "policy_value", {nn::tensor_field("policy", nn::TensorShapeSpec{.rows = 3, .cols = 2}), + nn::tensor_field("value", nn::TensorShapeSpec{.rows = 1, .cols = 2})}); CHECK(output_schema.has_value(), "PPO output schema"); if (!output_schema) return; nn::PpoObjective preset(cfg); - auto report = preset.evaluate(*owner, *output_schema, logits, value, - fixture->view); + auto report = preset.evaluate(*owner, *output_schema, logits, value, fixture->view); CHECK(report.has_value(), "PPO preset evaluates"); if (!report) return; nn::ClippedPolicyObjective policy(cfg.clip_epsilon); nn::EntropyBonusObjective entropy(cfg.entropy_coefficient); - nn::ClippedValueObjective value_obj(cfg.value_clip_epsilon, - cfg.value_coefficient); + nn::ClippedValueObjective value_obj(cfg.value_clip_epsilon, cfg.value_coefficient); auto p = policy.evaluate_context(*owner, logits, fixture->view); auto e = entropy.evaluate_context(*owner, logits, fixture->view); auto v = value_obj.evaluate_context(*owner, value, fixture->view); CHECK(p && e && v, "manual objective components evaluate"); if (!p || !e || !v) return; - CHECK(std::fabs(report->loss - (p->loss + e->loss + v->loss)) < - 1e-5f, + CHECK(std::fabs(report->loss - (p->loss + e->loss + v->loss)) < 1e-5f, "PPO preset loss equals component composition"); auto manual_policy_grad = NN_UNWRAP(owner->clone(p->gradient)); @@ -358,16 +334,14 @@ void test_missing_old_fields() auto mask = matrix_tensor(*owner, 1, 1, {1}); auto advantage = matrix_tensor(*owner, 1, 1, {1}); auto returns = matrix_tensor(*owner, 1, 1, {1}); - auto schema = nn::ExperienceSchema::create( - {"action", "action_mask", "advantage", "return"}); + auto schema = nn::ExperienceSchema::create({"action", "action_mask", "advantage", "return"}); CHECK(schema.has_value(), "minimal PPO schema"); if (!schema) return; auto record = nn::TensorRecord::create( - schema->schema(), - {nn::NamedTensor{.name = "action", .value = actions}, - nn::NamedTensor{.name = "action_mask", .value = mask}, - nn::NamedTensor{.name = "advantage", .value = advantage}, - nn::NamedTensor{.name = "return", .value = returns}}); + schema->schema(), {nn::NamedTensor{.name = "action", .value = actions}, + nn::NamedTensor{.name = "action_mask", .value = mask}, + nn::NamedTensor{.name = "advantage", .value = advantage}, + nn::NamedTensor{.name = "return", .value = returns}}); CHECK(record.has_value(), "minimal PPO record"); if (!record) return; auto view = nn::PolicyBatchView::create(*record); diff --git a/src/rl_training_integration_test.cpp b/src/rl_training_integration_test.cpp index 2d2ce3bd..00e94acf 100644 --- a/src/rl_training_integration_test.cpp +++ b/src/rl_training_integration_test.cpp @@ -21,25 +21,25 @@ int g_fail = 0; namespace { -nn::Tensor matrix_tensor(nn::ComputeEngine& engine, std::size_t rows, - std::size_t cols, +nn::Tensor matrix_tensor(nn::ComputeEngine& engine, std::size_t rows, std::size_t cols, std::initializer_list values) { nn::Matrix m(rows, cols); NN_ASSERT_ALWAYS(values.size() == m.size(), "matrix_tensor size"); auto span = m.span(); std::size_t i = 0; - for (float v : values) span[i++] = v; + for (float v : values) + span[i++] = v; return upload(engine, m); } -std::vector values_of(nn::ComputeEngine& engine, - const nn::Tensor& tensor) +std::vector values_of(nn::ComputeEngine& engine, const nn::Tensor& tensor) { auto m = download(engine, tensor); std::vector out; out.reserve(m.size()); - for (float v : m.span()) out.push_back(v); + for (float v : m.span()) + out.push_back(v); return out; } @@ -56,12 +56,8 @@ nn::TensorSchema policy_value_schema() { auto r = nn::TensorSchema::create( "rl_policy_value", - {nn::tensor_field("policy", - nn::TensorShapeSpec{.rows = 3, - .cols = std::nullopt}), - nn::tensor_field("value", - nn::TensorShapeSpec{.rows = 1, - .cols = std::nullopt})}); + {nn::tensor_field("policy", nn::TensorShapeSpec{.rows = 3, .cols = std::nullopt}), + nn::tensor_field("value", nn::TensorShapeSpec{.rows = 1, .cols = std::nullopt})}); NN_ASSERT_ALWAYS(r, "policy_value schema"); return std::move(*r); } @@ -69,8 +65,7 @@ nn::TensorSchema policy_value_schema() nn::ExperienceSchema experience_schema() { auto r = nn::ExperienceSchema::create( - {"action", "action_mask", "advantage", "return", "old_value", - "old_log_prob"}); + {"action", "action_mask", "advantage", "return", "old_value", "old_log_prob"}); NN_ASSERT_ALWAYS(r, "experience schema"); return std::move(*r); } @@ -84,10 +79,8 @@ void test_named_rl_session_update() nn::Model model(*engine); NN_UNWRAP(model.add(2, 2)); auto bundle = std::make_unique(output_schema); - NN_UNWRAP(bundle->add_port( - "policy", std::make_unique(2, 3))); - NN_UNWRAP(bundle->add_port( - "value", std::make_unique(2, 1))); + NN_UNWRAP(bundle->add_port("policy", std::make_unique(2, 3))); + NN_UNWRAP(bundle->add_port("value", std::make_unique(2, 1))); NN_UNWRAP(model.add_layer(std::move(bundle))); auto before = model.forward_named(input); @@ -140,20 +133,16 @@ void test_named_rl_session_update() if (!sample) return; auto mask = matrix_tensor(*engine, 1, 2, {1, 1}); - auto advantage = matrix_tensor(*engine, 1, 2, - {1.0f, -0.5f}); - auto returns = matrix_tensor(*engine, 1, 2, - {1.0f, -0.3f}); + auto advantage = matrix_tensor(*engine, 1, 2, {1.0f, -0.5f}); + auto returns = matrix_tensor(*engine, 1, 2, {1.0f, -0.3f}); auto record = nn::TensorRecord::create( - exp.schema(), - {nn::NamedTensor{.name = "action", .value = sample->action}, - nn::NamedTensor{.name = "action_mask", .value = mask}, - nn::NamedTensor{.name = "value_mask", .value = mask}, - nn::NamedTensor{.name = "advantage", .value = advantage}, - nn::NamedTensor{.name = "return", .value = returns}, - nn::NamedTensor{.name = "old_value", .value = **value}, - nn::NamedTensor{.name = "old_log_prob", - .value = sample->log_prob}}); + exp.schema(), {nn::NamedTensor{.name = "action", .value = sample->action}, + nn::NamedTensor{.name = "action_mask", .value = mask}, + nn::NamedTensor{.name = "value_mask", .value = mask}, + nn::NamedTensor{.name = "advantage", .value = advantage}, + nn::NamedTensor{.name = "return", .value = returns}, + nn::NamedTensor{.name = "old_value", .value = **value}, + nn::NamedTensor{.name = "old_log_prob", .value = sample->log_prob}}); CHECK(record.has_value(), "rollout experience record"); if (!record) return; auto view = nn::PolicyBatchView::create(*record); @@ -166,18 +155,15 @@ void test_named_rl_session_update() auto report = session.train_named_step( input, options, [&](nn::Model& active, nn::ComputeEngine& engine_ref, - const nn::Tensor& batch_input) - -> nn::Result + const nn::Tensor& batch_input) -> nn::Result { NN_TRY(outputs, active.forward_named(batch_input)); NN_TRY(policy_port, outputs->require("policy")); NN_TRY(value_port, outputs->require("value")); nn::PpoObjective ppo(ppo_cfg); NN_TRY(ppo_report, - ppo.evaluate(engine_ref, output_schema, - **policy_port, **value_port, *view)); - losses_finite = losses_finite && - std::isfinite(ppo_report->loss); + ppo.evaluate(engine_ref, output_schema, **policy_port, **value_port, *view)); + losses_finite = losses_finite && std::isfinite(ppo_report->loss); return nn::NamedTrainingStepOutput{ .gradients = std::move(ppo_report->gradients), .gradient_units = 2, diff --git a/src/rmsnorm_gradcheck.cpp b/src/rmsnorm_gradcheck.cpp index a9adeb8e..66c3b7e3 100644 --- a/src/rmsnorm_gradcheck.cpp +++ b/src/rmsnorm_gradcheck.cpp @@ -21,20 +21,20 @@ #include #include "test_common.hpp" -using nn::Scalar; +using nn::ComputeEngine; +using nn::Layer; using nn::Matrix; +using nn::RMSNorm; +using nn::Scalar; using nn::Tensor; using nn::TensorRef; -using nn::ComputeEngine; -using nn::RMSNorm; -using nn::Layer; -namespace { +namespace +{ // 用当前参数前向，返回 L = Σ y*go -Scalar eval_loss(ComputeEngine &engine, Layer &norm, - const Tensor &x, const Tensor &go) +Scalar eval_loss(ComputeEngine& engine, Layer& norm, const Tensor& x, const Tensor& go) { auto y = norm.forward(x); auto y_m = engine.to_matrix(*y); @@ -46,15 +46,12 @@ Scalar eval_loss(ComputeEngine &engine, Layer &norm, // 验证单个张量参数/输入的梯度（逐元素中心差分） // fwd_input: 前向传播使用的输入。参数验证时传固定 x；输入梯度验证时 // 传被扰动的输入副本（扰动 fwd_input 本身并重算 loss）。 -bool check_grad_tensor( - ComputeEngine &engine, Layer &norm, - const Tensor &fwd_input, - const Tensor &go, - Tensor ¶m, // 被扰动的张量（参数或输入） - const Matrix &base, // 扰动前的值 - const Matrix &grad_analytical, // backward 得到的梯度 - const std::string &name, - Scalar eps, Scalar tol) +bool check_grad_tensor(ComputeEngine& engine, Layer& norm, const Tensor& fwd_input, + const Tensor& go, + Tensor& param, // 被扰动的张量（参数或输入） + const Matrix& base, // 扰动前的值 + const Matrix& grad_analytical, // backward 得到的梯度 + const std::string& name, Scalar eps, Scalar tol) { bool ok = true; std::size_t bad = 0; @@ -68,18 +65,30 @@ bool check_grad_tensor( Matrix pp = base; pp.set_value_unchecked(r, c, orig + eps); auto r_p = engine.write(param, pp.span()); - if (!r_p) { std::cerr << "write(+) failed\n"; return false; } + if (!r_p) + { + std::cerr << "write(+) failed\n"; + return false; + } const Scalar lp = eval_loss(engine, norm, fwd_input, go); Matrix pm = base; pm.set_value_unchecked(r, c, orig - eps); auto r_m = engine.write(param, pm.span()); - if (!r_m) { std::cerr << "write(-) failed\n"; return false; } + if (!r_m) + { + std::cerr << "write(-) failed\n"; + return false; + } const Scalar lm = eval_loss(engine, norm, fwd_input, go); // 恢复原值 auto r_r = engine.write(param, base.span()); - if (!r_r) { std::cerr << "write(restore) failed\n"; return false; } + if (!r_r) + { + std::cerr << "write(restore) failed\n"; + return false; + } const Scalar num = (lp - lm) / (Scalar{2} * eps); const Scalar ana = grad_analytical.at_unchecked(r, c); @@ -89,31 +98,28 @@ bool check_grad_tensor( { ok = false; if (bad < 10) - std::cout << " [FAIL] " << name << "(" << r << "," << c - << ") num=" << num << " ana=" << ana << "\n"; + std::cout << " [FAIL] " << name << "(" << r << "," << c << ") num=" << num + << " ana=" << ana << "\n"; ++bad; } } } - std::cout << " " << name << ": " - << (ok ? "OK" : "FAIL") - << " max_err=" << max_err << " bad=" << bad << "\n"; + std::cout << " " << name << ": " << (ok ? "OK" : "FAIL") << " max_err=" << max_err + << " bad=" << bad << "\n"; return ok; } } // namespace -int main(int argc, char *argv[]) +int main(int argc, char* argv[]) { Scalar tol = 2e-2f; bool use_gpu = false; for (int i = 1; i < argc; ++i) { std::string a = argv[i]; - if (a == "--tol" && i + 1 < argc) - tol = static_cast(std::atof(argv[++i])); - else if (a == "--gpu") - use_gpu = true; + if (a == "--tol" && i + 1 < argc) tol = static_cast(std::atof(argv[++i])); + else if (a == "--gpu") use_gpu = true; else if (a == "--help") { std::cout << "用法: rmsnorm_gradcheck [--gpu] [--tol ]\n"; @@ -126,16 +132,15 @@ int main(int argc, char *argv[]) auto engine_res = nn::cli::create_engine(ecfg, std::cout); NN_EXIT(engine_res, 1, "引擎创建失败: "); auto engine = std::move(*engine_res); - ComputeEngine &eng = *engine; + ComputeEngine& eng = *engine; const std::size_t features = 16; - const std::size_t batch = 4; + const std::size_t batch = 4; std::cout << "========================================\n"; std::cout << " RMSNorm 数值梯度检查 (gradcheck)\n"; std::cout << "========================================\n"; - std::cout << " features=" << features << " batch=" << batch - << " tol=" << tol << "\n"; + std::cout << " features=" << features << " batch=" << batch << " tol=" << tol << "\n"; RMSNorm norm(features); NN_EXIT(norm.init(eng), 1, "RMSNorm init 失败: "); @@ -145,8 +150,10 @@ int main(int argc, char *argv[]) Matrix x_m(features, batch); Matrix go_m(features, batch); - 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); + 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); NN_EXIT(x, 1, "from_matrix(x) failed\n"); @@ -159,8 +166,8 @@ int main(int argc, char *argv[]) // ── 清零梯度并 backward ── auto params = norm.parameters(); // [gamma] - auto grads = norm.param_gradients(); // [grad_gamma] - for (auto &g : grads) + auto grads = norm.param_gradients(); // [grad_gamma] + for (auto& g : grads) { auto rz = eng.zero(g.get()); NN_EXIT(rz, 1, "zero failed\n"); @@ -178,8 +185,7 @@ int main(int argc, char *argv[]) auto g = eng.to_matrix(grads[pi].get()); NN_EXIT(g, 1, "to_matrix(grad) failed\n"); std::string name = (pi == 0) ? "gamma" : ("param[" + std::to_string(pi) + "]"); - all_pass &= check_grad_tensor( - eng, norm, *x, *go, params[pi].get(), *p, *g, name, eps, tol); + all_pass &= check_grad_tensor(eng, norm, *x, *go, params[pi].get(), *p, *g, name, eps, tol); } // ── 验证输入梯度（扰动输入副本本身，前向用该副本） ── @@ -188,8 +194,7 @@ int main(int argc, char *argv[]) NN_EXIT(gxm, 1, "to_matrix(grad_x) failed\n"); auto xt = eng.clone(*x); // 可写副本，避免扰动原始输入 Tensor NN_EXIT(xt, 1, "clone(x) failed\n"); - all_pass &= check_grad_tensor( - eng, norm, *xt, *go, *xt, x_m, *gxm, "grad_x", eps, tol); + all_pass &= check_grad_tensor(eng, norm, *xt, *go, *xt, x_m, *gxm, "grad_x", eps, tol); } std::cout << "----------------------------------------\n"; diff --git a/src/run_driver_test.cpp b/src/run_driver_test.cpp index 4bdd648d..38fa0802 100644 --- a/src/run_driver_test.cpp +++ b/src/run_driver_test.cpp @@ -27,12 +27,14 @@ #include static int g_fail = 0; -#define CHECK(cond, msg) \ - do { \ - if (!(cond)) { \ - std::cout << " [FAIL] " << (msg) << " (" << #cond << ")\n"; \ - ++g_fail; \ - } \ +#define CHECK(cond, msg) \ + do \ + { \ + if (!(cond)) \ + { \ + std::cout << " [FAIL] " << (msg) << " (" << #cond << ")\n"; \ + ++g_fail; \ + } \ } while (0) namespace @@ -77,8 +79,8 @@ struct Recorder nn::RunHooks make_hooks() { nn::RunHooks h; - h.do_update = [this](nn::TrainingSession&, const nn::UpdateRequest& r) - -> nn::Result + h.do_update = [this](nn::TrainingSession&, + const nn::UpdateRequest& r) -> nn::Result { reqs.push_back(r); const std::size_t call = do_update_calls++; @@ -94,8 +96,7 @@ struct Recorder } return o; }; - h.on_checkpoint = [this](nn::TrainingSession&, const nn::TrainState& st) - -> nn::Result + h.on_checkpoint = [this](nn::TrainingSession&, const nn::TrainState& st) -> nn::Result { if (fail_ckpt) NN_FAIL("模拟断点失败"); ++ckpt_calls; @@ -121,11 +122,14 @@ nn::IndexedBatchSource make_source(std::size_t n, std::size_t bs, std::uint64_t { auto r = nn::IndexedBatchSource::create( n, bs, seed, - [](std::span ix) - -> nn::Result> + [](std::span ix) -> nn::Result> { return std::make_pair(std::any(), ix.size()); }); // create 对合法参数必成功 - if (!r) { std::cout << " [FATAL] 源创建失败\n"; std::abort(); } + if (!r) + { + std::cout << " [FATAL] 源创建失败\n"; + std::abort(); + } return std::move(*r); } @@ -141,18 +145,16 @@ class StepOnlySource : public nn::BatchSource b.epoch_done = false; return b; } - [[nodiscard]] nn::Result rewind(std::size_t, std::size_t) override - { - return {}; - } + [[nodiscard]] nn::Result rewind(std::size_t, std::size_t) override { return {}; } }; // 空步钩子（0 步运行时不会被调用，但 fit 要求存在） static nn::RunHooks rec_empty_hooks() { nn::RunHooks h; - h.do_update = [](nn::TrainingSession&, const nn::UpdateRequest&) - -> nn::Result { return nn::UpdateOutcome{}; }; + h.do_update = [](nn::TrainingSession&, + const nn::UpdateRequest&) -> nn::Result + { return nn::UpdateOutcome{}; }; return h; } @@ -169,8 +171,7 @@ void test_epoch_loop() cfg.epochs = 2; auto rec = Recorder{}; - auto drv = nn::RunDriver::create(*fx.session, src, - std::make_unique(), cfg); + auto drv = nn::RunDriver::create(*fx.session, src, std::make_unique(), cfg); CHECK(drv.has_value(), "RunDriver::create"); if (!drv) return; drv->set_hooks(rec.make_hooks()); @@ -189,8 +190,7 @@ void test_epoch_loop() // epoch/step 网格 CHECK(rec.reqs[0].epoch == 0 && rec.reqs[0].step == 0, "首步 (0,0)"); CHECK(rec.reqs[3].epoch == 1 && rec.reqs[3].step == 0, "第 4 步 (1,0)"); - CHECK(rep->final_epoch == 2 && rep->final_step_in_epoch == 0, - "正常跑完 = (epochs,0)"); + CHECK(rep->final_epoch == 2 && rep->final_step_in_epoch == 0, "正常跑完 = (epochs,0)"); CHECK(rep->abort_code == 0, "无中止"); } @@ -210,8 +210,8 @@ void test_lr_wiring() cfg.lr = 0.1f; cfg.min_lr = 0.01f; auto rec = Recorder{}; - auto drv = nn::RunDriver::create(*fx.session, src, - std::make_unique(), cfg); + auto drv = + nn::RunDriver::create(*fx.session, src, std::make_unique(), cfg); if (drv) { drv->set_hooks(rec.make_hooks()); @@ -234,8 +234,8 @@ void test_lr_wiring() cfg.lr = 0.1f; cfg.min_lr = 0.01f; auto rec = Recorder{}; - auto drv = nn::RunDriver::create(*fx.session, src, - std::make_unique(), cfg); + auto drv = + nn::RunDriver::create(*fx.session, src, std::make_unique(), cfg); if (drv) { drv->set_hooks(rec.make_hooks()); @@ -260,8 +260,8 @@ void test_checkpoint_cadence() nn::RunDriverConfig cfg; cfg.epochs = 2; auto rec = Recorder{}; - auto drv = nn::RunDriver::create(*fx.session, src, - std::make_unique(), cfg); + auto drv = + nn::RunDriver::create(*fx.session, src, std::make_unique(), cfg); if (drv) { drv->set_hooks(rec.make_hooks()); @@ -282,8 +282,8 @@ void test_checkpoint_cadence() cfg.epochs = 2; cfg.save_interval = 1; auto rec = Recorder{}; - auto drv = nn::RunDriver::create(*fx.session, src, - std::make_unique(), cfg); + auto drv = + nn::RunDriver::create(*fx.session, src, std::make_unique(), cfg); if (drv) { drv->set_hooks(rec.make_hooks()); @@ -298,8 +298,8 @@ void test_checkpoint_cadence() cfg.epochs = 2; auto rec = Recorder{}; rec.fail_ckpt = true; - auto drv = nn::RunDriver::create(*fx.session, src, - std::make_unique(), cfg); + auto drv = + nn::RunDriver::create(*fx.session, src, std::make_unique(), cfg); if (drv) { drv->set_hooks(rec.make_hooks()); @@ -322,8 +322,7 @@ void test_max_steps() cfg.epochs = 10; cfg.max_steps = 4; auto rec = Recorder{}; - auto drv = nn::RunDriver::create(*fx.session, src, - std::make_unique(), cfg); + auto drv = nn::RunDriver::create(*fx.session, src, std::make_unique(), cfg); CHECK(drv.has_value(), "epochs>max_steps 时 create 仍合法"); if (!drv) return; drv->set_hooks(rec.make_hooks()); @@ -348,8 +347,7 @@ void test_stop_after_epoch() cfg.epochs = 5; cfg.stop_after_epoch = 2; // 训到 epoch 2（1-based，不含）→ 跑 2 个 epoch auto rec = Recorder{}; - auto drv = nn::RunDriver::create(*fx.session, src, - std::make_unique(), cfg); + auto drv = nn::RunDriver::create(*fx.session, src, std::make_unique(), cfg); if (!drv) return; drv->set_hooks(rec.make_hooks()); auto rep = drv->fit(); @@ -371,8 +369,7 @@ void test_abort() cfg.epochs = 3; auto rec = Recorder{}; rec.abort_at = 4; // 第 5 次 do_update 中止 - auto drv = nn::RunDriver::create(*fx.session, src, - std::make_unique(), cfg); + auto drv = nn::RunDriver::create(*fx.session, src, std::make_unique(), cfg); if (!drv) return; drv->set_hooks(rec.make_hooks()); auto rep = drv->fit(); @@ -394,8 +391,7 @@ void test_step_driven() cfg.epochs = 0; // 步进驱动 cfg.max_steps = 7; auto rec = Recorder{}; - auto drv = nn::RunDriver::create(*fx.session, src, - std::make_unique(), cfg); + auto drv = nn::RunDriver::create(*fx.session, src, std::make_unique(), cfg); CHECK(drv.has_value(), "epochs=0 + max_steps → create 合法"); if (!drv) return; drv->set_hooks(rec.make_hooks()); @@ -419,13 +415,11 @@ void test_validation() if (!fx.session) return; auto src = make_source(10, 4, 7); - CHECK(!nn::RunDriver::create(*fx.session, src, nullptr, {}).has_value(), - "state 为空 → 报错"); + CHECK(!nn::RunDriver::create(*fx.session, src, nullptr, {}).has_value(), "state 为空 → 报错"); { nn::RunDriverConfig bad; bad.epochs = -1; - CHECK(!nn::RunDriver::create(*fx.session, src, - std::make_unique(), bad) + CHECK(!nn::RunDriver::create(*fx.session, src, std::make_unique(), bad) .has_value(), "epochs=-1 → 报错"); } @@ -434,8 +428,8 @@ void test_validation() nn::RunDriverConfig zero; zero.epochs = 0; zero.max_steps = 0; - auto z = nn::RunDriver::create(*fx.session, src, - std::make_unique(), zero); + auto z = + nn::RunDriver::create(*fx.session, src, std::make_unique(), zero); CHECK(z.has_value(), "epochs=0 + max_steps=0 → 0 步运行合法"); if (z) { @@ -448,8 +442,7 @@ void test_validation() nn::RunDriverConfig bad; bad.epochs = 2; bad.stop_after_epoch = 5; - CHECK(!nn::RunDriver::create(*fx.session, src, - std::make_unique(), bad) + CHECK(!nn::RunDriver::create(*fx.session, src, std::make_unique(), bad) .has_value(), "stop_after_epoch > epochs → 报错"); } @@ -457,10 +450,9 @@ void test_validation() { nn::RunDriverConfig cfg; cfg.epochs = 1; - auto drv = nn::RunDriver::create(*fx.session, src, - std::make_unique(), cfg); - if (drv) - CHECK(!drv->fit().has_value(), "缺 do_update → fit 报错"); + auto drv = + nn::RunDriver::create(*fx.session, src, std::make_unique(), cfg); + if (drv) CHECK(!drv->fit().has_value(), "缺 do_update → fit 报错"); } } diff --git a/src/sampling_test.cpp b/src/sampling_test.cpp index 34e0d325..356d7c71 100644 --- a/src/sampling_test.cpp +++ b/src/sampling_test.cpp @@ -34,9 +34,11 @@ using nn::Scalar; int g_fail = 0; -namespace { +namespace +{ -std::mt19937_64 make_rng() { return std::mt19937_64{0xC0FFEEULL}; } +std::mt19937_64 make_rng() +{ return std::mt19937_64{0xC0FFEEULL}; } // 旧签名（无截断）与新签名默认配置的等价性：同 seed 下返回同一下标、 // 概率向量逐位相同、RNG 状态相同 @@ -70,8 +72,7 @@ void test_default_equals_legacy() } // 多次采样同一 logits：返回出现过的下标集合 -std::vector draw_support(const std::vector& logits, - const SamplingConfig& cfg, +std::vector draw_support(const std::vector& logits, const SamplingConfig& cfg, std::size_t draws) { std::vector seen; @@ -82,7 +83,8 @@ std::vector draw_support(const std::vector& logits, std::vector l = logits; const std::size_t t = nn::sample_next_token_(l, cfg, rng, dist, {}); bool found = false; - for (const std::size_t s : seen) if (s == t) found = true; + for (const std::size_t s : seen) + if (s == t) found = true; if (!found) seen.push_back(t); } return seen; @@ -90,7 +92,8 @@ std::vector draw_support(const std::vector& logits, bool in_set(const std::vector& set, std::size_t v) { - for (const std::size_t s : set) if (s == v) return true; + for (const std::size_t s : set) + if (s == v) return true; return false; } @@ -141,15 +144,15 @@ void test_top_p() cfg.top_p = 0.5; bool p_ok = true; auto seen = draw_support(peaked, cfg, 300); - for (const std::size_t s : seen) if (s != 0) p_ok = false; + for (const std::size_t s : seen) + if (s != 0) p_ok = false; CHECK(p_ok, "top_p=0.5 在近 one-hot 上只留峰（核采样）"); // 阈值极小 ⇒ 至少留 1 个（不会空集） SamplingConfig tiny; tiny.top_p = 1e-9; auto seen_tiny = draw_support(peaked, tiny, 50); - CHECK(!seen_tiny.empty() && in_set(seen_tiny, 0), - "top_p→0 仍保留至少 1 个候选（峰自身）"); + CHECK(!seen_tiny.empty() && in_set(seen_tiny, 0), "top_p→0 仍保留至少 1 个候选（峰自身）"); // 均匀分布 + top_p=0.5：4 个等概率 ⇒ 留 2 个 SamplingConfig half; @@ -188,8 +191,7 @@ void test_repeat_penalty() on.repeat_window = 0; l = {5.0, 4.0}; const std::size_t win0 = nn::sample_next_token_(l, on, rng, dist, hist10); - CHECK(win1 == 1 && win0 == 0, - "repeat_window=1 只看最近一个 token；0 = 全部上下文"); + CHECK(win1 == 1 && win0 == 0, "repeat_window=1 只看最近一个 token；0 = 全部上下文"); // 负 logit 走乘法分支（HF 语义）：{-1, -5} + 罚下标 1 ⇒ {-1, -50} ⇒ 取 0； // 若误用"除法"则为 {-1, -0.5} ⇒ 取 1（本断言语义专属，能区分两者）。 @@ -289,10 +291,8 @@ void test_gpt_model_generate() auto eng_owner = nn::CpuEngine::create(); nn::CpuEngine& eng = *eng_owner; const std::size_t vocab = 32, d_model = 16, seq = 8, heads = 2, d_ff = 32; - nn::GPTModel model(vocab, d_model, seq, heads, d_ff, /*layers=*/1, - nn::PosEncodingType::Learned, - nn::ActivationType::SwiGLU, - nn::NormType::RMSNorm); + nn::GPTModel model(vocab, d_model, seq, heads, d_ff, /*layers=*/1, nn::PosEncodingType::Learned, + nn::ActivationType::SwiGLU, nn::NormType::RMSNorm); NN_EXIT(model.init(eng), 1, "GPTModel init 失败: "); const std::vector prompt = {1, 2, 3}; @@ -316,8 +316,8 @@ void test_gpt_model_generate() SamplingConfig hot = seeded; hot.temperature = 5.0; std::vector> outs; - for (std::uint64_t s : {std::uint64_t{1}, std::uint64_t{2}, std::uint64_t{3}, - std::uint64_t{4}, std::uint64_t{5}, std::uint64_t{6}}) + for (std::uint64_t s : {std::uint64_t{1}, std::uint64_t{2}, std::uint64_t{3}, std::uint64_t{4}, + std::uint64_t{5}, std::uint64_t{6}}) { hot.seed = s; outs.push_back(run(hot)); @@ -326,7 +326,8 @@ void test_gpt_model_generate() for (std::size_t i = 0; i < outs.size(); ++i) { bool dup = false; - for (std::size_t j = 0; j < i; ++j) if (outs[i] == outs[j]) dup = true; + for (std::size_t j = 0; j < i; ++j) + if (outs[i] == outs[j]) dup = true; if (!dup) ++distinct; } CHECK(distinct >= 2, "不同 seed 在高温下给出不同序列（种子生效）"); diff --git a/src/scan_bitwise_test.cpp b/src/scan_bitwise_test.cpp index bdac2f7e..d2fc6df2 100644 --- a/src/scan_bitwise_test.cpp +++ b/src/scan_bitwise_test.cpp @@ -58,13 +58,13 @@ std::uint32_t float_as_bits(float f) // 逐位比较（同 u32 位型才算相等；-0.0 与 0.0 视为**不同**——它们位型不同， // 而本测试的语义正是"位型不动"） -void compare_bits(const char* label, const nn::Matrix& a, const nn::Matrix& b, - int& mismatches, float& max_abs) +void compare_bits(const char* label, const nn::Matrix& a, const nn::Matrix& b, int& mismatches, + float& max_abs) { if (a.rows() != b.rows() || a.cols() != b.cols()) { - std::printf(" [FAIL] %s: 形状不同 (%zu,%zu) vs (%zu,%zu)\n", label, a.rows(), - a.cols(), b.rows(), b.cols()); + std::printf(" [FAIL] %s: 形状不同 (%zu,%zu) vs (%zu,%zu)\n", label, a.rows(), a.cols(), + b.rows(), b.cols()); ++mismatches; return; } @@ -80,9 +80,8 @@ void compare_bits(const char* label, const nn::Matrix& a, const nn::Matrix& b, if (ba != bb) { if (mismatches < 3) - std::printf(" [bit] %s idx=%zu cpu=%.9g(0x%08x) gpu=%.9g(0x%08x)\n", - label, i, static_cast(sa[i]), ba, - static_cast(sb[i]), bb); + std::printf(" [bit] %s idx=%zu cpu=%.9g(0x%08x) gpu=%.9g(0x%08x)\n", label, i, + static_cast(sa[i]), ba, static_cast(sb[i]), bb); ++mismatches; } } @@ -92,7 +91,8 @@ nn::Matrix make_matrix(std::mt19937& rng, std::size_t rows, std::size_t cols) { nn::Matrix m(rows, cols); std::normal_distribution d(0.0f, 0.6f); - for (auto& v : m.span()) v = d(rng); + for (auto& v : m.span()) + v = d(rng); return m; } @@ -107,16 +107,15 @@ nn::Matrix make_boundary(std::mt19937& rng, std::size_t cols) } // 一次 prefix 比对（两个引擎跑同一份输入） -void run_prefix_case(nn::ComputeEngine& cpu, nn::ComputeEngine& gpu, std::size_t dk, - bool causal, bool has_state, bool has_bnd, bool f16) +void run_prefix_case(nn::ComputeEngine& cpu, nn::ComputeEngine& gpu, std::size_t dk, bool causal, + bool has_state, bool has_bnd, bool f16) { const std::size_t B = 2, H = 2, seq = 7; const std::size_t rows = B * H * dk; const nn::Precision prec = f16 ? nn::Precision::F16 : nn::Precision::F32; - std::mt19937 rng(20261015u + dk * 17u + (causal ? 1u : 0u) * 3u + - (has_state ? 1u : 0u) * 5u + (has_bnd ? 1u : 0u) * 7u + - (f16 ? 1u : 0u) * 11u); + std::mt19937 rng(20261015u + dk * 17u + (causal ? 1u : 0u) * 3u + (has_state ? 1u : 0u) * 5u + + (has_bnd ? 1u : 0u) * 7u + (f16 ? 1u : 0u) * 11u); const nn::Matrix Km = make_matrix(rng, rows, seq); const nn::Matrix Vm = make_matrix(rng, rows, seq); const nn::Matrix Pm = make_matrix(rng, rows, seq); @@ -125,7 +124,8 @@ void run_prefix_case(nn::ComputeEngine& cpu, nn::ComputeEngine& gpu, std::size_t const nn::Matrix B0m = make_matrix(rng, H * dk, dk); const nn::Matrix bnd = make_boundary(rng, B * seq); - auto run = [&](nn::ComputeEngine& e) -> nn::Result { + auto run = [&](nn::ComputeEngine& e) -> nn::Result + { NN_TRY(k, e.from_matrix(Km, prec)); NN_TRY(v, e.from_matrix(Vm, prec)); NN_TRY(p, e.from_matrix(Pm, prec)); @@ -133,23 +133,21 @@ void run_prefix_case(nn::ComputeEngine& cpu, nn::ComputeEngine& gpu, std::size_t NN_TRY(a0, e.from_matrix(A0m, prec)); NN_TRY(b0, e.from_matrix(B0m, prec)); NN_TRY(bd, e.from_matrix(bnd, prec)); - auto out = e.scan_prefix_outer(*k, *v, *p, *r, *a0, *b0, has_state, dk, H, causal, - *bd, has_bnd, prec); + auto out = e.scan_prefix_outer(*k, *v, *p, *r, *a0, *b0, has_state, dk, H, causal, *bd, + has_bnd, prec); NN_TRY_CHECK(out); return e.to_matrix(*out); }; ++g_cases; char label[160]; - std::snprintf(label, sizeof(label), "prefix dk=%zu causal=%d state=%d bnd=%d %s", - dk, causal ? 1 : 0, has_state ? 1 : 0, has_bnd ? 1 : 0, - f16 ? "f16" : "f32"); + std::snprintf(label, sizeof(label), "prefix dk=%zu causal=%d state=%d bnd=%d %s", dk, + causal ? 1 : 0, has_state ? 1 : 0, has_bnd ? 1 : 0, f16 ? "f16" : "f32"); auto cm = run(cpu); auto gm = run(gpu); if (!cm || !gm) { - std::printf(" [FAIL] %s: 运行失败 (%s%s)\n", label, - cm ? "" : cm.error().message.c_str(), + std::printf(" [FAIL] %s: 运行失败 (%s%s)\n", label, cm ? "" : cm.error().message.c_str(), gm ? "" : gm.error().message.c_str()); ++g_failures; return; @@ -169,21 +167,22 @@ void run_prefix_case(nn::ComputeEngine& cpu, nn::ComputeEngine& gpu, std::size_t } } -void run_suffix_case(nn::ComputeEngine& cpu, nn::ComputeEngine& gpu, std::size_t dk, - bool causal, bool has_bnd, bool f16) +void run_suffix_case(nn::ComputeEngine& cpu, nn::ComputeEngine& gpu, std::size_t dk, bool causal, + bool has_bnd, bool f16) { const std::size_t B = 2, H = 2, seq = 7; const std::size_t rows = B * H * dk; const nn::Precision prec = f16 ? nn::Precision::F16 : nn::Precision::F32; - std::mt19937 rng(777000u + dk * 13u + (causal ? 1u : 0u) * 3u + - (has_bnd ? 1u : 0u) * 7u + (f16 ? 1u : 0u) * 11u); + std::mt19937 rng(777000u + dk * 13u + (causal ? 1u : 0u) * 3u + (has_bnd ? 1u : 0u) * 7u + + (f16 ? 1u : 0u) * 11u); const nn::Matrix Dm = make_matrix(rng, rows * dk, seq); const nn::Matrix Xm = make_matrix(rng, rows, seq); const nn::Matrix Ym = make_matrix(rng, rows, seq); const nn::Matrix bnd = make_boundary(rng, B * seq); - auto run = [&](nn::ComputeEngine& e) -> nn::Result { + auto run = [&](nn::ComputeEngine& e) -> nn::Result + { NN_TRY(d, e.from_matrix(Dm, prec)); NN_TRY(x, e.from_matrix(Xm, prec)); NN_TRY(y, e.from_matrix(Ym, prec)); @@ -195,8 +194,8 @@ void run_suffix_case(nn::ComputeEngine& cpu, nn::ComputeEngine& gpu, std::size_t ++g_cases; char label[160]; - std::snprintf(label, sizeof(label), "suffix dk=%zu causal=%d bnd=%d %s", dk, - causal ? 1 : 0, has_bnd ? 1 : 0, f16 ? "f16" : "f32"); + std::snprintf(label, sizeof(label), "suffix dk=%zu causal=%d bnd=%d %s", dk, causal ? 1 : 0, + has_bnd ? 1 : 0, f16 ? "f16" : "f32"); auto cm = run(cpu); auto gm = run(gpu); if (!cm || !gm) diff --git a/src/session_checkpoint_test.cpp b/src/session_checkpoint_test.cpp index 39e1844f..e2e88955 100644 --- a/src/session_checkpoint_test.cpp +++ b/src/session_checkpoint_test.cpp @@ -22,20 +22,20 @@ int g_fail = 0; -namespace { +namespace +{ -constexpr std::size_t VOCAB = 32, D_MODEL = 16, SEQ_LEN = 8, HEADS = 2, - D_FF = 32, LAYERS = 1; +constexpr std::size_t VOCAB = 32, D_MODEL = 16, SEQ_LEN = 8, HEADS = 2, D_FF = 32, LAYERS = 1; std::filesystem::path scratch_dir() { - const auto dir = std::filesystem::temp_directory_path() / - "nn_session_ckpt_test"; + const auto dir = std::filesystem::temp_directory_path() / "nn_session_ckpt_test"; std::filesystem::create_directories(dir); return dir; } -nn::ModelSpec tiny_spec() { return nn::make_gpt_spec(VOCAB, D_MODEL, SEQ_LEN, HEADS, D_FF, LAYERS); } +nn::ModelSpec tiny_spec() +{ return nn::make_gpt_spec(VOCAB, D_MODEL, SEQ_LEN, HEADS, D_FF, LAYERS); } std::vector snapshot(nn::CpuEngine& e, nn::Model& m) { @@ -44,29 +44,32 @@ std::vector snapshot(nn::CpuEngine& e, nn::Model& m) { auto m2 = e.to_matrix(t.get(), nn::Precision::F32); NN_ASSERT_ALWAYS(m2, "snapshot: to_matrix"); - for (auto v : m2->span()) out.push_back(v); + for (auto v : m2->span()) + out.push_back(v); } return out; } // 走一遍 train_step（固定输入/标签，确定性） -nn::Result one_step(nn::TrainingSession& s, - const nn::Tensor& x, nn::Scalar lr) +nn::Result one_step(nn::TrainingSession& s, const nn::Tensor& x, + nn::Scalar lr) { nn::TrainingStepOptions opts; opts.lr = lr; opts.normalization_units = 1; opts.force_update = true; - return s.train_step(0, opts, [&](nn::Model& model, nn::ComputeEngine& eng, - int) -> nn::Result { - NN_TRY(logits, model.forward(x)); - // step_fn 契约：返回**对模型输出**的梯度；model.backward 由 session 执行 - //（本测试焦点是断点编排，直接给常数输出梯度，不经 loss 层） - auto go = eng.create_tensor(logits->rows(), logits->cols(), - nn::Precision::F32, nn::InitSpec::constant(0.01f)); - if (!go.valid()) NN_FAIL("one_step: go 分配失败"); - return nn::TrainingStepOutput{std::move(go), 1, {}, 0}; - }); + return s.train_step( + 0, opts, + [&](nn::Model& model, nn::ComputeEngine& eng, int) -> nn::Result + { + NN_TRY(logits, model.forward(x)); + // step_fn 契约：返回**对模型输出**的梯度；model.backward 由 session 执行 + //（本测试焦点是断点编排，直接给常数输出梯度，不经 loss 层） + auto go = eng.create_tensor(logits->rows(), logits->cols(), nn::Precision::F32, + nn::InitSpec::constant(0.01f)); + if (!go.valid()) NN_FAIL("one_step: go 分配失败"); + return nn::TrainingStepOutput{std::move(go), 1, {}, 0}; + }); } nn::TextTrainState sample_domain_state() @@ -117,8 +120,7 @@ void test_roundtrip() // ── [1][2][5] 保存（域字段预填）── auto domain = sample_domain_state(); - CHECK(session.save_checkpoint(path, tiny_spec(), domain).has_value(), - "[1] save_checkpoint"); + CHECK(session.save_checkpoint(path, tiny_spec(), domain).has_value(), "[1] save_checkpoint"); auto loaded_state = snapshot(engine, session.model()); // 保存不应改动参数 CHECK(loaded_state == mid, "[1] save_checkpoint 不改动参数"); @@ -186,8 +188,7 @@ void test_roundtrip() CHECK(!strict_r.has_value(), "[3] Strict：数据集身份不一致 fail-fast"); auto relaxed_r = session2.resume(path, &bad, nn::ResumePolicy::Relaxed); CHECK(relaxed_r.has_value(), "[3] Relaxed：放行"); - if (!relaxed_r) - std::printf(" %s\n", relaxed_r.error().message.c_str()); + if (!relaxed_r) std::printf(" %s\n", relaxed_r.error().message.c_str()); // ── [4] resume 后继续训练：参数确实移动（优化器状态已接上）── auto before = snapshot(engine, session2.model()); @@ -226,11 +227,19 @@ void test_recipe_mismatch() // 新会话用不同优化器名 → Strict 下配方不一致 fail-fast auto model2_r = nn::build_gpt_model_from_spec(engine, tiny_spec(), nn::PrecisionProfile{}); - if (!model2_r) { CHECK(false, "建模型 B"); return; } + if (!model2_r) + { + CHECK(false, "建模型 B"); + return; + } nn::TrainingConfig cfg2; cfg2.optimizer = "sgd"; // 与断点的 adam 不一致 auto ses2_r = nn::TrainingSession::create(std::move(*model2_r), cfg2); - if (!ses2_r) { CHECK(false, "create B"); return; } + if (!ses2_r) + { + CHECK(false, "create B"); + return; + } auto session2 = std::move(*ses2_r); auto cur = sample_domain_state(); session2.fill_recipe(cur.recipe); // cfg2.optimizer = "sgd" ≠ 断点 "adam" diff --git a/src/softmax_gradcheck.cpp b/src/softmax_gradcheck.cpp index d9f00c90..0b845b54 100644 --- a/src/softmax_gradcheck.cpp +++ b/src/softmax_gradcheck.cpp @@ -17,13 +17,14 @@ #include #include "test_common.hpp" -using nn::Scalar; -using nn::Matrix; -using nn::Tensor; using nn::ComputeEngine; +using nn::Matrix; +using nn::Scalar; using nn::Softmax; +using nn::Tensor; -namespace { +namespace +{ Scalar eval_loss(ComputeEngine& engine, Softmax& sm, const Tensor& x, const Tensor& go) @@ -35,8 +36,7 @@ Scalar eval_loss(ComputeEngine& engine, Softmax& sm, const Tensor& x, const Tens } -bool check_input_grad(ComputeEngine& engine, Softmax& sm, - const Tensor& go, const Tensor& grad_x, +bool check_input_grad(ComputeEngine& engine, Softmax& sm, const Tensor& go, const Tensor& grad_x, const Matrix& x_m, Scalar eps, Scalar tol) { bool ok = true; @@ -63,13 +63,13 @@ bool check_input_grad(ComputeEngine& engine, Softmax& sm, { ok = false; if (bad < 8) - std::cout << " [FAIL] (" << r << "," << c << ") num=" << num - << " ana=" << ana << " err=" << err << "\n"; + std::cout << " [FAIL] (" << r << "," << c << ") num=" << num << " ana=" << ana + << " err=" << err << "\n"; ++bad; } } - std::cout << " grad_x: " << (ok ? "OK" : "FAIL") - << " max_err=" << max_err << " bad=" << bad << "\n"; + std::cout << " grad_x: " << (ok ? "OK" : "FAIL") << " max_err=" << max_err << " bad=" << bad + << "\n"; return ok; } @@ -85,8 +85,10 @@ int main(int argc, char* argv[]) { std::string a = argv[i]; if (a == "--tol" && i + 1 < argc) tol = static_cast(std::atof(argv[++i])); - else if (a == "--rows" && i + 1 < argc) rows = static_cast(std::atoi(argv[++i])); - else if (a == "--cols" && i + 1 < argc) cols = static_cast(std::atoi(argv[++i])); + else if (a == "--rows" && i + 1 < argc) + rows = static_cast(std::atoi(argv[++i])); + else if (a == "--cols" && i + 1 < argc) + cols = static_cast(std::atoi(argv[++i])); else if (a == "--gpu") use_gpu = true; else if (a == "--help") { @@ -113,8 +115,10 @@ int main(int argc, char* argv[]) std::uniform_real_distribution dist(-2, 2); Matrix x_m(rows, cols); Matrix go_m(rows, cols); - 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); + 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); diff --git a/src/swiglu_gradcheck.cpp b/src/swiglu_gradcheck.cpp index f066c9da..ccdaa81a 100644 --- a/src/swiglu_gradcheck.cpp +++ b/src/swiglu_gradcheck.cpp @@ -24,21 +24,21 @@ #include #include "test_common.hpp" -using nn::Scalar; -using nn::Matrix; -using nn::Tensor; -using nn::TensorRef; +using nn::ActivationType; using nn::ComputeEngine; using nn::CpuEngine; -using nn::ActivationType; using nn::FeedForward; +using nn::Matrix; +using nn::Scalar; +using nn::Tensor; +using nn::TensorRef; -namespace { +namespace +{ // 用当前参数前向，返回 L = Σ y*go -Scalar eval_loss(ComputeEngine &engine, FeedForward &ff, - const Tensor &x, const Tensor &go) +Scalar eval_loss(ComputeEngine& engine, FeedForward& ff, const Tensor& x, const Tensor& go) { auto y = ff.forward(x); auto y_m = engine.to_matrix(*y); @@ -50,15 +50,12 @@ Scalar eval_loss(ComputeEngine &engine, FeedForward &ff, // 验证单个张量参数/输入的梯度（逐元素中心差分） // fwd_input: 前向传播使用的输入。参数验证时传固定 x；输入梯度验证时 // 传被扰动的输入副本（扰动 fwd_input 本身并重算 loss）。 -bool check_grad_tensor( - ComputeEngine &engine, FeedForward &ff, - const Tensor &fwd_input, - const Tensor &go, - Tensor ¶m, // 被扰动的张量（参数或输入） - const Matrix &base, // 扰动前的值 - const Matrix &grad_analytical, // backward 得到的梯度 - const std::string &name, - Scalar eps, Scalar tol) +bool check_grad_tensor(ComputeEngine& engine, FeedForward& ff, const Tensor& fwd_input, + const Tensor& go, + Tensor& param, // 被扰动的张量（参数或输入） + const Matrix& base, // 扰动前的值 + const Matrix& grad_analytical, // backward 得到的梯度 + const std::string& name, Scalar eps, Scalar tol) { bool ok = true; std::size_t bad = 0; @@ -72,18 +69,30 @@ bool check_grad_tensor( Matrix pp = base; pp.set_value_unchecked(r, c, orig + eps); auto r_p = engine.write(param, pp.span()); - if (!r_p) { std::cerr << "write(+) failed\n"; return false; } + if (!r_p) + { + std::cerr << "write(+) failed\n"; + return false; + } const Scalar lp = eval_loss(engine, ff, fwd_input, go); Matrix pm = base; pm.set_value_unchecked(r, c, orig - eps); auto r_m = engine.write(param, pm.span()); - if (!r_m) { std::cerr << "write(-) failed\n"; return false; } + if (!r_m) + { + std::cerr << "write(-) failed\n"; + return false; + } const Scalar lm = eval_loss(engine, ff, fwd_input, go); // 恢复原值 auto r_r = engine.write(param, base.span()); - if (!r_r) { std::cerr << "write(restore) failed\n"; return false; } + if (!r_r) + { + std::cerr << "write(restore) failed\n"; + return false; + } const Scalar num = (lp - lm) / (Scalar{2} * eps); const Scalar ana = grad_analytical.at_unchecked(r, c); @@ -93,31 +102,28 @@ bool check_grad_tensor( { ok = false; if (bad < 10) - std::cout << " [FAIL] " << name << "(" << r << "," << c - << ") num=" << num << " ana=" << ana << "\n"; + std::cout << " [FAIL] " << name << "(" << r << "," << c << ") num=" << num + << " ana=" << ana << "\n"; ++bad; } } } - std::cout << " " << name << ": " - << (ok ? "OK" : "FAIL") - << " max_err=" << max_err << " bad=" << bad << "\n"; + std::cout << " " << name << ": " << (ok ? "OK" : "FAIL") << " max_err=" << max_err + << " bad=" << bad << "\n"; return ok; } } // namespace -int main(int argc, char *argv[]) +int main(int argc, char* argv[]) { Scalar tol = 2e-2f; bool use_gpu = false; for (int i = 1; i < argc; ++i) { std::string a = argv[i]; - if (a == "--tol" && i + 1 < argc) - tol = static_cast(std::atof(argv[++i])); - else if (a == "--gpu") - use_gpu = true; + if (a == "--tol" && i + 1 < argc) tol = static_cast(std::atof(argv[++i])); + else if (a == "--gpu") use_gpu = true; else if (a == "--help") { std::cout << "用法: swiglu_gradcheck [--gpu] [--tol ]\n"; @@ -130,18 +136,18 @@ int main(int argc, char *argv[]) auto engine_res = nn::cli::create_engine(ecfg, std::cout); NN_EXIT(engine_res, 1, "引擎创建失败: "); auto engine = std::move(*engine_res); - ComputeEngine &eng = *engine; + ComputeEngine& eng = *engine; // 小规模，加速测试 const std::size_t d_model = 8; - const std::size_t d_ff = 16; // SwiGLU 时 fc1 → 2*d_ff=32 - const std::size_t batch = 4; + const std::size_t d_ff = 16; // SwiGLU 时 fc1 → 2*d_ff=32 + const std::size_t batch = 4; std::cout << "========================================\n"; std::cout << " SwiGLU 数值梯度检查 (gradcheck)\n"; std::cout << "========================================\n"; - std::cout << " d_model=" << d_model << " d_ff=" << d_ff - << " batch=" << batch << " tol=" << tol << "\n"; + std::cout << " d_model=" << d_model << " d_ff=" << d_ff << " batch=" << batch << " tol=" << tol + << "\n"; FeedForward ff(d_model, d_ff, ActivationType::SwiGLU); NN_EXIT(ff.init(eng), 1, "FeedForward init 失败: "); @@ -151,8 +157,10 @@ int main(int argc, char *argv[]) Matrix x_m(d_model, batch); Matrix go_m(d_model, batch); - 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); + 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); NN_EXIT(x, 1, "from_matrix(x) failed\n"); @@ -165,8 +173,8 @@ int main(int argc, char *argv[]) // ── 清零梯度并 backward ── auto params = ff.parameters(); - auto grads = ff.param_gradients(); - for (auto &g : grads) + auto grads = ff.param_gradients(); + for (auto& g : grads) { auto rz = eng.zero(g.get()); NN_EXIT(rz, 1, "zero failed\n"); @@ -186,8 +194,7 @@ int main(int argc, char *argv[]) auto g = eng.to_matrix(grads[pi].get()); NN_EXIT(g, 1, "to_matrix(grad) failed\n"); std::string name = pi < 4 ? std::string(pnames[pi]) : ("param[" + std::to_string(pi) + "]"); - all_pass &= check_grad_tensor( - eng, ff, *x, *go, params[pi].get(), *p, *g, name, eps, tol); + all_pass &= check_grad_tensor(eng, ff, *x, *go, params[pi].get(), *p, *g, name, eps, tol); } // ── 验证输入梯度（扰动输入副本本身，前向用该副本） ── @@ -196,8 +203,7 @@ int main(int argc, char *argv[]) NN_EXIT(gxm, 1, "to_matrix(grad_x) failed\n"); auto xt = eng.clone(*x); // 可写副本，避免扰动原始输入 Tensor NN_EXIT(xt, 1, "clone(x) failed\n"); - all_pass &= check_grad_tensor( - eng, ff, *xt, *go, *xt, x_m, *gxm, "grad_x", eps, tol); + all_pass &= check_grad_tensor(eng, ff, *xt, *go, *xt, x_m, *gxm, "grad_x", eps, tol); } std::cout << "----------------------------------------\n"; diff --git a/src/t1_58_test.cpp b/src/t1_58_test.cpp index 68afa9ad..d33ea1b8 100644 --- a/src/t1_58_test.cpp +++ b/src/t1_58_test.cpp @@ -51,7 +51,8 @@ std::string g_self; // argv[0]：子进程自举用 nn::Matrix m(rows, cols); std::mt19937 rng(seed); std::uniform_real_distribution d(lo, hi); - for (auto& v : m.span()) v = d(rng); + for (auto& v : m.span()) + v = d(rng); return upload(eng, m); } @@ -76,8 +77,8 @@ void test_precision_set() nn::CpuEngine& eng = *eng_owner; nn::Linear probe(4, 3); const nn::PrecisionSupport sup = probe.precision_support(); - CHECK(sup.param.mask == r.mask && sup.compute.mask == r.mask && - sup.stable.mask == r.mask && sup.optimizer.mask == r.mask, + CHECK(sup.param.mask == r.mask && sup.compute.mask == r.mask && sup.stable.mask == r.mask && + sup.optimizer.mask == r.mask, "默认 precision_support() = {f16,f32} × 四槽"); CHECK(std::string(probe.layer_name()) == "Linear", "Linear::layer_name()"); (void)eng; @@ -223,13 +224,13 @@ void ref_quant(double (&tau_ref)[kOUT], double (&wq_ref)[kOUT][kIN]) for (std::size_t o = 0; o < kOUT; ++o) { double s = 0.0; - for (std::size_t k = 0; k < kIN; ++k) s += std::fabs(static_cast(kRawW[o][k])); + for (std::size_t k = 0; k < kIN; ++k) + s += std::fabs(static_cast(kRawW[o][k])); tau_ref[o] = s / static_cast(kIN); for (std::size_t k = 0; k < kIN; ++k) { const double w = static_cast(kRawW[o][k]); - wq_ref[o][k] = (w > 0.5 * tau_ref[o]) ? 1.0 - : (w < -0.5 * tau_ref[o]) ? -1.0 : 0.0; + wq_ref[o][k] = (w > 0.5 * tau_ref[o]) ? 1.0 : (w < -0.5 * tau_ref[o]) ? -1.0 : 0.0; } } } @@ -282,8 +283,7 @@ void test_bitlinear_rules() { const nn::Matrix packed = download(eng, bl.quantized_weights_packed()); const std::size_t words = nn::trit_pack::words_per_row(IN); - CHECK(packed.rows() == OUT && packed.cols() == words, - "打包容器形状 (out, ceil(K/16))"); + CHECK(packed.rows() == OUT && packed.cols() == words, "打包容器形状 (out, ceil(K/16))"); for (std::size_t o = 0; o < OUT; ++o) { std::vector row(words); @@ -293,8 +293,7 @@ void test_bitlinear_rules() wq.set_value(o, k, nn::trit_pack::at(row.data(), k)); } } - else - wq = download(eng, bl.quantized_weights()); + else wq = download(eng, bl.quantized_weights()); bool exact = true, in_set = true; for (std::size_t o = 0; o < OUT; ++o) @@ -368,7 +367,8 @@ void test_bitlinear_rules() for (std::size_t o = 0; o < OUT; ++o) { double s = 0.0; - for (std::size_t b = 0; b < B; ++b) s += static_cast(dv[o][b]); + for (std::size_t b = 0; b < B; ++b) + s += static_cast(dv[o][b]); mxb = std::max(mxb, std::fabs(static_cast(gb.at(o, 0)) - s)); } CHECK(mxb < 1e-5, "db = Σ_batch dY"); @@ -400,9 +400,9 @@ void test_bitlinear_f16_compute() constexpr std::size_t IN = kIN, OUT = kOUT, B = kB; nn::PrecisionProfile prof; - prof.param = nn::Precision::T1_58; - prof.compute = nn::Precision::F16; - prof.stable = nn::Precision::F32; + prof.param = nn::Precision::T1_58; + prof.compute = nn::Precision::F16; + prof.stable = nn::Precision::F32; prof.optimizer = nn::Precision::F32; nn::BitLinear bl(IN, OUT); bl.set_precision_profile(prof); @@ -437,8 +437,7 @@ void test_bitlinear_f16_compute() if (bl.uses_packed_weights()) CHECK(bl.quantized_weights_packed().precision() == nn::Precision::F32, "打包容器是 f32 原始字容器"); - else - CHECK(bl.quantized_weights().precision() == nn::Precision::F16, "wq 缓冲恒 f16"); + else CHECK(bl.quantized_weights().precision() == nn::Precision::F16, "wq 缓冲恒 f16"); CHECK(bl.latent_weights().precision() == nn::Precision::F32, "latent 恒 f32（STE 落点）"); } @@ -454,9 +453,9 @@ void test_spec_roundtrip() const nn::Matrix X = known_input_matrix(); nn::ModelSpec spec; - auto &sp = spec.emplace(); - sp.layer_dims = {IN, OUT}; - sp.norm_type = nn::NormType::LayerNorm; + auto& sp = spec.emplace(); + sp.layer_dims = {IN, OUT}; + sp.norm_type = nn::NormType::LayerNorm; sp.weight_quant = nn::WeightQuant::T1_58; nn::PrecisionProfile t1; @@ -508,8 +507,7 @@ void test_spec_roundtrip() { CHECK(peeked->as()->weight_quant == nn::WeightQuant::T1_58, "规格记住 weight_quant=t1_58"); - CHECK(peeked->as()->layer_dims == - spec.as()->layer_dims, + CHECK(peeked->as()->layer_dims == spec.as()->layer_dims, "layer_dims 往返一致"); CHECK(nn::spec_matches(*peeked, spec), "spec_matches 含 weight_quant 比对"); CHECK(nn::spec_summary(*peeked).find("wq=t1_58") != std::string::npos, @@ -549,8 +547,8 @@ void test_spec_roundtrip() double mx = 0.0; for (std::size_t o = 0; o < OUT; ++o) for (std::size_t b = 0; b < B; ++b) - mx = std::max(mx, std::fabs(static_cast(Y2.at(o, b)) - - static_cast(Y_ref.at(o, b)))); + mx = std::max(mx, std::fabs(static_cast(Y2.at(o, b)) - + static_cast(Y_ref.at(o, b)))); CHECK(mx == 0.0, "重建 + 加载后前向与保存前**逐位一致**"); } std::remove(file.c_str()); @@ -569,7 +567,8 @@ bool test_training(nn::ComputeEngine& eng, const char* dev) for (auto& v : Wt.span()) v = static_cast(td(rng) - 1) * 0.5f; nn::Matrix Xm(IN, B); - for (auto& v : Xm.span()) v = ud(rng); + for (auto& v : Xm.span()) + v = ud(rng); nn::Matrix Ym(OUT, B); Wt.multiply_to(Ym, Xm); // Ym = Wt·Xm（教师权重**逐元素固定**，不随列变化） @@ -604,18 +603,46 @@ bool test_training(nn::ComputeEngine& eng, const char* dev) }; for (int step = 0; step < STEPS; ++step) { - if (auto r = model.zero_grad(); !r) { report("zero_grad", r.error().message); return false; } + if (auto r = model.zero_grad(); !r) + { + report("zero_grad", r.error().message); + return false; + } auto pred = model.forward(Xt); - if (!pred) { report("forward", pred.error().message); return false; } + if (!pred) + { + report("forward", pred.error().message); + return false; + } auto loss = mse.forward(eng, *pred, Yt); - if (!loss) { report("loss", loss.error().message); return false; } - if (!std::isfinite(*loss)) { CHECK(false, "loss 非有限"); return false; } + if (!loss) + { + report("loss", loss.error().message); + return false; + } + if (!std::isfinite(*loss)) + { + CHECK(false, "loss 非有限"); + return false; + } if (step == 0) first = *loss; last = *loss; auto grad = mse.backward(); - if (!grad) { report("loss.backward", "no grad"); return false; } - if (auto r = model.backward(*grad); !r) { report("model.backward", r.error().message); return false; } - if (auto r = opt->step(); !r) { report("optimizer.step", r.error().message); return false; } + if (!grad) + { + report("loss.backward", "no grad"); + return false; + } + if (auto r = model.backward(*grad); !r) + { + report("model.backward", r.error().message); + return false; + } + if (auto r = opt->step(); !r) + { + report("optimizer.step", r.error().message); + return false; + } } std::printf(" loss: %.6f → %.6f（%d 步）\n", first, last, STEPS); CHECK(last < first * 0.7f, "loss 显著下降（量化权重仍可训练）"); @@ -633,15 +660,14 @@ bool test_training(nn::ComputeEngine& eng, const char* dev) bool test_preset_t1_58(nn::ComputeEngine& eng, const char* dev) { std::printf("[9] 精度预设 profile_t1_58()（%s）：字段 + 真的接成 BitLinear\n", dev); - const nn::PrecisionProfile t1 = nn::profile_t1_58(); + const nn::PrecisionProfile t1 = nn::profile_t1_58(); const nn::PrecisionProfile f16 = nn::profile_f16(); CHECK(t1.param == nn::Precision::T1_58, "预设 param = T1_58"); CHECK(t1.compute == nn::Precision::F16, "预设 compute = F16（\"其他默认 f16\"）"); CHECK(t1.stable == nn::Precision::F32, "预设 stable = F32（数值敏感链）"); CHECK(t1.optimizer == nn::Precision::F32, "预设 optimizer = F32（Adam 的 m/v）"); - CHECK(t1.compute == f16.compute && t1.stable == f16.stable && - t1.optimizer == f16.optimizer, + CHECK(t1.compute == f16.compute && t1.stable == f16.stable && t1.optimizer == f16.optimizer, "预设与 --f16 只在 param 槽不同（另三槽逐字段相同 → 不会各自漂移）"); CHECK(!nn::is_profile_f32(t1), "预设不是全 f32（适配层必须启用）"); @@ -677,8 +703,10 @@ bool test_preset_t1_58(nn::ComputeEngine& eng, const char* dev) std::mt19937 rng(20261010u); std::uniform_real_distribution ud(-1.f, 1.f); nn::Matrix Xm(IN, B), Ym(OUT, B); - for (auto& v : Xm.span()) v = ud(rng); - for (auto& v : Ym.span()) v = ud(rng); + for (auto& v : Xm.span()) + v = ud(rng); + for (auto& v : Ym.span()) + v = ud(rng); auto opt = nn::create_optimizer("adamw", eng, model.parameters(), model.param_gradients(), /*lr=*/0.05f, /*wd=*/0.0f, t1); @@ -699,18 +727,46 @@ bool test_preset_t1_58(nn::ComputeEngine& eng, const char* dev) constexpr int STEPS = 60; for (int step = 0; step < STEPS; ++step) { - if (auto r = model.zero_grad(); !r) { report("zero_grad", r.error().message); return false; } + if (auto r = model.zero_grad(); !r) + { + report("zero_grad", r.error().message); + return false; + } auto pred = model.forward(Xt); - if (!pred) { report("forward", pred.error().message); return false; } + if (!pred) + { + report("forward", pred.error().message); + return false; + } auto loss = mse.forward(eng, *pred, Yt); - if (!loss) { report("loss", loss.error().message); return false; } - if (!std::isfinite(*loss)) { CHECK(false, "loss 非有限"); return false; } + if (!loss) + { + report("loss", loss.error().message); + return false; + } + if (!std::isfinite(*loss)) + { + CHECK(false, "loss 非有限"); + return false; + } if (step == 0) first = *loss; last = *loss; auto grad = mse.backward(); - if (!grad) { report("loss.backward", "no grad"); return false; } - if (auto r = model.backward(*grad); !r) { report("model.backward", r.error().message); return false; } - if (auto r = opt->step(); !r) { report("optimizer.step", r.error().message); return false; } + if (!grad) + { + report("loss.backward", "no grad"); + return false; + } + if (auto r = model.backward(*grad); !r) + { + report("model.backward", r.error().message); + return false; + } + if (auto r = opt->step(); !r) + { + report("optimizer.step", r.error().message); + return false; + } } std::printf(" loss: %.6f → %.6f（%d 步）\n", first, last, STEPS); CHECK(last < first * 0.5f, "预设（三值权重 + f16 存储）训练 loss 下降到一半以下"); diff --git a/src/tensor_expr_test.cpp b/src/tensor_expr_test.cpp index d1378941..bf1a154f 100644 --- a/src/tensor_expr_test.cpp +++ b/src/tensor_expr_test.cpp @@ -21,20 +21,21 @@ #include #include -using nn::Scalar; -using nn::Matrix; -using nn::Tensor; using nn::ComputeEngine; using nn::CpuEngine; +using nn::Matrix; +using nn::Scalar; +using nn::Tensor; using namespace nn::dsl; -namespace { +namespace +{ // 随机填充 -void fill_random(Matrix& m, std::mt19937& rng, - std::uniform_real_distribution& dist) +void fill_random(Matrix& m, std::mt19937& rng, std::uniform_real_distribution& dist) { - for (auto& v : m.span()) v = dist(rng); + for (auto& v : m.span()) + v = dist(rng); } // 最大绝对误差 @@ -54,40 +55,57 @@ Scalar max_abs_diff(const Matrix& a, const Matrix& b) // 从引擎取回 CPU Matrix bool to_cpu(ComputeEngine& eng, const nn::Result& r, Matrix& out) { - if (!r) { std::cerr << " dsl::compute 失败: " << r.error().message << "\n"; return false; } + if (!r) + { + std::cerr << " dsl::compute 失败: " << r.error().message << "\n"; + return false; + } auto m = eng.to_matrix(*r); - if (!m) { std::cerr << " to_matrix 失败: " << m.error().message << "\n"; return false; } + if (!m) + { + std::cerr << " to_matrix 失败: " << m.error().message << "\n"; + return false; + } out = std::move(*m); return true; } // 单个逐元素用例：计算并比对 template -int check(ComputeEngine& eng, const std::string& name, const Expr& expr, - std::size_t rows, std::size_t cols, const Matrix& ref, Scalar tol) +int check(ComputeEngine& eng, const std::string& name, const Expr& expr, std::size_t rows, + std::size_t cols, const Matrix& ref, Scalar tol) { auto r = nn::dsl::compute(eng, expr, rows, cols); Matrix got; NN_EXIT(to_cpu(eng, r, got), 1, "[FAIL] "); const Scalar err = max_abs_diff(ref, got); const bool ok = err <= tol; - std::cout << "[" << (ok ? "PASS" : "FAIL") << "] " << name - << " err=" << std::scientific << std::setprecision(2) << err << "\n"; + std::cout << "[" << (ok ? "PASS" : "FAIL") << "] " << name << " err=" << std::scientific + << std::setprecision(2) << err << "\n"; return ok ? 0 : 1; } // 逐元素参考：out = op(a, b) void ref_binary(Matrix& out, const Matrix& a, const Matrix& b, char op) { - const auto sa = a.span(), sb = b.span(); auto so = out.span(); + const auto sa = a.span(), sb = b.span(); + auto so = out.span(); for (std::size_t i = 0; i < sa.size(); ++i) { switch (op) { - case '+': so[i] = sa[i] + sb[i]; break; - case '-': so[i] = sa[i] - sb[i]; break; - case '*': so[i] = sa[i] * sb[i]; break; - case '/': so[i] = sa[i] / sb[i]; break; + case '+': + so[i] = sa[i] + sb[i]; + break; + case '-': + so[i] = sa[i] - sb[i]; + break; + case '*': + so[i] = sa[i] * sb[i]; + break; + case '/': + so[i] = sa[i] / sb[i]; + break; } } } @@ -111,8 +129,10 @@ int main() const std::size_t R = 4, C = 5; Matrix mA(R, C), mB(R, C), mC(R, C), mD(R, C); - fill_random(mA, rng, dist); fill_random(mB, rng, dist); - fill_random(mC, rng, dist); fill_random(mD, rng, dist); + fill_random(mA, rng, dist); + fill_random(mB, rng, dist); + fill_random(mC, rng, dist); + fill_random(mD, rng, dist); const Tensor ta = upload(eng, mA); const Tensor tb = upload(eng, mB); @@ -121,31 +141,40 @@ int main() // ── 1. 二元：A + B ──────────────────────────────────────────────── { - Matrix ref(R, C); ref_binary(ref, mA, mB, '+'); + Matrix ref(R, C); + ref_binary(ref, mA, mB, '+'); fail += check(eng, "A + B", leaf(ta) + leaf(tb), R, C, ref, kTol); } // ── 2. 融合：A * B + C（一次求值完成乘与加，编译期融合） ────────── { - Matrix ab(R, C); ref_binary(ab, mA, mB, '*'); - Matrix ref(R, C); ref_binary(ref, ab, mC, '+'); + Matrix ab(R, C); + ref_binary(ab, mA, mB, '*'); + Matrix ref(R, C); + ref_binary(ref, ab, mC, '+'); fail += check(eng, "A * B + C (fused)", leaf(ta) * leaf(tb) + leaf(tc), R, C, ref, kTol); } // ── 3. 嵌套融合：(A + B) * (C - D) ──────────────────────────────── { - Matrix ab(R, C); ref_binary(ab, mA, mB, '+'); - Matrix cd(R, C); ref_binary(cd, mC, mD, '-'); - Matrix ref(R, C); ref_binary(ref, ab, cd, '*'); - fail += check(eng, "(A+B)*(C-D)", (leaf(ta) + leaf(tb)) * (leaf(tc) - leaf(td)), R, C, ref, kTol); + Matrix ab(R, C); + ref_binary(ab, mA, mB, '+'); + Matrix cd(R, C); + ref_binary(cd, mC, mD, '-'); + Matrix ref(R, C); + ref_binary(ref, ab, cd, '*'); + fail += check(eng, "(A+B)*(C-D)", (leaf(ta) + leaf(tb)) * (leaf(tc) - leaf(td)), R, C, ref, + kTol); } // ── 4. 标量广播：A * 2 + B 与 2 * A ──────────────────────────────── { Matrix ref(R, C); - for (std::size_t i = 0; i < mA.size(); ++i) ref.span()[i] = mA.span()[i] * Scalar{2} + mB.span()[i]; + for (std::size_t i = 0; i < mA.size(); ++i) + ref.span()[i] = mA.span()[i] * Scalar{2} + mB.span()[i]; fail += check(eng, "A*2 + B (scalar)", leaf(ta) * Scalar{2} + leaf(tb), R, C, ref, kTol); } { Matrix ref(R, C); - for (std::size_t i = 0; i < mA.size(); ++i) ref.span()[i] = Scalar{2} * mA.span()[i]; + for (std::size_t i = 0; i < mA.size(); ++i) + ref.span()[i] = Scalar{2} * mA.span()[i]; fail += check(eng, "2 * A (scalar-left)", Scalar{2} * leaf(ta), R, C, ref, kTol); } // ── 5. 一元融合：relu(A) + tanh(B) ───────────────────────────────── @@ -166,17 +195,24 @@ int main() // ── 7. matmul 边界 + 逐元素：matmul(A2x3, B3x4) + C ─────────────── { Matrix a2(2, 3), b3(3, 4), c2(2, 4); - fill_random(a2, rng, dist); fill_random(b3, rng, dist); fill_random(c2, rng, dist); + fill_random(a2, rng, dist); + fill_random(b3, rng, dist); + fill_random(c2, rng, dist); Matrix ref(2, 4); for (std::size_t i = 0; i < 2; ++i) for (std::size_t j = 0; j < 4; ++j) { Scalar s = 0; - for (std::size_t k = 0; k < 3; ++k) s += a2.span()[i * 3 + k] * b3.span()[k * 4 + j]; + for (std::size_t k = 0; k < 3; ++k) + s += a2.span()[i * 3 + k] * b3.span()[k * 4 + j]; ref.span()[i * 4 + j] = s + c2.span()[i * 4 + j]; } auto mm = eng.matmul(upload(eng, a2), upload(eng, b3), false, false); - if (!mm) { std::cerr << "[FAIL] matmul 原语\n"; ++fail; } + if (!mm) + { + std::cerr << "[FAIL] matmul 原语\n"; + ++fail; + } else { const Tensor mmt = std::move(*mm); diff --git a/src/tensor_precision_test.cpp b/src/tensor_precision_test.cpp index e5ea6c97..ca4e3819 100644 --- a/src/tensor_precision_test.cpp +++ b/src/tensor_precision_test.cpp @@ -85,8 +85,7 @@ void test_f16_cpu_tensor(nn::ComputeEngine& eng) // 经 F16 槽位下载（降 cast 再升位）语义：仍是精确往返 auto back16_r = eng.to_matrix(t, nn::Precision::F16); CHECK(static_cast(back16_r), "to_matrix P=F16 succeeds"); - if (back16_r) - CHECK(back16_r->at(0, 0) == 1.5f, "f16 data[0,0] via P=F16 download"); + if (back16_r) CHECK(back16_r->at(0, 0) == 1.5f, "f16 data[0,0] via P=F16 download"); std::printf(" done\n"); } @@ -143,15 +142,15 @@ void test_create_tensor_factory(nn::ComputeEngine& eng) { auto r = eng.try_create_tensor(2, 2, P); CHECK(!r, "try_create_tensor 拒绝非存储精度"); - if (!r) - CHECK(!r.error().message.empty(), "拒绝原因非空（结构化错误）"); + if (!r) CHECK(!r.error().message.empty(), "拒绝原因非空（结构化错误）"); } // ② 形状溢出 → 拒绝（rows*cols 超过 size_t），而不是回绕后按错误元素数分配 auto ovf = eng.try_create_tensor(std::numeric_limits::max() / 2 + 1, 4); CHECK(!ovf, "try_create_tensor 拒绝形状溢出"); // ③ 正常请求仍返回可用张量（Result 成功路径） auto ok_r = eng.try_create_tensor(2, 3, nn::Precision::F32); - CHECK(static_cast(ok_r) && ok_r->valid() && ok_r->rows() == 2, "try_create_tensor 成功路径"); + CHECK(static_cast(ok_r) && ok_r->valid() && ok_r->rows() == 2, + "try_create_tensor 成功路径"); // ④ InitSpec 重载同协议 auto z_r = eng.try_create_tensor(2, 3, nn::Precision::F32, nn::InitSpec::zero()); CHECK(static_cast(z_r) && z_r->valid(), "try_create_tensor(InitSpec) 成功路径"); @@ -225,8 +224,7 @@ void test_f16_data_roundtrip(nn::ComputeEngine& eng) bool exact = true; for (std::size_t r = 0; r < 2 && exact; ++r) for (std::size_t c = 0; c < 2; ++c) - if (back.at(r, c) != static_cast(m.at(r, c))) - exact = false; + if (back.at(r, c) != static_cast(m.at(r, c))) exact = false; CHECK(exact, "f16 roundtrip bit-exact via f32 download"); std::printf(" done\n"); diff --git a/src/tensor_schema_test.cpp b/src/tensor_schema_test.cpp index cdb29bb7..139d9d5b 100644 --- a/src/tensor_schema_test.cpp +++ b/src/tensor_schema_test.cpp @@ -52,8 +52,7 @@ nn::TensorSchema output_schema(bool value_optional = false) return std::move(*r); } -nn::Tensor make_f16_tensor(nn::ComputeEngine& engine, - std::size_t rows, std::size_t cols) +nn::Tensor make_f16_tensor(nn::ComputeEngine& engine, std::size_t rows, std::size_t cols) { nn::Matrix host(rows, cols); auto r = engine.from_matrix(host, nn::Precision::F16); @@ -74,12 +73,10 @@ void test_schema_and_handles() CHECK(p && v, "field name resolves"); if (!p || !v) return; CHECK(p->index == 0 && v->index == 1, "field handle indexes"); - CHECK(p->schema_fingerprint == schema.fingerprint(), - "field handle carries schema fingerprint"); + CHECK(p->schema_fingerprint == schema.fingerprint(), "field handle carries schema fingerprint"); CHECK(schema.field(*p).has_value(), "handle lookup succeeds"); CHECK(!schema.field("missing").has_value(), "unknown field rejected"); - CHECK(!schema.optional_field("missing")->valid(), - "optional lookup reports missing"); + CHECK(!schema.optional_field("missing")->valid(), "optional lookup reports missing"); const nn::TensorSchema same = output_schema(); CHECK(same.fingerprint() == schema.fingerprint(), @@ -92,32 +89,22 @@ void test_schema_and_handles() auto other_r = nn::TensorSchema::create("policy_value", std::move(changed)); CHECK(other_r.has_value(), "reordered schema builds"); if (!other_r) return; - CHECK(other_r->fingerprint() != schema.fingerprint(), - "field order changes fingerprint"); + CHECK(other_r->fingerprint() != schema.fingerprint(), "field order changes fingerprint"); auto foreign = other_r->field("policy"); CHECK(foreign.has_value(), "foreign field resolves there"); - if (foreign) - CHECK(!schema.field(*foreign).has_value(), - "cross-schema handle rejected"); + if (foreign) CHECK(!schema.field(*foreign).has_value(), "cross-schema handle rejected"); nn::FieldHandle bad{schema.fingerprint(), 99}; CHECK(!schema.field(bad).has_value(), "out-of-range handle rejected"); - CHECK(!schema.field(nn::FieldHandle{}).has_value(), - "uninitialized handle rejected"); + CHECK(!schema.field(nn::FieldHandle{}).has_value(), "uninitialized handle rejected"); } void test_schema_construction_errors() { - CHECK(!nn::TensorSchema::create( - "", {simple_field("x")}), - "empty schema name rejected"); - CHECK(!nn::TensorSchema::create("empty", {}), - "empty schema rejected"); - CHECK(!nn::TensorSchema::create( - "blank", {simple_field("")}), - "empty field name rejected"); - CHECK(!nn::TensorSchema::create( - "dup", {simple_field("x"), simple_field("x")}), + CHECK(!nn::TensorSchema::create("", {simple_field("x")}), "empty schema name rejected"); + CHECK(!nn::TensorSchema::create("empty", {}), "empty schema rejected"); + CHECK(!nn::TensorSchema::create("blank", {simple_field("")}), "empty field name rejected"); + CHECK(!nn::TensorSchema::create("dup", {simple_field("x"), simple_field("x")}), "duplicate field name rejected"); } @@ -142,10 +129,9 @@ void test_named_tensor_set() nn::Tensor policy = make_tensor(*owner, 4, 3, 0.1f, 0.01f); nn::Tensor value = make_tensor(*owner, 1, 3, 0.2f, 0.01f); - auto set_r = nn::NamedTensorSet::create( - schema, - {nn::NamedTensor{.name = "value", .value = value}, - nn::NamedTensor{.name = "policy", .value = policy}}); + auto set_r = + nn::NamedTensorSet::create(schema, {nn::NamedTensor{.name = "value", .value = value}, + nn::NamedTensor{.name = "policy", .value = policy}}); CHECK(set_r.has_value(), "named set accepts out-of-order input"); if (!set_r) return; CHECK(set_r->size() == 2, "named set size"); @@ -156,42 +142,30 @@ void test_named_tensor_set() if (!p_h || !v_h) return; auto p = set_r->require(*p_h); auto v = set_r->require(*v_h); - CHECK(p && (*p)->rows() == policy.rows() && - (*p)->cols() == policy.cols(), + CHECK(p && (*p)->rows() == policy.rows() && (*p)->cols() == policy.cols(), "policy handle returns original tensor"); - CHECK(v && (*v)->rows() == value.rows() && - (*v)->cols() == value.cols(), + CHECK(v && (*v)->rows() == value.rows() && (*v)->cols() == value.cols(), "value handle returns original tensor"); CHECK(set_r->find("policy") == *p, "find by name"); CHECK(set_r->find(*v_h) == *v, "find by handle"); - CHECK(!nn::NamedTensorSet::create( - schema, - {nn::NamedTensor{.name = "policy", .value = policy}}), + CHECK(!nn::NamedTensorSet::create(schema, {nn::NamedTensor{.name = "policy", .value = policy}}), "missing required field rejected"); - CHECK(!nn::NamedTensorSet::create( - schema, - {nn::NamedTensor{.name = "policy", .value = policy}, - nn::NamedTensor{.name = "policy", .value = policy}, - nn::NamedTensor{.name = "value", .value = value}}), + CHECK(!nn::NamedTensorSet::create(schema, {nn::NamedTensor{.name = "policy", .value = policy}, + nn::NamedTensor{.name = "policy", .value = policy}, + nn::NamedTensor{.name = "value", .value = value}}), "duplicate named tensor rejected"); - CHECK(!nn::NamedTensorSet::create( - schema, - {nn::NamedTensor{.name = "policy", .value = policy}, - nn::NamedTensor{.name = "value", .value = value}, - nn::NamedTensor{.name = "extra", .value = value}}), + CHECK(!nn::NamedTensorSet::create(schema, {nn::NamedTensor{.name = "policy", .value = policy}, + nn::NamedTensor{.name = "value", .value = value}, + nn::NamedTensor{.name = "extra", .value = value}}), "unknown named tensor rejected"); CHECK(!nn::NamedTensorSet::create( - schema, - {nn::NamedTensor{.name = "policy", .value = policy}, - nn::NamedTensor{.name = "value", - .value = make_tensor(*owner, 2, 3)}}), + schema, {nn::NamedTensor{.name = "policy", .value = policy}, + nn::NamedTensor{.name = "value", .value = make_tensor(*owner, 2, 3)}}), "shape constraint enforced"); CHECK(!nn::NamedTensorSet::create( - schema, - {nn::NamedTensor{.name = "policy", - .value = make_f16_tensor(*owner, 4, 3)}, - nn::NamedTensor{.name = "value", .value = value}}), + schema, {nn::NamedTensor{.name = "policy", .value = make_f16_tensor(*owner, 4, 3)}, + nn::NamedTensor{.name = "value", .value = value}}), "precision constraint enforced"); } @@ -201,15 +175,14 @@ void test_optional_and_gradients() const nn::TensorSchema schema = output_schema(/*value_optional=*/true); nn::Tensor policy = make_tensor(*owner, 5, 2, 0.3f, 0.01f); - auto only_policy = nn::NamedTensorSet::create( - schema, {nn::NamedTensor{.name = "policy", .value = policy}}); + auto only_policy = + nn::NamedTensorSet::create(schema, {nn::NamedTensor{.name = "policy", .value = policy}}); CHECK(only_policy.has_value(), "optional field may be absent"); if (only_policy) - CHECK(only_policy->find("value") == nullptr, - "absent optional field has no tensor"); + CHECK(only_policy->find("value") == nullptr, "absent optional field has no tensor"); - auto grads = nn::GradientBundle::create( - schema, {nn::NamedTensor{.name = "policy", .value = policy}}); + auto grads = + nn::GradientBundle::create(schema, {nn::NamedTensor{.name = "policy", .value = policy}}); CHECK(grads.has_value(), "gradient bundle uses output schema"); if (!grads) return; auto gh = grads->field("policy"); @@ -217,8 +190,7 @@ void test_optional_and_gradients() if (gh) { auto g = grads->require(*gh); - CHECK(g && (*g)->rows() == policy.rows(), - "gradient bundle exposes tensor by handle"); + CHECK(g && (*g)->rows() == policy.rows(), "gradient bundle exposes tensor by handle"); } } diff --git a/src/ternary_gemm_test.cpp b/src/ternary_gemm_test.cpp index d90b3383..4bda189f 100644 --- a/src/ternary_gemm_test.cpp +++ b/src/ternary_gemm_test.cpp @@ -53,13 +53,12 @@ void test_layout() CHECK(bytes_per_row(512) == 128, "K=512 → 128 B/行 = 0.25 B/权重"); // 位序：第 j 个权重的码在 (word >> 2j) & 3 - const float w[16] = {1.0f, -1.0f, 0.0f, 0.4f, -0.4f, 1.0f, -1.0f, 0.0f, - 0.9f, -0.9f, 0.1f, -0.1f, 0.6f, -0.6f, 0.0f, 0.0f}; + const float w[16] = {1.0f, -1.0f, 0.0f, 0.4f, -0.4f, 1.0f, -1.0f, 0.0f, + 0.9f, -0.9f, 0.1f, -0.1f, 0.6f, -0.6f, 0.0f, 0.0f}; std::uint32_t word[1] = {0}; pack_row(w, 16, /*tau=*/1.0f, word); // τ=1 → 阈值 0.5：1.0 → +1(码3)、-1.0 → -1(码1)、|w|≤0.5 → 0(码2) - const std::uint32_t expect_codes[16] = {3, 1, 2, 2, 2, 3, 1, 2, - 3, 1, 2, 2, 3, 1, 2, 2}; + const std::uint32_t expect_codes[16] = {3, 1, 2, 2, 2, 3, 1, 2, 3, 1, 2, 2, 3, 1, 2, 2}; bool bits_ok = true; for (std::uint32_t j = 0; j < 16; ++j) if (((word[0] >> (2 * j)) & 3u) != expect_codes[j]) bits_ok = false; @@ -76,9 +75,8 @@ void test_layout() CHECK(roundtrip, "解码往返（尾部位 = 码 2，不是未定义码 0）"); // ── base-3（部署编码，P2）：5 digit/字节，digit = 码 - 1 ───────────── - CHECK(kBase3PerByte == 5 && base3_bytes_per_row(128) == 26 && - base3_bytes_per_row(5) == 1 && base3_bytes_per_row(6) == 2 && - base3_bytes_per_row(17) == 4, + CHECK(kBase3PerByte == 5 && base3_bytes_per_row(128) == 26 && base3_bytes_per_row(5) == 1 && + base3_bytes_per_row(6) == 2 && base3_bytes_per_row(17) == 4, "base-3 尺寸账：ceil(K/5)（128 → 26 B = 1.625 bpw）"); { // 覆盖三种码 + 长度非 5 倍数的尾部：base3 → 2-bit 容器必须与直接打包逐位一致 @@ -96,8 +94,7 @@ void test_layout() for (std::size_t i = 0; i < wpr * nn::trit_pack::kPerWord; ++i) { const std::uint32_t code = (i < K) ? w[i] : 2u; - ref[i / nn::trit_pack::kPerWord] |= - code << (2 * (i % nn::trit_pack::kPerWord)); + ref[i / nn::trit_pack::kPerWord] |= code << (2 * (i % nn::trit_pack::kPerWord)); } // 容器 → base-3 → 容器（部署落盘/读盘的宿主侧往返，不涉浮点） nn::trit_pack::pack_row_base3(ref.data(), K, b3.data()); @@ -105,8 +102,8 @@ void test_layout() if (back != ref) b3_ok = false; // 单点解码（参考读法）与真值一致 for (std::size_t i = 0; i < K; ++i) - if (nn::trit_pack::at_base3(b3.data(), i) != - nn::trit_pack::value_of(w[i])) b3_ok = false; + if (nn::trit_pack::at_base3(b3.data(), i) != nn::trit_pack::value_of(w[i])) + b3_ok = false; // 尾部字必须是"码 2 = 三值 0" for (std::size_t i = K; i < wpr * nn::trit_pack::kPerWord; ++i) if (nn::trit_pack::at(back.data(), i) != 0.0f) b3_ok = false; @@ -117,9 +114,8 @@ void test_layout() // ── 工具：宿主参考 GEMM（入参是**三值**（已量化），乘逐行 τ；transA=0/1）── // GEMM 消费的是量化后的三值权重 —— 参考必须用 trit 而不是原始 latent。 -void ref_gemm(const std::vector& trit, const std::vector& tau, - std::size_t rows, std::size_t k, const nn::Matrix& B, bool transA, - nn::Matrix& out) +void ref_gemm(const std::vector& trit, const std::vector& tau, std::size_t rows, + std::size_t k, const nn::Matrix& B, bool transA, nn::Matrix& out) { if (!transA) { @@ -183,12 +179,15 @@ void test_primitives(nn::ComputeEngine& eng, const char* tag) for (std::size_t r = 0; r < ROWS; ++r) { double s = 0.0; - for (std::size_t c = 0; c < K; ++c) s += std::fabs(static_cast(W.at(r, c))); + for (std::size_t c = 0; c < K; ++c) + s += std::fabs(static_cast(W.at(r, c))); tau[r] = static_cast(s / static_cast(K)); - for (std::size_t c = 0; c < K; ++c) w_flat[r * K + c] = W.at(r, c); + for (std::size_t c = 0; c < K; ++c) + w_flat[r * K + c] = W.at(r, c); } nn::Matrix tau_m(ROWS, 1); - for (std::size_t r = 0; r < ROWS; ++r) tau_m.set_value(r, 0, tau[r]); + for (std::size_t r = 0; r < ROWS; ++r) + tau_m.set_value(r, 0, tau[r]); nn::Tensor tau_t = upload(eng, tau_m); // 三值（宿主，作为后续参考与"逐位一致"的基准） @@ -203,8 +202,7 @@ void test_primitives(nn::ComputeEngine& eng, const char* tag) if (!packed_r.has_value()) return; nn::Tensor packed = std::move(*packed_r); const std::size_t words = nn::trit_pack::words_per_row(K); - CHECK(packed.rows() == ROWS && packed.cols() == words, - "打包容器形状 (rows, ceil(K/16))"); + CHECK(packed.rows() == ROWS && packed.cols() == words, "打包容器形状 (rows, ceil(K/16))"); { const nn::Matrix pm = download(eng, packed); bool same = true; @@ -246,7 +244,8 @@ void test_primitives(nn::ComputeEngine& eng, const char* tag) // 与"f16 稠密缓冲路径"（matmul(wq,·)·τ）等价 nn::Matrix wq_m(ROWS, K); for (std::size_t r = 0; r < ROWS; ++r) - for (std::size_t c = 0; c < K; ++c) wq_m.set_value(r, c, trit[r * K + c]); + for (std::size_t c = 0; c < K; ++c) + wq_m.set_value(r, c, trit[r * K + c]); nn::Tensor wq_t = upload(eng, wq_m, nn::Precision::F16); auto dense = eng.matmul(wq_t, Xt, false, false, nn::Precision::F32); CHECK(dense.has_value(), "稠密 wq matmul 成功"); @@ -256,9 +255,9 @@ void test_primitives(nn::ComputeEngine& eng, const char* tag) double dmax = 0.0; for (std::size_t r = 0; r < ROWS; ++r) for (std::size_t n = 0; n < N; ++n) - dmax = std::max(dmax, - std::fabs(static_cast(d.at(r, n)) * static_cast(tau[r]) - - static_cast(ref.at(r, n)))); + dmax = std::max(dmax, std::fabs(static_cast(d.at(r, n)) * + static_cast(tau[r]) - + static_cast(ref.at(r, n)))); CHECK(dmax < 1e-5, "f16 稠密路径（P1）同一口径（三值逐位一致）"); } } @@ -275,9 +274,11 @@ void test_primitives(nn::ComputeEngine& eng, const char* tag) a_flat[k * M + m] = 0.19f * static_cast((k * 5 + m * 3) % 7) - 0.5f; nn::Matrix A(K, M); for (std::size_t k = 0; k < K; ++k) - for (std::size_t m = 0; m < M; ++m) A.set_value(k, m, a_flat[k * M + m]); + for (std::size_t m = 0; m < M; ++m) + A.set_value(k, m, a_flat[k * M + m]); nn::Matrix tau_k_m(K, 1); - for (std::size_t k = 0; k < K; ++k) tau_k_m.set_value(k, 0, tau_k[k]); + for (std::size_t k = 0; k < K; ++k) + tau_k_m.set_value(k, 0, tau_k[k]); nn::Tensor At = upload(eng, A); nn::Tensor tau_kt = upload(eng, tau_k_m); auto pk2_r = eng.trit_pack(At, tau_kt); @@ -332,7 +333,8 @@ void test_bitlinear_paths(nn::ComputeEngine& eng, const char* tag) nn::PrecisionProfile prof; prof.param = nn::Precision::T1_58; - auto make = [&](bool packed) { + auto make = [&](bool packed) + { auto bl = std::make_unique(IN, OUT); bl->set_precision_profile(prof); auto r = bl->init(eng); @@ -364,7 +366,8 @@ void test_bitlinear_paths(nn::ComputeEngine& eng, const char* tag) } // 反向：dX 与 dW 两路对比（同一上游梯度） - auto make_grad = [&] { + auto make_grad = [&] + { nn::Matrix G(OUT, B); for (std::size_t o = 0; o < OUT; ++o) for (std::size_t c = 0; c < B; ++c) @@ -405,8 +408,8 @@ void test_bitlinear_paths(nn::ComputeEngine& eng, const char* tag) // 2. **尾块**：M/N 非 64 倍数、K 非 16 倍数（trit 越界必须判 0 = 三值 0）。 // 3. **单列**：N=1（自回归解码形态；列块大面积空转，网格与守卫都必须对）。 // 覆盖 transA=0/1 × f32/f16 存储四种组合。 -void check_gemm_shape(nn::ComputeEngine& eng, std::size_t M, std::size_t K, - std::size_t N, bool transA, bool f16_mode) +void check_gemm_shape(nn::ComputeEngine& eng, std::size_t M, std::size_t K, std::size_t N, + bool transA, bool f16_mode) { const std::size_t sr = transA ? K : M; // A 的存储行数 const std::size_t al = transA ? M : K; // A 的存储行长（= 一行三值个数） @@ -417,24 +420,27 @@ void check_gemm_shape(nn::ComputeEngine& eng, std::size_t M, std::size_t K, a_flat[r * al + c] = 0.17f * static_cast((r * 7 + c * 3) % 13) - 0.9f; nn::Matrix A_lat(sr, al); for (std::size_t r = 0; r < sr; ++r) - for (std::size_t c = 0; c < al; ++c) A_lat.set_value(r, c, a_flat[r * al + c]); + for (std::size_t c = 0; c < al; ++c) + A_lat.set_value(r, c, a_flat[r * al + c]); // τ = 逐存储行 absmean（宿主；DSL 那条归约由 t1_58_test 覆盖） std::vector tau(sr, 0.0f); for (std::size_t r = 0; r < sr; ++r) { double s = 0.0; - for (std::size_t c = 0; c < al; ++c) s += std::fabs(a_flat[r * al + c]); + for (std::size_t c = 0; c < al; ++c) + s += std::fabs(a_flat[r * al + c]); tau[r] = static_cast(s / static_cast(al)); } nn::Matrix tau_m(sr, 1); - for (std::size_t r = 0; r < sr; ++r) tau_m.set_value(r, 0, tau[r]); + for (std::size_t r = 0; r < sr; ++r) + tau_m.set_value(r, 0, tau[r]); std::vector trit(sr * al, 0.0f); for (std::size_t r = 0; r < sr; ++r) for (std::size_t c = 0; c < al; ++c) - trit[r * al + c] = nn::trit_pack::value_of( - nn::trit_pack::code_of(a_flat[r * al + c], tau[r])); + trit[r * al + c] = + nn::trit_pack::value_of(nn::trit_pack::code_of(a_flat[r * al + c], tau[r])); // B (K,N)：f16 存储时宿主参考必须用**同一批** f16 舍入值，否则容差里混进 // "输入表示差"（既会把真实错误判成噪声、也会把噪声判成错误）。 @@ -463,8 +469,8 @@ void check_gemm_shape(nn::ComputeEngine& eng, std::size_t M, std::size_t K, CHECK(packed_r.has_value(), "跨形状：trit_pack 成功"); if (!packed_r.has_value()) return; - auto got_r = transA ? eng.matmul_q_t(*packed_r, tau_t, Bt, M, P) - : eng.matmul_q(*packed_r, tau_t, Bt, P); + auto got_r = + transA ? eng.matmul_q_t(*packed_r, tau_t, Bt, M, P) : eng.matmul_q(*packed_r, tau_t, Bt, P); CHECK(got_r.has_value(), "跨形状：matmul_q(_t) 成功"); if (!got_r.has_value()) return; CHECK(got_r->rows() == M && got_r->cols() == N, "跨形状：输出形状 (M,N)"); @@ -481,8 +487,8 @@ void check_gemm_shape(nn::ComputeEngine& eng, std::size_t M, std::size_t K, // 容差按**全局幅度**定标：三值点积可正负相消，逐元素相对判据在近零点爆炸 const double scale = std::max(1.0, rmax); const double tol = f16_mode ? 4e-3 : 1e-4; - std::printf(" M=%-5zu K=%-5zu N=%-4zu transA=%u %-3s max_err=%.3e\n", - M, K, N, transA ? 1u : 0u, f16_mode ? "f16" : "f32", emax); + std::printf(" M=%-5zu K=%-5zu N=%-4zu transA=%u %-3s max_err=%.3e\n", M, K, N, + transA ? 1u : 0u, f16_mode ? "f16" : "f32", emax); CHECK(emax <= tol * scale, "跨形状：matmul_q(_t) = 宿主参考"); } @@ -494,7 +500,8 @@ void test_shapes(nn::ComputeEngine& eng, const char* tag, bool gpu) std::puts(" （引擎无打包能力 → 跳过）"); return; } - struct Sh { std::size_t M, K, N; }; + struct Sh + { std::size_t M, K, N; }; const Sh base[] = { {64, 16, 64}, // 单块方形（旧极小形状 = 掩盖轴向缺陷的那一类） {128, 32, 64}, // 行 2 块 / 列 1 块 @@ -529,8 +536,8 @@ void test_shapes(nn::ComputeEngine& eng, const char* tag, bool gpu) // [5] 用的是 OUT=9/IN=24/B=4（两向都只有 1 块）→ 天然测不出分块/网格缺陷。 // 这里把形状撑过块边界，并**同时覆盖 f32 与 f16 compute**（后者才是训练路径： // 2B4T 形状的 GPT 训练发散正发生在 f16 compute 的打包内核上）。 -void check_bitlinear_shape(nn::ComputeEngine& eng, std::size_t IN, std::size_t OUT, - std::size_t B, const nn::PrecisionProfile& prof) +void check_bitlinear_shape(nn::ComputeEngine& eng, std::size_t IN, std::size_t OUT, std::size_t B, + const nn::PrecisionProfile& prof) { nn::Matrix W(OUT, IN); for (std::size_t o = 0; o < OUT; ++o) @@ -543,7 +550,8 @@ void check_bitlinear_shape(nn::ComputeEngine& eng, std::size_t IN, std::size_t O const nn::Precision P = (prof.compute == nn::Precision::F16) ? nn::Precision::F16 : nn::Precision::F32; - auto make = [&](bool packed) { + auto make = [&](bool packed) + { auto bl = std::make_unique(IN, OUT); bl->set_precision_profile(prof); auto r = bl->init(eng); @@ -563,7 +571,8 @@ void check_bitlinear_shape(nn::ComputeEngine& eng, std::size_t IN, std::size_t O CHECK(y_pk.has_value() && y_dn.has_value(), "两条路径 forward 都成功"); if (!y_pk.has_value() || !y_dn.has_value()) return; - auto cmp = [&](const nn::Tensor& a, const nn::Tensor& b, const char* what) { + auto cmp = [&](const nn::Tensor& a, const nn::Tensor& b, const char* what) + { const nn::Matrix am = download(eng, a), bm = download(eng, b); const auto pa = am.span(); const auto pb = bm.span(); @@ -575,8 +584,7 @@ void check_bitlinear_shape(nn::ComputeEngine& eng, std::size_t IN, std::size_t O } const double scale = std::max(1.0, rmax); const double tol = (P == nn::Precision::F16) ? 4e-3 : 1e-5; - std::printf(" [%s] IN=%-4zu OUT=%-4zu B=%-4zu max_err=%.3e\n", - what, IN, OUT, B, emax); + std::printf(" [%s] IN=%-4zu OUT=%-4zu B=%-4zu max_err=%.3e\n", what, IN, OUT, B, emax); CHECK(emax <= tol * scale, "BitLinear 跨形状：打包 = 稠密"); }; cmp(*y_pk, *y_dn, "fwd"); @@ -605,7 +613,8 @@ void test_bitlinear_shapes(nn::ComputeEngine& eng, const char* tag) p32.param = nn::Precision::T1_58; const nn::PrecisionProfile profs[2] = {p32, nn::profile_t1_58()}; const char* ptags[2] = {"t1_58/f32", "t1_58/f16"}; - struct Sh { std::size_t IN, OUT, B; }; + struct Sh + { std::size_t IN, OUT, B; }; const Sh shapes[] = { {70, 150, 100}, // 两向都多块 + 尾块 {256, 64, 9}, // 窄 N 刚过分派门（9 > 8 → 分块内核）+ K 整块 @@ -659,7 +668,8 @@ void test_dequant() for (const std::vector& w : rows) { double acc = 0.0; - for (float v : w) acc += std::fabs(static_cast(v)); + for (float v : w) + acc += std::fabs(static_cast(v)); const float tau = static_cast(acc / static_cast(K)); std::vector words(wpr, 0u); pack_row(w.data(), K, tau, words.data()); @@ -676,7 +686,8 @@ void test_dequant() if (again != words) repack_ok = false; // ③ 逐行 absmean(L) == τ（浮点链只有 ±1 ulp；这里给 2e-6 相对带） double l1 = 0.0; - for (float v : lat) l1 += std::fabs(static_cast(v)); + for (float v : lat) + l1 += std::fabs(static_cast(v)); const double mean = l1 / static_cast(K); const double rel = std::fabs(mean - static_cast(tau)) / std::max(1e-30, std::fabs(static_cast(tau))); @@ -728,7 +739,6 @@ int main(int argc, char* argv[]) test_bitlinear_shapes(eng, "CPU"); } - if (g_failures == 0) - std::puts("ternary_gemm_test: 全部通过"); + if (g_failures == 0) std::puts("ternary_gemm_test: 全部通过"); return g_failures == 0 ? 0 : 1; } diff --git a/src/test_common.hpp b/src/test_common.hpp index 478fdb0f..ad9f0cd7 100644 --- a/src/test_common.hpp +++ b/src/test_common.hpp @@ -38,12 +38,14 @@ #define NN_TEST_COUNTER g_fail #endif -#define CHECK(cond, msg) \ - do { \ - if (!(cond)) { \ - std::printf("[FAIL] %s:%d: %s\n", __FILE__, __LINE__, (msg)); \ - ++NN_TEST_COUNTER; \ - } \ +#define CHECK(cond, msg) \ + do \ + { \ + if (!(cond)) \ + { \ + std::printf("[FAIL] %s:%d: %s\n", __FILE__, __LINE__, (msg)); \ + ++NN_TEST_COUNTER; \ + } \ } while (0) // ── M1 访问收口（docs/internal/17 §4.1）：库外一律经引擎上传/下载 ─────── @@ -61,8 +63,7 @@ inline nn::Tensor upload(nn::ComputeEngine& e, const nn::Matrix& m, } // f16 宿主矩阵上传（f16→f32 逐元素升位后按 F16 落盘：往返精确无损） -inline nn::Tensor upload(nn::ComputeEngine& e, - const nn::MatrixT& m) +inline nn::Tensor upload(nn::ComputeEngine& e, const nn::MatrixT& m) { nn::Matrix m32(m.rows(), m.cols()); const auto s = m.span(); @@ -80,9 +81,8 @@ inline nn::Matrix download(nn::ComputeEngine& e, const nn::Tensor& t) } // ── 构造确定性测试张量（base + i*step 填充；经引擎上传，出生即绑定）──────── -inline nn::Tensor make_tensor(nn::ComputeEngine& e, std::size_t rows, - std::size_t cols, float base = 0.0f, - float step = 0.01f) +inline nn::Tensor make_tensor(nn::ComputeEngine& e, std::size_t rows, std::size_t cols, + float base = 0.0f, float step = 0.01f) { nn::Matrix m(rows, cols); auto sp = m.span(); @@ -92,13 +92,13 @@ inline nn::Tensor make_tensor(nn::ComputeEngine& e, std::size_t rows, } // ── 逐元素比对 Tensor（容差 1e-4 相对；失败计数进 fails = 调用方的 g_fail）── -inline void check_close(nn::ComputeEngine& e, const nn::Tensor& got, - const nn::Tensor& ref, const char* msg, int& fails) +inline void check_close(nn::ComputeEngine& e, const nn::Tensor& got, const nn::Tensor& ref, + const char* msg, int& fails) { if (got.rows() != ref.rows() || got.cols() != ref.cols()) { - std::printf("[FAIL] %s: shape mismatch got %zux%zu ref %zux%zu\n", - msg, got.rows(), got.cols(), ref.rows(), ref.cols()); + std::printf("[FAIL] %s: shape mismatch got %zux%zu ref %zux%zu\n", msg, got.rows(), + got.cols(), ref.rows(), ref.cols()); ++fails; return; } @@ -119,15 +119,12 @@ inline void check_close(nn::ComputeEngine& e, const nn::Tensor& got, } } if (ok) std::printf("[ OK ] %s\n", msg); - else ++fails; + else ++fails; } // ── 数值 vs 解析梯度的相对容差判据（gradcheck 通用）─────────────────────── inline bool approx(nn::Scalar num, nn::Scalar ana, nn::Scalar tol) -{ - return std::fabs(num - ana) <= - tol * (nn::Scalar{1} + std::fabs(num) + std::fabs(ana)); -} +{ return std::fabs(num - ana) <= tol * (nn::Scalar{1} + std::fabs(num) + std::fabs(ana)); } // ── 矩阵点积（gradcheck 计算参考损失用）─────────────────────────────────── inline nn::Scalar dot(const nn::Matrix& a, const nn::Matrix& b) @@ -144,8 +141,7 @@ inline nn::Scalar dot(const nn::Matrix& a, const nn::Matrix& b) inline bool close_to(const nn::Matrix& a, const nn::Matrix& b, nn::Scalar tol, const std::string& name, std::size_t idx) { - NN_ASSERT_ALWAYS(a.rows() == b.rows() && a.cols() == b.cols(), - "close_to: shape mismatch"); + NN_ASSERT_ALWAYS(a.rows() == b.rows() && a.cols() == b.cols(), "close_to: shape mismatch"); nn::Scalar max_abs = 0; nn::Scalar max_rel = 0; const auto& sa = a.span(); @@ -159,8 +155,7 @@ inline bool close_to(const nn::Matrix& a, const nn::Matrix& b, nn::Scalar tol, if (rel > max_rel) max_rel = rel; } const bool pass = (max_abs <= tol) || (max_rel <= tol); - std::cout << " [" << idx << "] " << name - << " max_abs=" << max_abs << " max_rel=" << max_rel - << (pass ? " OK" : " FAIL") << "\n"; + std::cout << " [" << idx << "] " << name << " max_abs=" << max_abs + << " max_rel=" << max_rel << (pass ? " OK" : " FAIL") << "\n"; return pass; } diff --git a/src/text_infer.cpp b/src/text_infer.cpp index 9b5d6dfb..2aa7f8fe 100644 --- a/src/text_infer.cpp +++ b/src/text_infer.cpp @@ -26,7 +26,7 @@ using nn::Scalar; // ==================== 帮助信息 ==================== -void print_usage(const char *prog) +void print_usage(const char* prog) { nn::cli::Help help(std::cout, prog, "GPT 文本生成推理程序"); @@ -36,11 +36,17 @@ void print_usage(const char *prog) "多轮对话（历史由调用方自带，一次一问）"); help.section("选项"); - help.opt("--model ", "模型文件路径 (默认: gpt_model.bin)\nV3 格式模型自动读取规格和嵌入 tokenizer"); - help.opt("--vocab ", "词表 .nnvocab 路径 (默认: bpe_vocab.nnvocab)\n仅当模型未嵌入 tokenizer 时使用；自动识别分词器类型 (bpe / charbpe)"); + help.opt("--model ", + "模型文件路径 (默认: gpt_model.bin)\nV3 格式模型自动读取规格和嵌入 tokenizer"); + help.opt("--vocab ", "词表 .nnvocab 路径 (默认: bpe_vocab.nnvocab)\n仅当模型未嵌入 " + "tokenizer 时使用；自动识别分词器类型 (bpe / charbpe)"); help.opt("--prompt ", "输入提示文本"); - help.opt("--turn :", "追加一轮对话历史；role = user|assistant，可重复。\n给了 --turn 即进入多轮对话模式：按词表的对话标记拼 prompt\n（末条必须是 user），并从最旧的一问一答开始截断到模型 seq_len"); - help.opt("--reply-out ", "多轮模式：把本轮回复（纯文本）另写一份到该文件\n（GUI 靠它取回复，无需解析日志）"); + help.opt("--turn :", + "追加一轮对话历史；role = user|assistant，可重复。\n给了 --turn " + "即进入多轮对话模式：按词表的对话标记拼 prompt\n（末条必须是 " + "user），并从最旧的一问一答开始截断到模型 seq_len"); + help.opt("--reply-out ", + "多轮模式：把本轮回复（纯文本）另写一份到该文件\n（GUI 靠它取回复，无需解析日志）"); help.opt("--system ", "对话模式的系统提示（默认: 你是一个有用的AI助手。）"); help.opt("--interactive", "交互式生成模式（多轮，上下文累积保留）"); help.opt("--max-tokens ", "最大生成 token 数 (默认: 200)"); @@ -51,7 +57,9 @@ void print_usage(const char *prog) help.opt("--repeat-window ", "重复惩罚只看最近 n 个 token (默认: 0=全部上下文)"); help.opt("--seed ", "随机种子 (给了就固定种子 ⇒ 同 prompt 输出可复现)"); help.opt("--show-tokens", "显示 token ID (调试用)"); - help.opt("--gpu [索引|名称]", "启用 GPU 加速 (需要 Vulkan SDK)\n空格形式只收枚举索引；名称子串用 --gpu=<名称>，如 --gpu=NVIDIA / --gpu=40HX"); + help.opt("--gpu [索引|名称]", + "启用 GPU 加速 (需要 Vulkan SDK)\n空格形式只收枚举索引；名称子串用 --gpu=<名称>，如 " + "--gpu=NVIDIA / --gpu=40HX"); help.opt("--help, -h", "显示此帮助信息"); } @@ -79,7 +87,7 @@ struct InferConfig std::string gpu_device; // --gpu 的可选设备选择子（空 = 自动选卡） }; -InferConfig parse_args(int argc, char *argv[]) +InferConfig parse_args(int argc, char* argv[]) { InferConfig cfg; for (int i = 1; i < argc; ++i) @@ -90,19 +98,15 @@ InferConfig parse_args(int argc, char *argv[]) print_usage(argv[0]); std::exit(0); } - else if (arg == "--model" && i + 1 < argc) - cfg.model_path = argv[++i]; - else if (arg == "--vocab" && i + 1 < argc) - cfg.vocab_path = argv[++i]; + else if (arg == "--model" && i + 1 < argc) cfg.model_path = argv[++i]; + else if (arg == "--vocab" && i + 1 < argc) cfg.vocab_path = argv[++i]; else if (arg == "--prompt" && i + 1 < argc) { cfg.prompt = argv[++i]; cfg.prompt_given = true; } - else if (arg == "--system" && i + 1 < argc) - cfg.system_prompt = argv[++i]; - else if (arg == "--reply-out" && i + 1 < argc) - cfg.reply_out = argv[++i]; + else if (arg == "--system" && i + 1 < argc) cfg.system_prompt = argv[++i]; + else if (arg == "--reply-out" && i + 1 < argc) cfg.reply_out = argv[++i]; else if (arg == "--turn" && i + 1 < argc) { // :：取**第一个**冒号分段（正文里的冒号不转义） @@ -120,19 +124,30 @@ InferConfig parse_args(int argc, char *argv[]) t.text = val.substr(sep + 1); cfg.turns.push_back(std::move(t)); } - else if (arg == "--interactive") - cfg.interactive = true; + else if (arg == "--interactive") cfg.interactive = true; else if (arg == "--max-tokens" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --max-tokens\n"; std::exit(1); } - if (*v < 1) { std::cerr << "--max-tokens 必须 >= 1\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --max-tokens\n"; + std::exit(1); + } + if (*v < 1) + { + std::cerr << "--max-tokens 必须 >= 1\n"; + std::exit(1); + } cfg.max_tokens = *v; } else if (arg == "--temperature" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --temperature\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --temperature\n"; + std::exit(1); + } cfg.temperature = *v; } else if (arg == "--top-k" && i + 1 < argc) @@ -191,8 +206,7 @@ InferConfig parse_args(int argc, char *argv[]) cfg.gpu_enabled = true; if (!gpu_dev->empty()) cfg.gpu_device = *gpu_dev; } - else if (arg == "--show-tokens") - cfg.show_tokens = true; + else if (arg == "--show-tokens") cfg.show_tokens = true; else if (!arg.starts_with("--")) { cfg.prompt = arg; @@ -229,8 +243,8 @@ InferConfig parse_args(int argc, char *argv[]) const bool expect_user = (i % 2 == 0); if (cfg.turns[i].user != expect_user) { - std::cerr << "--turn 角色必须 user/assistant 交替（第 " << (i + 1) - << " 条应为 " << (expect_user ? "user" : "assistant") << "）\n"; + std::cerr << "--turn 角色必须 user/assistant 交替（第 " << (i + 1) << " 条应为 " + << (expect_user ? "user" : "assistant") << "）\n"; std::exit(1); } } @@ -269,15 +283,15 @@ nn::Result> generate_text(nn::Model& model, nn::Compute const nn::SamplingConfig& sc, std::size_t eos_token_id) { - auto &layer_ref = model.layer_at(0); + auto& layer_ref = model.layer_at(0); // 传入 EOS_ID，生成遇到 EOS 自动停止 // min_new_tokens = max_new_tokens/2，至少生成一半 token 才允许 EOS 停止， // 避免模型因训练偏置一上来就输出 EOS 导致无输出。 const std::size_t min_new = max_new_tokens / 2; - if (auto *gpt_ptr = dynamic_cast(&layer_ref)) + if (auto* gpt_ptr = dynamic_cast(&layer_ref)) return gpt_ptr->generate(engine, prompt_tokens, max_new_tokens, sc, eos_token_id, min_new); - if (auto *rapt_ptr = dynamic_cast(&layer_ref)) + if (auto* rapt_ptr = dynamic_cast(&layer_ref)) return rapt_ptr->generate(engine, prompt_tokens, max_new_tokens, sc, eos_token_id, min_new); NN_FAIL("Model does not contain a GPTModel or RAPTModel layer"); } @@ -286,43 +300,35 @@ nn::Result> generate_text(nn::Model& model, nn::Compute // 多轮（P0-5 尾项）：每轮把**本轮及之前的所有轮次**拼进 prompt（`chat_prompt`）， // 因此连续 N 轮不丢上下文；超 `seq_len` 时从最旧的一问一答开始整轮丢弃。 // 助手轮用**生成的原 token** 回灌（不经 decode → encode 往返）。 -void interactive_mode(nn::Model &model, nn::ComputeEngine &engine, - const nn::Tokenizer &tokenizer, const InferConfig &cfg, - std::size_t seq_len) +void interactive_mode(nn::Model& model, nn::ComputeEngine& engine, const nn::Tokenizer& tokenizer, + const InferConfig& cfg, std::size_t seq_len) { const std::size_t eos_id = tokenizer.eos_id(); const bool has_dialogue = tokenizer.has_dialogue_markers(); std::vector history; // 交替：user → assistant → … - if (has_dialogue) - std::cout << "GPT 交互式对话（多轮，上下文保留；输入 'quit' 退出）\n\n"; - else - std::cout << "GPT 交互式生成（多轮，上下文保留；输入 'quit' 退出）\n\n"; + if (has_dialogue) std::cout << "GPT 交互式对话（多轮，上下文保留；输入 'quit' 退出）\n\n"; + else std::cout << "GPT 交互式生成（多轮，上下文保留；输入 'quit' 退出）\n\n"; while (true) { std::cout << ">>> "; std::string prompt; - if (!std::getline(std::cin, prompt)) - break; - if (prompt == "quit" || prompt == "exit") - break; - if (prompt.empty()) - prompt = " "; + if (!std::getline(std::cin, prompt)) break; + if (prompt == "quit" || prompt == "exit") break; + if (prompt.empty()) prompt = " "; // 本轮用户输入入列 → 整个历史（含之前 N-1 轮）一起进 prompt history.push_back(nn::Tokenizer::ChatTurn{true, prompt, {}}); std::size_t dropped = 0; - auto prompt_tokens = tokenizer.chat_prompt(history, cfg.system_prompt, seq_len, - &dropped); + auto prompt_tokens = tokenizer.chat_prompt(history, cfg.system_prompt, seq_len, &dropped); if (dropped > 0) { // 历史按"一问一答"整轮丢弃（丢弃后仍以 user 起头） - history.erase(history.begin(), - history.begin() + static_cast(dropped)); - std::cout << "[上下文] 已达 seq_len=" << seq_len << " 上限：丢弃最旧 " - << dropped << " 轮，当前保留 " << history.size() << " 轮\n"; + history.erase(history.begin(), history.begin() + static_cast(dropped)); + std::cout << "[上下文] 已达 seq_len=" << seq_len << " 上限：丢弃最旧 " << dropped + << " 轮，当前保留 " << history.size() << " 轮\n"; } std::cout << "生成中...\n"; @@ -343,7 +349,11 @@ void interactive_mode(nn::Model &model, nn::ComputeEngine &engine, const std::size_t end_asst = tokenizer.end_assistant_marker_id(); if (end_asst != nn::Tokenizer::npos) for (std::size_t i = 0; i < shown.size(); ++i) - if (shown[i] == end_asst) { shown.resize(i); break; } + if (shown[i] == end_asst) + { + shown.resize(i); + break; + } } std::cout << "\n" << tokenizer.decode(shown) << "\n\n"; @@ -370,7 +380,7 @@ void interactive_mode(nn::Model &model, nn::ComputeEngine &engine, } // ==================== 主函数 ==================== -int main(int argc, char *argv[]) +int main(int argc, char* argv[]) { InferConfig cfg = parse_args(argc, argv); @@ -382,25 +392,31 @@ int main(int argc, char *argv[]) NN_EXIT(role_r, 1, "读取模型文件失败: "); if (!spec.is_gpt() && !spec.is_rapt()) { - std::cerr << "模型文件不是 GPT/RAPT 类型 (type=" - << spec.type_name() << ")\n"; + std::cerr << "模型文件不是 GPT/RAPT 类型 (type=" << spec.type_name() << ")\n"; return 1; } // GPT/RAPT 共享的规格视图（两者类型参数同名，按实际类型取值） - std::size_t v_vocab = 0, v_d_model = 0, v_heads = 0, v_layers = 0, v_d_ff = 0, - v_seq_len = 0; + std::size_t v_vocab = 0, v_d_model = 0, v_heads = 0, v_layers = 0, v_d_ff = 0, v_seq_len = 0; nn::PosEncodingType v_pos = nn::PosEncodingType::Learned; - if (const auto *g = spec.as()) + if (const auto* g = spec.as()) { - v_vocab = g->vocab_size; v_d_model = g->d_model; v_heads = g->num_heads; - v_layers = g->num_layers; v_d_ff = g->d_ff; v_seq_len = g->seq_len; + v_vocab = g->vocab_size; + v_d_model = g->d_model; + v_heads = g->num_heads; + v_layers = g->num_layers; + v_d_ff = g->d_ff; + v_seq_len = g->seq_len; v_pos = g->pos_encoding; } - else if (const auto *r = spec.as()) + else if (const auto* r = spec.as()) { - v_vocab = r->vocab_size; v_d_model = r->d_model; v_heads = r->num_heads; - v_layers = r->num_layers; v_d_ff = r->d_ff; v_seq_len = r->seq_len; + v_vocab = r->vocab_size; + v_d_model = r->d_model; + v_heads = r->num_heads; + v_layers = r->num_layers; + v_d_ff = r->d_ff; + v_seq_len = r->seq_len; v_pos = r->pos_encoding; } @@ -413,29 +429,17 @@ int main(int argc, char *argv[]) auto engine = std::move(*engine_res); // ── 构建 GPT 模型 ──────────────────────────────────────── - std::cout << "模型规格: vocab=" << v_vocab - << " d_model=" << v_d_model - << " heads=" << v_heads - << " layers=" << v_layers - << " d_ff=" << v_d_ff - << " seq_len=" << v_seq_len; - if (spec.is_rapt()) - { - std::cout << " [RAPT (ReLU 线性注意力)]"; - } - if (spec.is_alibi_gpt() || v_pos == nn::PosEncodingType::ALiBi) - std::cout << " [ALiBi]"; - else if (v_pos == nn::PosEncodingType::Sinusoidal) - std::cout << " [Sinusoidal]"; - else if (v_pos == nn::PosEncodingType::RoPE) - std::cout << " [RoPE]"; - else - std::cout << " [Learned]"; + std::cout << "模型规格: vocab=" << v_vocab << " d_model=" << v_d_model << " heads=" << v_heads + << " layers=" << v_layers << " d_ff=" << v_d_ff << " seq_len=" << v_seq_len; + if (spec.is_rapt()) { std::cout << " [RAPT (ReLU 线性注意力)]"; } + if (spec.is_alibi_gpt() || v_pos == nn::PosEncodingType::ALiBi) std::cout << " [ALiBi]"; + else if (v_pos == nn::PosEncodingType::Sinusoidal) std::cout << " [Sinusoidal]"; + else if (v_pos == nn::PosEncodingType::RoPE) std::cout << " [RoPE]"; + else std::cout << " [Learned]"; std::cout << "\n"; nn::Result model_result; - if (spec.is_rapt()) - model_result = nn::build_rapt_model_from_spec(*engine, spec); + if (spec.is_rapt()) model_result = nn::build_rapt_model_from_spec(*engine, spec); else // 统一的 GPTModel 通过 pos_encoding 区分 Learned/Sinusoidal/ALiBi model_result = nn::build_gpt_model_from_spec(*engine, spec); NN_EXIT(model_result, 1, "构建模型失败: "); @@ -449,8 +453,7 @@ int main(int argc, char *argv[]) auto load_result = nn::load_model(cfg.model_path, model, &model_version); NN_EXIT(load_result, 1, "加载模型失败: "); embedded_vocab = std::move(*load_result); - std::cout << "模型已加载（角色: " << nn::model_role_name(*role_r) << "）" - << std::endl; + std::cout << "模型已加载（角色: " << nn::model_role_name(*role_r) << "）" << std::endl; } // ── 加载 tokenizer（按模型版本分派：v5=JSON / v6=.nnvocab kvrec）── @@ -471,7 +474,8 @@ int main(int argc, char *argv[]) std::cerr << "加载词表失败: " << cfg.vocab_path << '\n' << tok_r.error().message << "\n请使用 --vocab 指定 .nnvocab 词表路径，" - "或使用嵌入 tokenizer 的模型" << std::endl; + "或使用嵌入 tokenizer 的模型" + << std::endl; return 1; } tokenizer = std::move(*tok_r); @@ -494,8 +498,8 @@ int main(int argc, char *argv[]) if (!cfg.turns.empty()) { std::size_t dropped = 0; - auto prompt_tokens = tokenizer->chat_prompt(cfg.turns, cfg.system_prompt, seq_len, - &dropped); + auto prompt_tokens = + tokenizer->chat_prompt(cfg.turns, cfg.system_prompt, seq_len, &dropped); std::cout << "多轮对话: " << cfg.turns.size() << " 轮 (" << (tokenizer->has_dialogue_markers() ? "对话标记模板" : "纯文本拼接") << ")" << " → 进入上下文 " << (cfg.turns.size() - dropped) << " 轮"; @@ -528,7 +532,11 @@ int main(int argc, char *argv[]) const std::size_t end_asst = tokenizer->end_assistant_marker_id(); if (end_asst != nn::Tokenizer::npos) for (std::size_t i = 0; i < shown.size(); ++i) - if (shown[i] == end_asst) { shown.resize(i); break; } + if (shown[i] == end_asst) + { + shown.resize(i); + break; + } } const std::string reply = tokenizer->decode(shown); std::cout << reply << std::endl; @@ -551,11 +559,10 @@ int main(int argc, char *argv[]) } std::cout << "\n----------------------------------------" << std::endl; - std::cout << "生成 " << generated.size() << " tokens, 耗时 " - << std::fixed << std::setprecision(1) << gen_sec << "s" - << " (" << std::setprecision(0) - << (gen_sec > 0 ? generated.size() / gen_sec : 0) << " tokens/s)" - << std::endl; + std::cout << "生成 " << generated.size() << " tokens, 耗时 " << std::fixed + << std::setprecision(1) << gen_sec << "s" + << " (" << std::setprecision(0) << (gen_sec > 0 ? generated.size() / gen_sec : 0) + << " tokens/s)" << std::endl; // GUI 取回复的唯一通道（日志可能被其他行插花，不做解析） if (!cfg.reply_out.empty()) @@ -580,8 +587,7 @@ int main(int argc, char *argv[]) const std::size_t bos_id = tokenizer->bos_id(); auto prompt_tokens = tokenizer->encode(cfg.prompt); // 添加 BOS 前缀：tokenizer 提供 bos_id 时，用作 prompt 的起始标记 - if (bos_id != nn::Tokenizer::npos) - prompt_tokens.insert(prompt_tokens.begin(), bos_id); + if (bos_id != nn::Tokenizer::npos) prompt_tokens.insert(prompt_tokens.begin(), bos_id); std::cout << "提示: \"" << cfg.prompt << "\"\n"; std::cout << "Prompt tokens: " << prompt_tokens.size() << " 个\n"; std::cout << "生成 " << cfg.max_tokens << " 个 token" @@ -616,11 +622,10 @@ int main(int argc, char *argv[]) } std::cout << "\n----------------------------------------" << std::endl; - std::cout << "生成 " << generated.size() << " tokens, 耗时 " - << std::fixed << std::setprecision(1) << gen_sec << "s" - << " (" << std::setprecision(0) - << (gen_sec > 0 ? generated.size() / gen_sec : 0) << " tokens/s)" - << std::endl; + std::cout << "生成 " << generated.size() << " tokens, 耗时 " << std::fixed + << std::setprecision(1) << gen_sec << "s" + << " (" << std::setprecision(0) << (gen_sec > 0 ? generated.size() / gen_sec : 0) + << " tokens/s)" << std::endl; return 0; } diff --git a/src/text_train.cpp b/src/text_train.cpp index cb0d05f2..532e08e1 100644 --- a/src/text_train.cpp +++ b/src/text_train.cpp @@ -64,15 +64,14 @@ namespace // "精度 × activation offload" 组合，否则 rc0 承诺的配置期拒绝会被绕开 // （残余路径只能到首个 block 才由库层 fail-fast）。这里经 Result 上抛 // （不用 std::exit），保持三模块不终止进程的不变量。 - if (auto combo_err = nn::cli::precision_offload_combo_error(cfg)) - NN_FAIL(*combo_err); + if (auto combo_err = nn::cli::precision_offload_combo_error(cfg)) NN_FAIL(*combo_err); } // ── 打开数据集（.nndataset；token 流/掩码/词表都在生成期定好）────── std::cout << "加载数据集: " << cfg.text_path << " ..." << std::endl; NN_TRY_MSG(ds_r, "打开数据集失败", nn::Dataset::open(cfg.text_path)); - const nn::Dataset &dataset = *ds_r; - const std::string &loss_scope = dataset.info().loss_scope; + const nn::Dataset& dataset = *ds_r; + const std::string& loss_scope = dataset.info().loss_scope; // 词表来自数据集配套（内嵌三层 kvrec）；保存模型时原样嵌出（模型 v6）。 NN_TRY_MSG(tokenizer_r, "加载数据集内嵌词表失败", @@ -98,15 +97,12 @@ namespace { NN_TRY_MSG(test_samples_r, "读取测试样本集失败", nn::TextSampleSet::create(dataset, nn::DatasetSubset::Test, window)); - if (test_samples_r->sample_count() > 0) - test_samples = std::move(*test_samples_r); - else - empty_test_tokens = test_samples_r->token_count(); + if (test_samples_r->sample_count() > 0) test_samples = std::move(*test_samples_r); + else empty_test_tokens = test_samples_r->token_count(); } std::cout << "训练子集: " << train_samples.doc_count() << " docs / " - << train_samples.token_count() << " tokens loss_scope=" << loss_scope - << "\n"; + << train_samples.token_count() << " tokens loss_scope=" << loss_scope << "\n"; // 文档感知：doc_ids 非空（每 doc = 一篇文档）即启用块对角掩码。 if (!train_samples.corpus().doc_ids.empty()) std::cout << "文档感知掩码已启用（每行 = 一篇文档，" @@ -117,15 +113,15 @@ namespace if (loss_scope == "assistant") { std::size_t asst_tokens = 0; - for (unsigned char m : train_samples.corpus().loss_mask) asst_tokens += m; + for (unsigned char m : train_samples.corpus().loss_mask) + asst_tokens += m; std::cout << "assistant loss 掩码已启用（" << asst_tokens << " / " - << train_samples.token_count() - << " tokens 作为预测目标参与 loss）\n"; + << train_samples.token_count() << " tokens 作为预测目标参与 loss）\n"; } - std::cout << "滑动窗口: seq_len=" << cfg.seq_len - << " stride=" << train_samples.stride() - << " 样本数=" << train_samples.sample_count() << "\n" << std::endl; + std::cout << "滑动窗口: seq_len=" << cfg.seq_len << " stride=" << train_samples.stride() + << " 样本数=" << train_samples.sample_count() << "\n" + << std::endl; // 空 test 子集的启动信息（v1.9.0 锚：仍打印"样本数: 0"；评估跳过由 // TextTrainer 的 nullopt test 语义保证）。 @@ -134,8 +130,8 @@ namespace if (loss_scope == "assistant") std::cout << "测试集 assistant 掩码: 0 / " << *empty_test_tokens << " tokens 参与评估\n"; - std::cout << "测试集: " << dataset.info().test_docs << " docs 样本数: 0 tokens: " - << *empty_test_tokens << std::endl; + std::cout << "测试集: " << dataset.info().test_docs + << " docs 样本数: 0 tokens: " << *empty_test_tokens << std::endl; } // ── 数据集身份（P1-2 断点门禁的"内容"侧）──────────────────────────── @@ -151,8 +147,7 @@ namespace dataset_identity.num_samples = train_samples.sample_count(); dataset_identity.loss_scope = loss_scope; // steps_per_epoch 是**派生量**（样本数 ÷ batch），与训练工厂同一公式。 - dataset_identity.steps_per_epoch = - train_samples.steps_per_epoch(cfg.batch_size); + dataset_identity.steps_per_epoch = train_samples.steps_per_epoch(cfg.batch_size); // ── 数据集身份门禁（硬错、无逃生阀；docs/internal/23 §2.1.1）──────────── // 数据集变 ⇒ 样本数变 ⇒ steps_per_epoch 变 ⇒ step 级调度的 total_steps、 @@ -160,8 +155,8 @@ namespace // 的续训语义。合法出口 = 从 .nnmodel 开新一轮（下面的 warm start 分支）。 if (adoption.adopted) { - NN_TRY_CHECK(nn::cli::verify_dataset_identity( - adoption.state.dataset, dataset_identity, cfg.resume_path)); + NN_TRY_CHECK(nn::cli::verify_dataset_identity(adoption.state.dataset, dataset_identity, + cfg.resume_path)); } // ── 打印配置 ───────────────────────────────────────────── @@ -174,24 +169,22 @@ namespace std::cout << " Transformer 层数: " << cfg.num_layers << "\n"; std::cout << " FFN 维度: " << cfg.d_ff << "\n"; std::cout << " 序列长度: " << cfg.seq_len << "\n"; - std::cout << " 模型架构: " - << (cfg.model_type == "rapt" ? "RAPT (ReLU 线性注意力)" - : "GPT") << "\n"; + std::cout << " 模型架构: " << (cfg.model_type == "rapt" ? "RAPT (ReLU 线性注意力)" : "GPT") + << "\n"; if (cfg.model_type == "rapt") { std::cout << " 位置编码: RoPE（RLA 强制，施加在 ReLU 之前）\n"; } std::cout << " 优化器: " << cfg.optimizer_name << " 学习率: " << cfg.lr << "\n"; // β1/β2 的消费方 = 表中 beta = true 的优化器（sgd 家族/muon 忽略） - if (const auto *om = nn::find_optimizer(cfg.optimizer_name); om != nullptr && om->beta) + if (const auto* om = nn::find_optimizer(cfg.optimizer_name); om != nullptr && om->beta) { // 前面的计时打印固定过 precision(1)，这里恢复足够位数（0.95 才不会印成 0.9） - std::cout << std::defaultfloat << std::setprecision(6) - << " Adam β1/β2: " << cfg.beta1 << " / " << cfg.beta2 << "\n"; + std::cout << std::defaultfloat << std::setprecision(6) << " Adam β1/β2: " << cfg.beta1 + << " / " << cfg.beta2 << "\n"; } std::cout << " loss 范围: " << loss_scope << "（生成期定好，随数据集）\n"; - if (cfg.max_steps > 0) - std::cout << " 最大步数: " << cfg.max_steps << "\n"; + if (cfg.max_steps > 0) std::cout << " 最大步数: " << cfg.max_steps << "\n"; std::cout << " 轮数: " << cfg.epochs << " 批大小: " << cfg.batch_size << "\n"; std::cout << " GPU: " << (cfg.gpu_enabled ? "启用" : "禁用") << "\n"; std::cout << " 梯度日志: " << (cfg.grad_log ? "启用" : "禁用") << "\n"; @@ -213,10 +206,10 @@ namespace // 触发条件 = **真的有 f16 存储**（四个槽里任一为 f16）。不能用 // `!is_profile_f32()`：`param = t1_58` 时三值并不产生 f16 存储（latent 恒 f32、 // 只有量化缓冲是 f16），那样会打出误导性的"f16 存储已启用"。 - const bool f16_storage = (cfg.precision.param == nn::Precision::F16 || - cfg.precision.compute == nn::Precision::F16 || - cfg.precision.stable == nn::Precision::F16 || - cfg.precision.optimizer == nn::Precision::F16); + const bool f16_storage = + (cfg.precision.param == nn::Precision::F16 || cfg.precision.compute == nn::Precision::F16 || + cfg.precision.stable == nn::Precision::F16 || + cfg.precision.optimizer == nn::Precision::F16); if (f16_storage) { std::cout << "[精度] f16 存储已启用（基类边界 cast + in-kernel f16 变体 + f16 GEMM）\n" @@ -249,20 +242,17 @@ namespace "见 docs/internal/21-quantized-weights.md §4.8.3。\n" " · 三值 GPT：改用 --model gpt\n" " · 三值 MLP：mnist_train --arch mlp --precision-param t1_58"); - model_build = nn::build_rapt_model(*engine, nn::RAPTConfig{ - tokenizer->vocab_size(), cfg.d_model, cfg.seq_len, - cfg.num_heads, cfg.d_ff, cfg.num_layers, - cfg.pos_encoding, cfg.activation, cfg.norm_type, - /*causal=*/true, cfg.precision}); + model_build = nn::build_rapt_model( + *engine, nn::RAPTConfig{tokenizer->vocab_size(), cfg.d_model, cfg.seq_len, + cfg.num_heads, cfg.d_ff, cfg.num_layers, cfg.pos_encoding, + cfg.activation, cfg.norm_type, + /*causal=*/true, cfg.precision}); } else { model_build = nn::build_gpt_model( - *engine, - tokenizer->vocab_size(), cfg.d_model, cfg.seq_len, - cfg.num_heads, cfg.d_ff, cfg.num_layers, - cfg.pos_encoding, cfg.activation, cfg.norm_type, - cfg.precision); + *engine, tokenizer->vocab_size(), cfg.d_model, cfg.seq_len, cfg.num_heads, cfg.d_ff, + cfg.num_layers, cfg.pos_encoding, cfg.activation, cfg.norm_type, cfg.precision); } NN_TRY_MSG(model_build_r, "构建模型失败", std::move(model_build)); auto model = std::move(*model_build_r); @@ -319,9 +309,11 @@ namespace else { model.set_activation_offload(true); - std::cout << "activation offload 已启用: 每块激活搬 host-visible，backward 拷回（不重算）\n"; + std::cout + << "activation offload 已启用: 每块激活搬 host-visible，backward 拷回（不重算）\n"; if (cfg.checkpoint_every > 0) - std::cout << " （混合模式）与 checkpoint 共存：checkpoint 块重算，其余块 offload\n"; + std::cout + << " （混合模式）与 checkpoint 共存：checkpoint 块重算，其余块 offload\n"; std::cout << " offload slab 在首次 forward 时按激活形状分配（RAM ≈ 激活实际体积）\n"; } } @@ -329,22 +321,20 @@ namespace // ── 构建规格（用于保存） ───────────────────────────────── nn::ModelSpec spec; if (cfg.model_type == "rapt") - spec = nn::make_rapt_spec( - tokenizer->vocab_size(), cfg.d_model, cfg.seq_len, - cfg.num_heads, cfg.d_ff, cfg.num_layers, - cfg.pos_encoding, cfg.activation, cfg.norm_type); + spec = nn::make_rapt_spec(tokenizer->vocab_size(), cfg.d_model, cfg.seq_len, cfg.num_heads, + cfg.d_ff, cfg.num_layers, cfg.pos_encoding, cfg.activation, + cfg.norm_type); else - spec = nn::make_gpt_spec( - tokenizer->vocab_size(), cfg.d_model, cfg.seq_len, - cfg.num_heads, cfg.d_ff, cfg.num_layers, - cfg.pos_encoding, cfg.activation, cfg.norm_type); + spec = nn::make_gpt_spec(tokenizer->vocab_size(), cfg.d_model, cfg.seq_len, cfg.num_heads, + cfg.d_ff, cfg.num_layers, cfg.pos_encoding, cfg.activation, + cfg.norm_type); // 三值权重（P1.5，docs 21 §4.7/§4.8.3）：param=t1_58 → 线性层用 BitLinear。 // 规格是**加载时的权威来源**（缺键 = None → 旧文件零破坏）；--resume 时下面 // 会用文件里的规格覆盖本行。 - spec.as()->weight_quant = - (cfg.precision.param == nn::Precision::T1_58) ? nn::WeightQuant::T1_58 - : nn::WeightQuant::None; + spec.as()->weight_quant = (cfg.precision.param == nn::Precision::T1_58) + ? nn::WeightQuant::T1_58 + : nn::WeightQuant::None; // ── 断点恢复（--resume / --resume-epoch / --resume-step）────────────── @@ -382,8 +372,8 @@ namespace // 位置必须在优化器构造之后：状态张量按优化器自身的缓冲形状回写。 if (adoption.adopted) { - NN_TRY_CHECK(nn::cli::load_optimizer_state(*engine, cfg.resume_path, - session.optimizer(), std::cout)); + NN_TRY_CHECK(nn::cli::load_optimizer_state(*engine, cfg.resume_path, session.optimizer(), + std::cout)); } // ── 训练工厂：数据集（样本集）直接接入 ─────────────────────────────── @@ -407,8 +397,7 @@ namespace if (report.abort_code != 0) return report.abort_code; // ── step 剖析报告（NN_PROFILE=1 / NN_PROFILE_TRACE 时生效）──────────── - if (nn::prof::enabled()) - nn::prof::dump(stderr, report.elapsed_ns); + if (nn::prof::enabled()) nn::prof::dump(stderr, report.elapsed_ns); // ── 保存模型（含规格 + 嵌入 tokenizer） ────────────────── // 推理模型（Model 角色）：三值层写打包形态 + τ（不含训练参数）。 @@ -417,8 +406,8 @@ namespace nn::cli::save_text_train_model(cfg.save_path, trainer.session().model(), spec, tokenizer_bytes)); (void)save_result; - std::cout << "\n训练完成! 总耗时: " << std::fixed << std::setprecision(1) - << report.seconds << "s" + std::cout << "\n训练完成! 总耗时: " << std::fixed << std::setprecision(1) << report.seconds + << "s" << " 词表已嵌入模型文件" << std::endl; // 退出码语义（P1-1）：有断点没写出去 ⇒ 非零退出。推理模型文件已落盘， // 但"可续训"这一性质不成立，上层不得把它当成功（文案在 checkpoint 模块）。 @@ -430,15 +419,14 @@ namespace // 本函数是唯一的进程出口（`main` 返回退出码）：解析层内部按 CLI 语义 // exit(0)/exit(1)（--help / 参数错误），其余一切经 `run_text_train` 的 Result // 回报后在此统一打印并返回 1。 -int main(int argc, char *argv[]) +int main(int argc, char* argv[]) { nn::cli::TrainConfig cfg = nn::cli::parse_text_train_args(argc, argv); // ── 危险精度组合启动期告警（仅警告，不拒绝；docs/internal/24 §2.3 F1）──── // optimizer=f16（含四字段全 f16）只打印一次到 stderr：不改变退出码、不阻断 // 训练、不静默降级；`--f16` 与合法组合返回空串（静默）。 - if (const std::string warn = nn::cli::dangerous_precision_warning(cfg.precision); - !warn.empty()) + if (const std::string warn = nn::cli::dangerous_precision_warning(cfg.precision); !warn.empty()) std::cerr << warn; auto run = run_text_train(cfg); diff --git a/src/text_train_modules_test.cpp b/src/text_train_modules_test.cpp index 8f4ae5c2..ca4626e5 100644 --- a/src/text_train_modules_test.cpp +++ b/src/text_train_modules_test.cpp @@ -58,20 +58,22 @@ std::filesystem::path scratch_dir() } [[nodiscard]] nn::ModelSpec tiny_spec() -{ - return nn::make_gpt_spec(VOCAB, D_MODEL, SEQ_LEN, HEADS, D_FF, LAYERS); -} +{ return nn::make_gpt_spec(VOCAB, D_MODEL, SEQ_LEN, HEADS, D_FF, LAYERS); } // 按参数顺序拍参数快照（宿主侧 f32；用于"参数确实移动"的断言） -[[nodiscard]] std::vector params_snapshot(nn::ComputeEngine& eng, - nn::Model& m) +[[nodiscard]] std::vector params_snapshot(nn::ComputeEngine& eng, nn::Model& m) { std::vector out; for (auto& t : m.parameters()) { auto mat = eng.to_matrix(t.get()); - if (!mat) { out.push_back(-1e30f); continue; } - for (auto v : mat->span()) out.push_back(v); + if (!mat) + { + out.push_back(-1e30f); + continue; + } + for (auto v : mat->span()) + out.push_back(v); } return out; } @@ -123,8 +125,7 @@ void test_role_gate_and_roundtrip(const std::filesystem::path& dir) const std::string model_file = (dir / "gate.nnmodel").string(); const std::string ckpt_file = (dir / "gate.nnckpt").string(); - NN_ASSERT_ALWAYS(nn::cli::save_text_train_model(model_file, original, spec, - kTokenizerBytes), + NN_ASSERT_ALWAYS(nn::cli::save_text_train_model(model_file, original, spec, kTokenizerBytes), "text_train_modules_test: 写推理模型失败"); nn::cli::TextTrainCheckpointLog log; nn::TextTrainState ckpt_state; // P1-2：断点带训练状态段（本用例不挂优化器状态张量） @@ -167,13 +168,14 @@ void test_role_gate_and_roundtrip(const std::filesystem::path& dir) // 解析层登记（不看值：给了 --lr 且值与断点相同也登记） { - std::vector args{"text_train", "/tmp/ds.nndataset", "--lr", "0.001", - "--epochs", "3", "--save", "/tmp/m.nnmodel"}; + std::vector args{"text_train", "/tmp/ds.nndataset", "--lr", + "0.001", "--epochs", "3", + "--save", "/tmp/m.nnmodel"}; std::vector argv; argv.reserve(args.size()); - for (std::string& a : args) argv.push_back(a.data()); - auto cfg_p = nn::cli::parse_text_train_args(static_cast(argv.size()), - argv.data()); + for (std::string& a : args) + argv.push_back(a.data()); + auto cfg_p = nn::cli::parse_text_train_args(static_cast(argv.size()), argv.data()); CHECK(cfg_p.explicit_train_flags.size() == 2 && cfg_p.explicit_train_flags[0] == "--lr" && cfg_p.explicit_train_flags[1] == "--epochs", @@ -218,15 +220,15 @@ void test_role_gate_and_roundtrip(const std::filesystem::path& dir) legacy_state.dataset.steps_per_epoch = 5; const std::string legacy_ckpt = (dir / "legacy.nnckpt").string(); nn::cli::TextTrainCheckpointLog legacy_log; - legacy_log.save(legacy_ckpt, original, spec, kTokenizerBytes, legacy_state, - no_blobs, std::cerr); + legacy_log.save(legacy_ckpt, original, spec, kTokenizerBytes, legacy_state, no_blobs, + std::cerr); nn::cli::TrainConfig legacy_cfg; legacy_cfg.resume_path = legacy_ckpt; legacy_cfg.load_existing = true; std::ostringstream legacy_out, legacy_err; auto legacy_r = nn::cli::adopt_resume_state(legacy_cfg, legacy_out, legacy_err); - CHECK(legacy_r && legacy_cfg.batch_size == 4 && - legacy_cfg.accum_steps == 2 && legacy_cfg.start_step == 1, + CHECK(legacy_r && legacy_cfg.batch_size == 4 && legacy_cfg.accum_steps == 2 && + legacy_cfg.start_step == 1, "旧 micro-batch 断点 → global=4 / accum=2 / update step=1"); if (legacy_r) CHECK(legacy_r->state.recipe.batch_size_is_global && @@ -275,8 +277,7 @@ void test_role_gate_and_roundtrip(const std::filesystem::path& dir) } CHECK(out.str().find("已加载模型: " + ckpt_file) != std::string::npos, "恢复成功应打印已加载模型路径"); - CHECK(params_identical(eng, original, restored), - "恢复后的参数应与保存前逐位一致"); + CHECK(params_identical(eng, original, restored), "恢复后的参数应与保存前逐位一致"); // 角色语义变更后（2026-10-16 裁定）的反例形态：dense `.nnmodel` ⇒ **warm start** // 走通整条恢复流程（SFT 入口）；三值 `.nnmodel` 见下面的 P1-7 反量化块。 @@ -289,9 +290,7 @@ void test_role_gate_and_roundtrip(const std::filesystem::path& dir) auto warm = nn::cli::resume_text_train(eng, warm_cfg, m2, spec_out2, out2, err2); CHECK(warm.has_value(), "dense 推理模型应能走通恢复流程（warm start）"); CHECK(err2.str().empty(), "warm start 不应在 stderr 上留痕（错误才经 Result/stderr）"); - if (warm) - CHECK(params_identical(eng, original, m2), - "warm start 载入的参数应与写出前逐位一致"); + if (warm) CHECK(params_identical(eng, original, m2), "warm start 载入的参数应与写出前逐位一致"); // ── P1-7：三值（T1_58）`.nnmodel` 的续训入口（打包 → latent 反量化）────── // 三值推理文件的载荷是打包形态（没有 latent，部署态 `backward()` 会报错）。 @@ -307,8 +306,7 @@ void test_role_gate_and_roundtrip(const std::filesystem::path& dir) nn::Model src = std::move(*src_r); const std::string t_file = (dir / "ternary.nnmodel").string(); - NN_ASSERT_ALWAYS(nn::cli::save_text_train_model(t_file, src, t_spec, - kTokenizerBytes), + NN_ASSERT_ALWAYS(nn::cli::save_text_train_model(t_file, src, t_spec, kTokenizerBytes), "text_train_modules_test: 写三值推理模型失败"); { auto role = nn::peek_model_role(t_file); @@ -329,8 +327,7 @@ void test_role_gate_and_roundtrip(const std::filesystem::path& dir) nn::ModelSpec t_spec_out; nn::Model t_model; std::ostringstream t_out, t_err; - auto t_res = - nn::cli::resume_text_train(eng, t_cfg, t_model, t_spec_out, t_out, t_err); + auto t_res = nn::cli::resume_text_train(eng, t_cfg, t_model, t_spec_out, t_out, t_err); CHECK(t_res.has_value(), "三值 .nnmodel 应能按 warm start 载入（P1-7：不再被拒）"); if (t_res.has_value()) { @@ -370,8 +367,7 @@ void test_role_gate_and_roundtrip(const std::filesystem::path& dir) nn::ModelSpec d_spec_out; nn::Model d_model; std::ostringstream d_out, d_err; - auto d_res = nn::cli::resume_text_train(eng, d_cfg, d_model, d_spec_out, d_out, - d_err); + auto d_res = nn::cli::resume_text_train(eng, d_cfg, d_model, d_spec_out, d_out, d_err); CHECK(d_res.has_value(), "不给精度旗标也应能载入三值 warm start（规格定 BitLinear）"); if (d_res.has_value()) { @@ -384,7 +380,7 @@ void test_role_gate_and_roundtrip(const std::filesystem::path& dir) // 8 步 Adam：backward 可跑（不是只读部署态）、参数确实移动（P1-7 验收） auto opt = nn::create_optimizer("adam", eng, t_model.parameters(), - t_model.param_gradients(), 1e-2f, 0.0f, t1); + t_model.param_gradients(), 1e-2f, 0.0f, t1); NN_ASSERT_ALWAYS(opt != nullptr, "text_train_modules_test: 建三值优化器失败"); const std::vector before = params_snapshot(eng, t_model); @@ -400,16 +396,35 @@ void test_role_gate_and_roundtrip(const std::filesystem::path& dir) { steps_ok = t_model.zero_grad().has_value(); auto logits = t_model.forward(x_t); - if (!logits) { steps_ok = false; break; } + if (!logits) + { + steps_ok = false; + break; + } auto loss = ce.forward(eng, *logits, tgt_t); - if (!loss) { steps_ok = false; break; } + if (!loss) + { + steps_ok = false; + break; + } auto grad = ce.backward(); - if (!grad) { steps_ok = false; break; } - if (!t_model.backward(*grad)) { steps_ok = false; break; } - if (!opt->step()) { steps_ok = false; break; } + if (!grad) + { + steps_ok = false; + break; + } + if (!t_model.backward(*grad)) + { + steps_ok = false; + break; + } + if (!opt->step()) + { + steps_ok = false; + break; + } } - CHECK(steps_ok, - "三值 warm start 后 8 步 Adam（forward/loss/backward/step）全部成功"); + CHECK(steps_ok, "三值 warm start 后 8 步 Adam（forward/loss/backward/step）全部成功"); const std::vector after = params_snapshot(eng, t_model); std::size_t moved = 0, total = 0; for (std::size_t i = 0; i < before.size() && i < after.size(); ++i) @@ -447,12 +462,10 @@ void test_ckpt_accounting(const std::filesystem::path& dir) std::ostringstream err_code; CHECK(nn::cli::text_train_ckpt_exit_code(log.failures(), err_code) == 1, "有断点写失败时收尾退出码必须为 1"); - CHECK(err_code.str().find("退出码 1") != std::string::npos, - "收尾错误行应说明退出码"); + CHECK(err_code.str().find("退出码 1") != std::string::npos, "收尾错误行应说明退出码"); // 成功路径：不增计数、退出码 0、不打印 - log.save((dir / "ok.nnckpt").string(), model, spec, kTokenizerBytes, ckpt_state, - no_blobs, err); + log.save((dir / "ok.nnckpt").string(), model, spec, kTokenizerBytes, ckpt_state, no_blobs, err); CHECK(log.failures() == 1, "成功保存不应增加失败计数"); std::ostringstream err_ok; CHECK(nn::cli::text_train_ckpt_exit_code(0, err_ok) == 0, "无失败时退出码必须为 0"); @@ -469,9 +482,8 @@ void test_train_state_wiring(const std::filesystem::path& dir) NN_ASSERT_ALWAYS(built, "test_train_state_wiring: 构建小 GPT 失败"); nn::Model model = std::move(*built); - auto opt = nn::create_optimizer("adam", eng, model.parameters(), - model.param_gradients(), 1e-3f, 0.0f, - nn::PrecisionProfile{}, 0.9f, 0.999f); + auto opt = nn::create_optimizer("adam", eng, model.parameters(), model.param_gradients(), 1e-3f, + 0.0f, nn::PrecisionProfile{}, 0.9f, 0.999f); NN_ASSERT_ALWAYS(opt != nullptr, "test_train_state_wiring: 建优化器失败"); NN_ASSERT_ALWAYS(opt->step(), "test_train_state_wiring: 一步失败"); @@ -496,10 +508,9 @@ void test_train_state_wiring(const std::filesystem::path& dir) nn::cli::TextTrainCheckpointLog log; // P2b：状态装配唯一实现 = nn::make_text_train_state（TextTrainer 内部同一函数） const nn::TextTrainState state = nn::make_text_train_state( - cfg, id, /*steps_per_epoch=*/2, /*epoch=*/1, /*step_done=*/2, 1e-3, - opt->state_step()); - nn::cli::save_ckpt_with_state(path.string(), model, spec, kTokenizerBytes, state, *opt, - log, err); + cfg, id, /*steps_per_epoch=*/2, /*epoch=*/1, /*step_done=*/2, 1e-3, opt->state_step()); + nn::cli::save_ckpt_with_state(path.string(), model, spec, kTokenizerBytes, state, *opt, log, + err); CHECK(log.failures() == 0 && err.str().empty(), "带状态的断点应一次写成功"); auto st = nn::peek_train_state(path.string()); @@ -527,24 +538,22 @@ void test_train_state_wiring(const std::filesystem::path& dir) // 数据集身份门禁走的就是断点里的这份身份 nn::DatasetIdentity changed = id; changed.num_samples = 5; - CHECK(!nn::cli::verify_dataset_identity(ts->dataset, changed, path.string()) - .has_value(), - "断点记录的身份 vs 变化的网格 ⇒ 硬报错（接线一致）"); + CHECK( + !nn::cli::verify_dataset_identity(ts->dataset, changed, path.string()).has_value(), + "断点记录的身份 vs 变化的网格 ⇒ 硬报错（接线一致）"); } } // 重启：新优化器（状态零初始化）→ 载入断点状态 → t 与张量恢复 - auto opt2 = nn::create_optimizer("adam", eng, model.parameters(), - model.param_gradients(), 1e-3f, 0.0f, - nn::PrecisionProfile{}, 0.9f, 0.999f); + auto opt2 = nn::create_optimizer("adam", eng, model.parameters(), model.param_gradients(), + 1e-3f, 0.0f, nn::PrecisionProfile{}, 0.9f, 0.999f); NN_ASSERT_ALWAYS(opt2 != nullptr, "test_train_state_wiring: 重建优化器失败"); CHECK(opt2->state_step() == 0, "新优化器的步计数器从 0 起（前置条件）"); std::ostringstream out2; auto lr = nn::cli::load_optimizer_state(eng, path.string(), *opt2, out2); CHECK(lr.has_value(), "load_optimizer_state 应成功"); CHECK(opt2->state_step() == 1, "载入后 Adam 的 t 恢复（偏置修正计数不丢）"); - CHECK(out2.str().find("已恢复优化器状态") != std::string::npos, - "载入成功应打印一行进度"); + CHECK(out2.str().find("已恢复优化器状态") != std::string::npos, "载入成功应打印一行进度"); } // ── [5] 推理模型角色（text_infer 侧前提）─────────────────────────────── diff --git a/src/text_trainer_test.cpp b/src/text_trainer_test.cpp index f294cf5b..e2e1662b 100644 --- a/src/text_trainer_test.cpp +++ b/src/text_trainer_test.cpp @@ -134,8 +134,7 @@ struct Fixture } [[nodiscard]] nn::TrainingSession make_session(Fixture& f, nn::Scalar lr = 1e-2f, - nn::Scalar clip = 0, - std::size_t accum_steps = 1) + nn::Scalar clip = 0, std::size_t accum_steps = 1) { auto built = nn::build_gpt_model_from_spec(*f.engine, f.spec); NN_ASSERT_ALWAYS(built, "text_trainer_test: 构建小 GPT 失败"); @@ -249,8 +248,7 @@ void test_step_range(const std::string& dir) CHECK(r && r->steps_executed == 1, "max_steps=1 只执行 1 步"); CHECK(out1.str().find("已达 --max-steps=1") != std::string::npos, "max_steps 提前结束文案"); // P1：--max-steps 提前停时状态必须是"真实下一步"，不能谎报 (epochs,0) - CHECK(r && r->final_epoch == 0 && r->final_step_in_epoch == 1, - "report 记录真实停止位置"); + CHECK(r && r->final_epoch == 0 && r->final_step_in_epoch == 1, "report 记录真实停止位置"); CHECK(t1->train_state().counters.epoch == 0 && t1->train_state().counters.step_in_epoch == 1 && t1->train_state().counters.global_step == 1, @@ -291,11 +289,10 @@ void test_global_batch_accumulation(const std::string& dir) std::size_t updates = 0; nn::TextTrainerHooks hooks; hooks.out = nullptr; - hooks.on_batch_ready = [&](nn::TrainingSession&, const nn::TextSampleBatch& batch, - std::size_t) { micro_batches += batch.batch; }; - hooks.on_update = [&](nn::TrainingSession&, std::size_t, std::size_t, nn::Scalar) { - ++updates; - }; + hooks.on_batch_ready = [&](nn::TrainingSession&, const nn::TextSampleBatch& batch, std::size_t) + { micro_batches += batch.batch; }; + hooks.on_update = [&](nn::TrainingSession&, std::size_t, std::size_t, nn::Scalar) + { ++updates; }; auto trainer_r = nn::TextTrainer::create(std::move(session), *ds, cfg, hooks); CHECK(trainer_r, "accum: create"); @@ -303,12 +300,10 @@ void test_global_batch_accumulation(const std::string& dir) CHECK(trainer_r->micro_batch_size() == 2, "accum: micro = global / accum"); const std::size_t samples = trainer_r->sample_count(); const std::size_t expected_steps = (samples + 3) / 4; - CHECK(trainer_r->steps_per_epoch() == expected_steps, - "accum: steps_per_epoch 按 global batch"); + CHECK(trainer_r->steps_per_epoch() == expected_steps, "accum: steps_per_epoch 按 global batch"); auto report = trainer_r->fit(); - CHECK(report && report->steps_executed == expected_steps, - "accum: update step 数正确"); + CHECK(report && report->steps_executed == expected_steps, "accum: update step 数正确"); CHECK(report && report->samples_executed == samples, "accum: 样本数正确"); CHECK(trainer_r->session().step_count() == expected_steps, "accum: Session step_count 按 update 计"); @@ -316,8 +311,7 @@ void test_global_batch_accumulation(const std::string& dir) CHECK(micro_batches == samples, "accum: micro-batch 样本列数总和"); CHECK(trainer_r->train_state().recipe.batch_size_is_global, "accum: 状态标记 global batch 语义"); - CHECK(trainer_r->train_state().recipe.micro_batch_size == 2, - "accum: 状态记录派生 micro-batch"); + CHECK(trainer_r->train_state().recipe.micro_batch_size == 2, "accum: 状态记录派生 micro-batch"); } // ── [3b] 高性能入口（CLI 实际路径）：样本集 + nullopt test ───────────────── @@ -345,8 +339,8 @@ void test_fast_path_sets(const std::string& dir) std::ostringstream out; nn::TextTrainerHooks hooks; hooks.out = &out; - auto trainer_r = nn::TextTrainer::create(std::move(session), std::move(*train_r), - std::nullopt, cfg, hooks); + auto trainer_r = + nn::TextTrainer::create(std::move(session), std::move(*train_r), std::nullopt, cfg, hooks); CHECK(trainer_r, "fast-path: create(session, train_set, nullopt, cfg)"); if (!trainer_r) return; CHECK(trainer_r->test_samples() == nullptr, "fast-path: nullopt test ⇒ 不评估"); diff --git a/src/tokenizer_infer.cpp b/src/tokenizer_infer.cpp index e672b1cb..e9f12818 100644 --- a/src/tokenizer_infer.cpp +++ b/src/tokenizer_infer.cpp @@ -24,7 +24,7 @@ #include // ==================== 帮助信息 ==================== -void print_usage(const char *prog) +void print_usage(const char* prog) { nn::cli::Help help(std::cout, prog, "BPE 分词器推理程序 (编码/解码)"); @@ -38,13 +38,15 @@ void print_usage(const char *prog) help.item("", "直接给文本文件即统计最长行 token 数"); help.section("选项"); - help.opt("--vocab ", "词表 .nnvocab 路径 (默认: bpe_vocab.nnvocab)\n自动识别分词器类型 (bpe / charbpe)"); + help.opt("--vocab ", + "词表 .nnvocab 路径 (默认: bpe_vocab.nnvocab)\n自动识别分词器类型 (bpe / charbpe)"); help.opt("--encode ", "编码文本为 token IDs"); help.opt("--decode ", "解码 token IDs (逗号分隔) 为文本"); help.opt("--encode-file ", "编码整个文件 (同时自动统计最长行 token 数)"); help.opt("--interactive", "交互模式 (输入 'quit' 退出)"); help.opt("--top ", "最长行排行榜行数 (默认: 10)"); - help.opt("--threads ", "encode 并行度: 0=自动(默认) 1=顺序 >1=指定\n保序并行：任何并行度下输出与顺序执行逐字节一致"); + help.opt("--threads ", "encode 并行度: 0=自动(默认) 1=顺序 " + ">1=指定\n保序并行：任何并行度下输出与顺序执行逐字节一致"); help.opt("--show-bytes", "显示原始字节 (调试用)"); help.opt("--help, -h", "显示此帮助信息"); } @@ -63,7 +65,7 @@ struct Config bool show_bytes = false; }; -Config parse_args(int argc, char *argv[]) +Config parse_args(int argc, char* argv[]) { Config cfg; for (int i = 1; i < argc; ++i) @@ -74,18 +76,12 @@ Config parse_args(int argc, char *argv[]) print_usage(argv[0]); std::exit(0); } - else if (arg == "--vocab" && i + 1 < argc) - cfg.vocab_path = argv[++i]; - else if (arg == "--encode" && i + 1 < argc) - cfg.encode_text = argv[++i]; - else if (arg == "--decode" && i + 1 < argc) - cfg.decode_ids = argv[++i]; - else if (arg == "--encode-file" && i + 1 < argc) - cfg.encode_file = argv[++i]; - else if (arg == "--interactive") - cfg.interactive = true; - else if (arg == "--show-bytes") - cfg.show_bytes = true; + else if (arg == "--vocab" && i + 1 < argc) cfg.vocab_path = argv[++i]; + else if (arg == "--encode" && i + 1 < argc) cfg.encode_text = argv[++i]; + else if (arg == "--decode" && i + 1 < argc) cfg.decode_ids = argv[++i]; + else if (arg == "--encode-file" && i + 1 < argc) cfg.encode_file = argv[++i]; + else if (arg == "--interactive") cfg.interactive = true; + else if (arg == "--show-bytes") cfg.show_bytes = true; else if (arg == "--top" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); @@ -106,8 +102,7 @@ Config parse_args(int argc, char *argv[]) } cfg.threads = *v; } - else if (!arg.starts_with("--")) - cfg.text_path = arg; + else if (!arg.starts_with("--")) cfg.text_path = arg; else { std::cerr << "未知参数: " << arg << "\n使用 --help 查看用法\n"; @@ -118,7 +113,7 @@ Config parse_args(int argc, char *argv[]) } // ==================== 解析 token IDs 字符串 ==================== -std::vector parse_ids(const std::string &s) +std::vector parse_ids(const std::string& s) { std::vector ids; std::stringstream ss(s); @@ -140,7 +135,7 @@ std::vector parse_ids(const std::string &s) // ==================== 统一输出 token IDs ==================== // 辅助函数：以 "Tokens (): [id0, id1, ...]" 格式输出 ID 列表， // 消除交互模式 / 命令行模式 / 默认编码路径中的重复打印代码。 -void print_token_ids(const std::vector &ids, std::ostream &os = std::cout) +void print_token_ids(const std::vector& ids, std::ostream& os = std::cout) { os << " Tokens (" << ids.size() << "): ["; for (std::size_t i = 0; i < ids.size(); ++i) @@ -152,7 +147,7 @@ void print_token_ids(const std::vector &ids, std::ostream &os = std } // ==================== 交互模式 ==================== -void interactive_mode(const nn::Tokenizer &tokenizer, bool show_bytes) +void interactive_mode(const nn::Tokenizer& tokenizer, bool show_bytes) { std::cout << "Tokenizer 交互模式 (输入 'quit' 退出)\n"; std::cout << " encode → 编码文本\n"; @@ -162,12 +157,9 @@ void interactive_mode(const nn::Tokenizer &tokenizer, bool show_bytes) { std::cout << ">>> "; std::string line; - if (!std::getline(std::cin, line)) - break; - if (line == "quit" || line == "exit") - break; - if (line.empty()) - continue; + if (!std::getline(std::cin, line)) break; + if (line == "quit" || line == "exit") break; + if (line.empty()) continue; if (line.starts_with("encode ")) { @@ -176,7 +168,7 @@ void interactive_mode(const nn::Tokenizer &tokenizer, bool show_bytes) print_token_ids(ids); if (show_bytes) { - const auto &v = tokenizer.vocab(); + const auto& v = tokenizer.vocab(); std::cout << " Bytes: "; for (auto id : ids) { @@ -212,7 +204,8 @@ void interactive_mode(const nn::Tokenizer &tokenizer, bool show_bytes) // 逐行编码文本，找出 token 数最多的行，并输出 Top-N 排行榜。 // 无需额外开关：位置参数传入文本文件、或 --encode-file 都会自动执行。 -namespace { +namespace +{ // 统计 UTF-8 字符串的字符数（跳过连续字节） std::size_t utf8_char_count(std::string_view s) noexcept @@ -226,7 +219,7 @@ std::size_t utf8_char_count(std::string_view s) noexcept } // namespace -void longest_lines_mode(const nn::Tokenizer &tokenizer, const std::string &path, std::size_t top_n) +void longest_lines_mode(const nn::Tokenizer& tokenizer, const std::string& path, std::size_t top_n) { std::ifstream in(path, std::ios::binary); if (!in) @@ -235,7 +228,8 @@ void longest_lines_mode(const nn::Tokenizer &tokenizer, const std::string &path, std::exit(1); } - struct LineStat { + struct LineStat + { std::size_t line_no = 0; std::size_t tokens = 0; std::size_t bytes = 0; @@ -245,8 +239,9 @@ void longest_lines_mode(const nn::Tokenizer &tokenizer, const std::string &path, // 小顶堆：始终保留 token 数最大的 top_n 行（堆顶 = 当前第 top_n 大）。 // 每行只需 O(log top_n) 的堆操作，而不是对整表 std::sort（O(top_n log top_n)）。 - struct MinHeap { - bool operator()(const LineStat &a, const LineStat &b) const noexcept + struct MinHeap + { + bool operator()(const LineStat& a, const LineStat& b) const noexcept { return a.tokens > b.tokens; // 小顶堆：tokens 越小越靠堆顶 } @@ -266,17 +261,12 @@ void longest_lines_mode(const nn::Tokenizer &tokenizer, const std::string &path, auto ids = tokenizer.encode(line); total_tokens += ids.size(); - if (!line.empty()) - ++nonempty_lines; - if (ids.empty() && line.empty()) - continue; // 空行不参与排行榜 + if (!line.empty()) ++nonempty_lines; + if (ids.empty() && line.empty()) continue; // 空行不参与排行榜 LineStat st{line_no, ids.size(), line.size(), utf8_char_count(line), line}; - if (best.size() < top_n) - { - best.push(std::move(st)); - } + if (best.size() < top_n) { best.push(std::move(st)); } else if (st.tokens > best.top().tokens) { best.pop(); // 踢掉当前第 top_n 大 @@ -310,17 +300,14 @@ void longest_lines_mode(const nn::Tokenizer &tokenizer, const std::string &path, << static_cast(total_tokens) / nonempty_lines << " tokens/行\n"; if (!top.empty()) { - std::cout << "最长行 : " << top.front().tokens << " tokens (第 " - << top.front().line_no << " 行, " - << top.front().bytes << " 字节, " - << top.front().chars << " 字符)\n"; + std::cout << "最长行 : " << top.front().tokens << " tokens (第 " << top.front().line_no + << " 行, " << top.front().bytes << " 字节, " << top.front().chars << " 字符)\n"; } std::cout << "\n── 最长行 Top " << top.size() << " ──\n"; for (std::size_t i = 0; i < top.size(); ++i) { - const auto &s = top[i]; - std::cout << "#" << (i + 1) - << " 第 " << s.line_no << " 行 | " << s.tokens << " tokens | " + const auto& s = top[i]; + std::cout << "#" << (i + 1) << " 第 " << s.line_no << " 行 | " << s.tokens << " tokens | " << s.bytes << " 字节 | " << s.chars << " 字符\n"; const std::size_t preview_limit = 80; std::string_view preview = s.text; @@ -333,11 +320,10 @@ void longest_lines_mode(const nn::Tokenizer &tokenizer, const std::string &path, { const auto c = static_cast(s.text[i]); std::size_t width = 1; - if ((c & 0xE0) == 0xC0) width = 2; // 2 字节字符 + if ((c & 0xE0) == 0xC0) width = 2; // 2 字节字符 else if ((c & 0xF0) == 0xE0) width = 3; // 3 字节字符（中文等） else if ((c & 0xF8) == 0xF0) width = 4; // 4 字节字符 - if (i + width > preview_limit) - break; // 放不下一个完整字符 + if (i + width > preview_limit) break; // 放不下一个完整字符 i += width; end = i; } @@ -349,7 +335,7 @@ void longest_lines_mode(const nn::Tokenizer &tokenizer, const std::string &path, } // ==================== 主函数 ==================== -int main(int argc, char *argv[]) +int main(int argc, char* argv[]) { Config cfg = parse_args(argc, argv); @@ -412,8 +398,7 @@ int main(int argc, char *argv[]) return 0; } std::cout << "压缩率: " << std::fixed << std::setprecision(2) - << static_cast(text_result->size()) / ids.size() - << " bytes/token\n"; + << static_cast(text_result->size()) / ids.size() << " bytes/token\n"; std::cout << '\n'; // 自动附加最长行 token 统计 longest_lines_mode(*tokenizer, cfg.encode_file, cfg.top); diff --git a/src/tokenizer_train.cpp b/src/tokenizer_train.cpp index f78fe129..46a58fae 100644 --- a/src/tokenizer_train.cpp +++ b/src/tokenizer_train.cpp @@ -21,7 +21,7 @@ using nn::Scalar; // ==================== 帮助信息 ==================== -void print_usage(const char *prog) +void print_usage(const char* prog) { nn::cli::Help help(std::cout, prog, "BPE 分词器训练程序"); @@ -31,7 +31,11 @@ void print_usage(const char *prog) help.item("", "训练文本文件路径 (UTF-8 编码)"); help.section("选项"); - help.opt("--tokenizer ", "分词器类型: bpe/charbpe (默认: bpe)\nbpe: 字节级 BBPE，兼容性最好，任意字节均可编码；\n但中文按 3 字节切分，小模型难学到词义\ncharbpe: 字符级 BPE，每个汉字为 1 个基础 token，\nBPE 学会合并常见词组；中文场景推荐使用"); + help.opt( + "--tokenizer ", + "分词器类型: bpe/charbpe (默认: bpe)\nbpe: 字节级 " + "BBPE，兼容性最好，任意字节均可编码；\n但中文按 3 字节切分，小模型难学到词义\ncharbpe: " + "字符级 BPE，每个汉字为 1 个基础 token，\nBPE 学会合并常见词组；中文场景推荐使用"); help.opt("--output ", "词表输出路径 (默认: bpe_vocab.nnvocab)"); help.opt("--vocab-size ", "目标词表大小 (默认: 5000；最小 258 = 256 字节 + BOS/EOS)"); help.opt("--min-freq ", "最小合并频率 (默认: 2)"); @@ -50,7 +54,7 @@ struct Config std::uint32_t threads = 0; // 预分词并行: 0=自动, 1=顺序, >1=指定 }; -Config parse_args(int argc, char *argv[]) +Config parse_args(int argc, char* argv[]) { Config cfg; for (int i = 1; i < argc; ++i) @@ -66,17 +70,19 @@ Config parse_args(int argc, char *argv[]) cfg.tokenizer_type = argv[++i]; if (cfg.tokenizer_type != "bpe" && cfg.tokenizer_type != "charbpe") { - std::cerr << "未知分词器类型: " << cfg.tokenizer_type - << "，可选: bpe, charbpe\n"; + std::cerr << "未知分词器类型: " << cfg.tokenizer_type << "，可选: bpe, charbpe\n"; std::exit(1); } } - else if (arg == "--output" && i + 1 < argc) - cfg.output_path = argv[++i]; + else if (arg == "--output" && i + 1 < argc) cfg.output_path = argv[++i]; else if (arg == "--vocab-size" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --vocab-size: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --vocab-size: " << v.error().message << "\n"; + std::exit(1); + } if (*v < 258) { // BPE 词表 = 256 字节 + BOS/EOS；更小值会在 @@ -89,17 +95,24 @@ Config parse_args(int argc, char *argv[]) else if (arg == "--min-freq" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --min-freq: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --min-freq: " << v.error().message << "\n"; + std::exit(1); + } cfg.min_freq = *v; } else if (arg == "--threads" && i + 1 < argc) { auto v = nn::parse_number(argv[++i]); - if (!v) { std::cerr << "无效 --threads: " << v.error().message << "\n"; std::exit(1); } + if (!v) + { + std::cerr << "无效 --threads: " << v.error().message << "\n"; + std::exit(1); + } cfg.threads = *v; } - else if (!arg.starts_with("--")) - cfg.text_path = arg; + else if (!arg.starts_with("--")) cfg.text_path = arg; else { std::cerr << "未知参数: " << arg << "\n使用 --help 查看用法\n"; @@ -115,7 +128,7 @@ Config parse_args(int argc, char *argv[]) } // ==================== 主函数 ==================== -int main(int argc, char *argv[]) +int main(int argc, char* argv[]) { Config cfg = parse_args(argc, argv); @@ -146,13 +159,11 @@ int main(int argc, char *argv[]) // ── 训练 ────────────────────────────────────────────────── std::cout << "========================================\n"; - std::cout << " " << (cfg.tokenizer_type == "charbpe" ? "CharBPE" : "BPE") - << " 分词器训练\n"; + std::cout << " " << (cfg.tokenizer_type == "charbpe" ? "CharBPE" : "BPE") << " 分词器训练\n"; std::cout << "========================================\n"; std::cout << " 目标词表: " << cfg.vocab_size << "\n"; std::cout << " 最小频率: " << cfg.min_freq << "\n"; - std::cout << " 预分词线程: " << cfg.threads - << (cfg.threads == 0 ? " (自动)" : "") << "\n"; + std::cout << " 预分词线程: " << cfg.threads << (cfg.threads == 0 ? " (自动)" : "") << "\n"; std::cout << "========================================\n" << std::endl; // 根据分词器类型调用对应训练接口（BPE 和 CharBPE 接口签名一致） @@ -187,8 +198,8 @@ int main(int argc, char *argv[]) auto t_end = std::chrono::steady_clock::now(); Scalar train_sec = std::chrono::duration(t_end - t_start).count(); - std::cout << "\n训练完成! 耗时: " << std::fixed << std::setprecision(1) - << train_sec << "s" << std::endl; + std::cout << "\n训练完成! 耗时: " << std::fixed << std::setprecision(1) << train_sec << "s" + << std::endl; std::cout << "词表大小: " << tokenizer->vocab_size() << std::endl; // ── 保存词表 ───────────────────────────────────────────── @@ -205,9 +216,8 @@ int main(int argc, char *argv[]) sample = nl == std::string::npos ? text : text.substr(0, nl); } auto ids = tokenizer->encode(sample); - std::cout << "\n示例编码 (" << sample.size() << " 字节 → " - << ids.size() << " tokens, 压缩率 " - << std::fixed << std::setprecision(2) + std::cout << "\n示例编码 (" << sample.size() << " 字节 → " << ids.size() + << " tokens, 压缩率 " << std::fixed << std::setprecision(2) << static_cast(sample.size()) / std::max(ids.size(), 1) << " bytes/token):\n"; std::cout << " 原文: \"" << sample << "\"\n"; diff --git a/src/train_state_test.cpp b/src/train_state_test.cpp index 80b6ed1b..9d7a86fb 100644 --- a/src/train_state_test.cpp +++ b/src/train_state_test.cpp @@ -57,25 +57,21 @@ std::filesystem::path scratch_dir() } [[nodiscard]] nn::ModelSpec tiny_spec() -{ - return nn::make_gpt_spec(VOCAB, D_MODEL, SEQ_LEN, HEADS, D_FF, LAYERS); -} +{ return nn::make_gpt_spec(VOCAB, D_MODEL, SEQ_LEN, HEADS, D_FF, LAYERS); } -[[nodiscard]] nn::Result make_model(nn::ComputeEngine &engine) -{ - return nn::build_gpt_model_from_spec(engine, tiny_spec(), nn::PrecisionProfile{}); -} +[[nodiscard]] nn::Result make_model(nn::ComputeEngine& engine) +{ return nn::build_gpt_model_from_spec(engine, tiny_spec(), nn::PrecisionProfile{}); } // ── 参数快照 / 逐位比较（[6] 用；顺序 = parameters()，逐元素 ==）───────── -[[nodiscard]] std::vector snapshot_params(nn::ComputeEngine &engine, nn::Model &m) +[[nodiscard]] std::vector snapshot_params(nn::ComputeEngine& engine, nn::Model& m) { std::vector out; - for (auto &t : m.parameters()) out.push_back(download(engine, t.get())); + for (auto& t : m.parameters()) + out.push_back(download(engine, t.get())); return out; } -[[nodiscard]] bool same_matrices(const std::vector &a, - const std::vector &b) +[[nodiscard]] bool same_matrices(const std::vector& a, const std::vector& b) { if (a.size() != b.size()) return false; for (std::size_t i = 0; i < a.size(); ++i) @@ -130,13 +126,12 @@ std::filesystem::path scratch_dir() return s; } -[[nodiscard]] bool same_state(const nn::TrainState &a, const nn::TrainState &b) +[[nodiscard]] bool same_state(const nn::TrainState& a, const nn::TrainState& b) { if (std::string(a.state_kind()) != std::string(b.state_kind())) return false; return a.counters.global_step == b.counters.global_step && a.counters.epoch == b.counters.epoch && - a.counters.step_in_epoch == b.counters.step_in_epoch && - a.counters.lr == b.counters.lr && + a.counters.step_in_epoch == b.counters.step_in_epoch && a.counters.lr == b.counters.lr && a.optimizer_step == b.optimizer_step && a.optimizer_blob_count == b.optimizer_blob_count && a.recipe.precision.param == b.recipe.precision.param && @@ -144,12 +139,10 @@ std::filesystem::path scratch_dir() a.recipe.precision.stable == b.recipe.precision.stable && a.recipe.precision.optimizer == b.recipe.precision.optimizer && a.recipe.optimizer == b.recipe.optimizer && - a.recipe.lr_schedule == b.recipe.lr_schedule && - a.recipe.lr == b.recipe.lr && a.recipe.min_lr == b.recipe.min_lr && - a.recipe.weight_decay == b.recipe.weight_decay && + a.recipe.lr_schedule == b.recipe.lr_schedule && a.recipe.lr == b.recipe.lr && + a.recipe.min_lr == b.recipe.min_lr && a.recipe.weight_decay == b.recipe.weight_decay && a.recipe.beta1 == b.recipe.beta1 && a.recipe.beta2 == b.recipe.beta2 && - a.recipe.max_norm == b.recipe.max_norm && - a.recipe.epochs == b.recipe.epochs && + a.recipe.max_norm == b.recipe.max_norm && a.recipe.epochs == b.recipe.epochs && a.recipe.total_steps == b.recipe.total_steps && a.recipe.warmup_epochs == b.recipe.warmup_epochs && a.recipe.warmup_steps == b.recipe.warmup_steps && @@ -164,8 +157,7 @@ std::filesystem::path scratch_dir() const auto* tb = dynamic_cast(&b); if ((ta == nullptr) != (tb == nullptr)) return false; if (ta == nullptr) return true; // 两份都无文本扩展 ⇒ 基类字段一致即一致 - return ta->stride == tb->stride && - ta->dataset.path == tb->dataset.path && + return ta->stride == tb->stride && ta->dataset.path == tb->dataset.path && ta->dataset.sha256 == tb->dataset.sha256 && ta->dataset.num_tokens == tb->dataset.num_tokens && ta->dataset.num_docs == tb->dataset.num_docs && @@ -175,19 +167,19 @@ std::filesystem::path scratch_dir() } // ── 宿主文件辅助（测试自用；库内字节 cast 收敛在 core_file.hpp）───────── -[[nodiscard]] std::string read_file(const std::filesystem::path &p) +[[nodiscard]] std::string read_file(const std::filesystem::path& p) { std::ifstream ifs(p, std::ios::binary); return std::string{std::istreambuf_iterator(ifs), std::istreambuf_iterator()}; } -void write_file(const std::filesystem::path &p, const std::string &bytes) +void write_file(const std::filesystem::path& p, const std::string& bytes) { std::ofstream ofs(p, std::ios::binary); ofs.write(bytes.data(), static_cast(bytes.size())); } -[[nodiscard]] std::uint64_t rd_u64(const std::string &b, std::size_t off) +[[nodiscard]] std::uint64_t rd_u64(const std::string& b, std::size_t off) { std::uint64_t v = 0; std::memcpy(&v, b.data() + off, sizeof(v)); @@ -206,7 +198,8 @@ void test_kvrec_roundtrip() auto back = nn::train_state_from_kv(*kv); CHECK(back.has_value(), "TrainState 读回成功"); if (!back) return; - CHECK(back && same_state(s, **back), "全字段逐项一致（含精度四槽 / lr_per_epoch / 数据集身份）"); + CHECK(back && same_state(s, **back), + "全字段逐项一致（含精度四槽 / lr_per_epoch / 数据集身份）"); // 缺键 → 默认值（老文件兜底口径） nn::KeyValueRecord empty; @@ -223,7 +216,7 @@ void test_kvrec_roundtrip() } // ── [2][3][4][6][8] 落盘往返与 role 门禁 ─────────────────────────────── -void test_checkpoint_roundtrip(const std::filesystem::path &dir) +void test_checkpoint_roundtrip(const std::filesystem::path& dir) { std::cout << "\n── [2] Checkpoint 往返（状态段 + 优化器张量）──\n"; auto engine_owner = nn::CpuEngine::create(); @@ -231,16 +224,23 @@ void test_checkpoint_roundtrip(const std::filesystem::path &dir) auto model_r = make_model(engine); CHECK(model_r.has_value(), "建模型"); if (!model_r) return; - nn::Model &model = *model_r; + nn::Model& model = *model_r; const nn::ModelSpec spec = tiny_spec(); const auto params_before = snapshot_params(engine, model); // 优化器状态张量：一个 f32（值可精确比较）+ 一个 f16（验证精度记录） nn::Matrix m32(3, 2); - m32.span()[0] = 0.5f; m32.span()[1] = 1.5f; m32.span()[2] = -2.0f; - m32.span()[3] = 0.25f; m32.span()[4] = 8.0f; m32.span()[5] = -0.125f; + m32.span()[0] = 0.5f; + m32.span()[1] = 1.5f; + m32.span()[2] = -2.0f; + m32.span()[3] = 0.25f; + m32.span()[4] = 8.0f; + m32.span()[5] = -0.125f; nn::Matrix m16(2, 2); - m16.span()[0] = 0.5f; m16.span()[1] = 1.0f; m16.span()[2] = 2.0f; m16.span()[3] = -3.0f; + m16.span()[0] = 0.5f; + m16.span()[1] = 1.0f; + m16.span()[2] = 2.0f; + m16.span()[3] = -3.0f; std::vector blobs{ upload(engine, m32, nn::Precision::F32), upload(engine, m16, nn::Precision::F16), @@ -251,10 +251,13 @@ void test_checkpoint_roundtrip(const std::filesystem::path &dir) s.optimizer_blob_count = blobs.size(); // 写侧按实参回填，读回应等于它 nn::CheckpointState ckpt{&s, blobs}; const auto ckpt_path = dir / "state.nnckpt"; - auto sr = nn::save_model(ckpt_path.string(), model, spec, tok, - nn::ModelRole::Checkpoint, ckpt); + auto sr = nn::save_model(ckpt_path.string(), model, spec, tok, nn::ModelRole::Checkpoint, ckpt); CHECK(sr.has_value(), "写 Checkpoint（含状态段）"); - if (!sr) { std::cout << " " << sr.error().message << "\n"; return; } + if (!sr) + { + std::cout << " " << sr.error().message << "\n"; + return; + } auto pr = nn::peek_train_state(ckpt_path.string()); CHECK(pr.has_value(), "peek_train_state 成功"); @@ -379,13 +382,17 @@ void test_checkpoint_roundtrip(const std::filesystem::path &dir) } // ── [5] 尾部校验和（P2-1）────────────────────────────────────────────── -void test_integrity(const std::filesystem::path &dir) +void test_integrity(const std::filesystem::path& dir) { std::cout << "\n── [5] 整文件校验和 / 尾部完整性标记 ──\n"; auto engine_owner = nn::CpuEngine::create(); nn::CpuEngine& engine = *engine_owner; auto model_r = make_model(engine); - if (!model_r) { CHECK(false, "建模型"); return; } + if (!model_r) + { + CHECK(false, "建模型"); + return; + } const auto path = dir / "integrity.nnmodel"; auto sr = nn::save_model(path.string(), *model_r, tiny_spec(), "tok", nn::ModelRole::Model); CHECK(sr.has_value(), "写文件"); @@ -420,12 +427,12 @@ void test_integrity(const std::filesystem::path &dir) auto tv = nn::verify_model_integrity(trunc_path.string()); CHECK(!tv.has_value(), "截断 → 明确报错"); if (!tv) - CHECK(tv.error().message.find("NNEND") != std::string::npos, - "文案点明缺少尾部完整性标记"); + CHECK(tv.error().message.find("NNEND") != std::string::npos, "文案点明缺少尾部完整性标记"); { auto back = make_model(engine); if (back) - CHECK(!nn::load_model(trunc_path.string(), *back).has_value(), "截断文件 load_model 失败"); + CHECK(!nn::load_model(trunc_path.string(), *back).has_value(), + "截断文件 load_model 失败"); } } @@ -434,7 +441,7 @@ void test_integrity(const std::filesystem::path &dir) // kvrec 字段顺序固定、摘要覆盖确定、参数载荷同精度回写。 // （f16 状态张量经 load 物化为 f32 ⇒ 再 save 会被提升为 f32：值逐位不变、 // 编码变化；生产路径不存在这种再写——save 一律从活状态导出，见 23 §2.1.1。） -void test_resave_byte_identity(const std::filesystem::path &dir) +void test_resave_byte_identity(const std::filesystem::path& dir) { std::cout << "\n── [9] save → load → save 逐字节一致（f32 状态）──\n"; auto engine_owner = nn::CpuEngine::create(); @@ -442,11 +449,15 @@ void test_resave_byte_identity(const std::filesystem::path &dir) auto model_r = make_model(engine); CHECK(model_r.has_value(), "建模型"); if (!model_r) return; - nn::Model &model = *model_r; + nn::Model& model = *model_r; nn::Matrix m(2, 3); - m.span()[0] = 0.125f; m.span()[1] = -1.5f; m.span()[2] = 3.0f; - m.span()[3] = -0.25f; m.span()[4] = 7.5f; m.span()[5] = 0.5f; + m.span()[0] = 0.125f; + m.span()[1] = -1.5f; + m.span()[2] = 3.0f; + m.span()[3] = -0.25f; + m.span()[4] = 7.5f; + m.span()[5] = 0.5f; std::vector blobs{upload(engine, m, nn::Precision::F32)}; nn::TextTrainState st = sample_state(); @@ -455,22 +466,21 @@ void test_resave_byte_identity(const std::filesystem::path &dir) const auto path1 = dir / "resave1.nnckpt"; const auto path2 = dir / "resave2.nnckpt"; CHECK(nn::save_model(path1.string(), model, tiny_spec(), tok, nn::ModelRole::Checkpoint, - nn::CheckpointState{&st, blobs}).has_value(), + nn::CheckpointState{&st, blobs}) + .has_value(), "写第一份断点"); auto loaded = nn::load_train_state(engine, path1.string()); CHECK(loaded.has_value(), "读回状态 + 状态张量"); if (loaded) { - CHECK(nn::save_model(path2.string(), model, tiny_spec(), tok, - nn::ModelRole::Checkpoint, - nn::CheckpointState{loaded->state.get(), - loaded->optimizer_blobs}).has_value(), + CHECK(nn::save_model(path2.string(), model, tiny_spec(), tok, nn::ModelRole::Checkpoint, + nn::CheckpointState{loaded->state.get(), loaded->optimizer_blobs}) + .has_value(), "用读回的内容重写第二份断点"); const std::string b1 = read_file(path1); const std::string b2 = read_file(path2); - CHECK(b1.size() == b2.size() && b1 == b2, - "两份文件逐字节一致（save → load → save）"); + CHECK(b1.size() == b2.size() && b1 == b2, "两份文件逐字节一致（save → load → save）"); CHECK(b1.size() > nn::MODEL_TAIL_BYTES, "文件含尾部标记"); } } diff --git a/src/training_group_session_test.cpp b/src/training_group_session_test.cpp index 5a775998..c46869a3 100644 --- a/src/training_group_session_test.cpp +++ b/src/training_group_session_test.cpp @@ -22,14 +22,12 @@ namespace nn::TensorSchema role_schema(std::string name) { - auto r = nn::TensorSchema::create( - std::move(name), {nn::tensor_field("output")}); + auto r = nn::TensorSchema::create(std::move(name), {nn::tensor_field("output")}); NN_ASSERT_ALWAYS(r, "role_schema"); return std::move(*r); } -nn::Result make_model(nn::CpuEngine& engine, std::string role, - std::size_t in = 3, +nn::Result make_model(nn::CpuEngine& engine, std::string role, std::size_t in = 3, std::size_t out = 2) { nn::Model model(engine); @@ -38,8 +36,7 @@ nn::Result make_model(nn::CpuEngine& engine, std::string role, return model; } -nn::Result make_session( - nn::CpuEngine& engine, std::string role) +nn::Result make_session(nn::CpuEngine& engine, std::string role) { NN_TRY(model, make_model(engine, std::move(role))); nn::TrainingConfig cfg; @@ -53,25 +50,25 @@ void train_one_microstep(nn::TrainingSession& session) const auto x = make_tensor(session.engine(), 3, 2, 0.2f, 0.1f); nn::TrainingStepOptions opts; opts.normalization_units = 2; - auto report = session.train_step( - x, opts, - [](nn::Model& model, nn::ComputeEngine& engine, - const nn::Tensor& input) -> nn::Result { - NN_TRY(out, model.forward(input)); - nn::MSELoss loss; - auto target = make_tensor(engine, out->rows(), out->cols(), - 0.0f, 0.0f); - NN_TRY_CHECK(loss.forward(engine, *out, target)); - NN_TRY(grad, loss.backward()); - return nn::TrainingStepOutput{ - .grad = std::move(*grad), - .gradient_units = 2, - .loss_sum = nn::Tensor{}, - .loss_units = 0, - }; - }); - CHECK(report && report->parameters_updated, - "group role executes one optimizer update"); + auto report = + session.train_step(x, opts, + [](nn::Model& model, nn::ComputeEngine& engine, + const nn::Tensor& input) -> nn::Result + { + NN_TRY(out, model.forward(input)); + nn::MSELoss loss; + auto target = + make_tensor(engine, out->rows(), out->cols(), 0.0f, 0.0f); + NN_TRY_CHECK(loss.forward(engine, *out, target)); + NN_TRY(grad, loss.backward()); + return nn::TrainingStepOutput{ + .grad = std::move(*grad), + .gradient_units = 2, + .loss_sum = nn::Tensor{}, + .loss_units = 0, + }; + }); + CHECK(report && report->parameters_updated, "group role executes one optimizer update"); } void test_roles_and_lifecycle() @@ -88,26 +85,20 @@ void test_roles_and_lifecycle() "add actor training role"); CHECK(group.add_training_role("critic", std::move(*critic)).has_value(), "add critic training role"); - CHECK(group.add_readonly_role("reference", - std::move(*reference)).has_value(), + CHECK(group.add_readonly_role("reference", std::move(*reference)).has_value(), "add readonly reference role"); CHECK(group.size() == 3, "group role count"); - CHECK((group.role_names() == - std::vector{"actor", "critic", "reference"}), + CHECK((group.role_names() == std::vector{"actor", "critic", "reference"}), "group preserves add order"); - CHECK(group.engine_bound() && &group.engine() == engine.get(), - "group reports shared engine"); + CHECK(group.engine_bound() && &group.engine() == engine.get(), "group reports shared engine"); auto actor_trainable = group.is_trainable("actor"); auto ref_trainable = group.is_trainable("reference"); CHECK(actor_trainable && *actor_trainable, "actor is trainable"); CHECK(ref_trainable && !*ref_trainable, "reference is readonly"); - CHECK(group.training_session("actor").has_value(), - "training role resolves Session"); - CHECK(!group.training_session("reference").has_value(), - "readonly role cannot resolve Session"); - CHECK(group.model("reference").has_value(), - "readonly role resolves Model"); + CHECK(group.training_session("actor").has_value(), "training role resolves Session"); + CHECK(!group.training_session("reference").has_value(), "readonly role cannot resolve Session"); + CHECK(group.model("reference").has_value(), "readonly role resolves Model"); CHECK(!group.model("missing").has_value(), "unknown role rejected"); auto actor_schema = group.schema("actor"); @@ -118,12 +109,10 @@ void test_roles_and_lifecycle() "different role schemas may coexist"); nn::TrainingGroupSession moved = std::move(group); - CHECK(moved.size() == 3 && group.empty(), - "group move transfers role ownership"); + CHECK(moved.size() == 3 && group.empty(), "group move transfers role ownership"); auto moved_actor = moved.training_session("actor"); CHECK(moved_actor.has_value(), "moved group resolves training role"); - if (moved_actor) - train_one_microstep(**moved_actor); + if (moved_actor) train_one_microstep(**moved_actor); } void test_rejections() @@ -133,26 +122,21 @@ void test_rejections() CHECK(actor.has_value(), "rejection actor"); if (!actor) return; nn::TrainingGroupSession group; - CHECK(group.add_training_role("actor", std::move(*actor)).has_value(), - "rejection actor added"); + CHECK(group.add_training_role("actor", std::move(*actor)).has_value(), "rejection actor added"); auto duplicate = make_session(*engine, "actor"); CHECK(duplicate.has_value(), "duplicate source session"); if (duplicate) - CHECK(!group.add_training_role("actor", std::move(*duplicate)), - "duplicate role rejected"); + CHECK(!group.add_training_role("actor", std::move(*duplicate)), "duplicate role rejected"); auto blank = make_session(*engine, "blank"); CHECK(blank.has_value(), "blank source session"); - if (blank) - CHECK(!group.add_training_role("", std::move(*blank)), - "empty role rejected"); + if (blank) CHECK(!group.add_training_role("", std::move(*blank)), "empty role rejected"); auto other_engine = nn::CpuEngine::create(); auto other = make_session(*other_engine, "other"); CHECK(other.has_value(), "other-engine session"); if (other) - CHECK(!group.add_training_role("other", std::move(*other)), - "cross-engine role rejected"); + CHECK(!group.add_training_role("other", std::move(*other)), "cross-engine role rejected"); auto readonly = make_model(*engine, "help"); CHECK(readonly.has_value(), "readonly model"); @@ -160,8 +144,7 @@ void test_rejections() { CHECK(group.add_readonly_role("help", std::move(*readonly)).has_value(), "readonly role added"); - CHECK(!group.add_readonly_role("help", - NN_UNWRAP(make_model(*engine, "help2"))), + CHECK(!group.add_readonly_role("help", NN_UNWRAP(make_model(*engine, "help2"))), "duplicate readonly role rejected"); } } diff --git a/src/training_session_test.cpp b/src/training_session_test.cpp index 09742f30..14d1779c 100644 --- a/src/training_session_test.cpp +++ b/src/training_session_test.cpp @@ -31,7 +31,8 @@ std::vector snapshot_params(nn::CpuEngine& eng, nn::Model& model) { auto m = eng.to_matrix(t.get(), nn::Precision::F32); NN_ASSERT_ALWAYS(m, "snapshot_params: to_matrix failed"); - for (auto v : m->span()) out.push_back(v); + for (auto v : m->span()) + out.push_back(v); } return out; } @@ -99,30 +100,32 @@ void test_create_and_train() CHECK(infer.has_value(), "session.model().forward 推理可用"); } -[[nodiscard]] nn::Result -run_mse_step(nn::TrainingSession& session, const nn::Tensor& x, - std::size_t normalization_units, bool force_update = false) +[[nodiscard]] nn::Result run_mse_step(nn::TrainingSession& session, + const nn::Tensor& x, + std::size_t normalization_units, + bool force_update = false) { nn::TrainingStepOptions opts; opts.normalization_units = normalization_units; opts.force_update = force_update; - return session.train_step( - x, opts, - [&](nn::Model& model, nn::ComputeEngine& engine, const nn::Tensor& input) - -> nn::Result { - auto logits_r = model.forward(input); - NN_TRY_CHECK(logits_r); - nn::MSELoss loss; - auto target = make_tensor(engine, logits_r->rows(), logits_r->cols(), 0.f, 0.f); - auto loss_r = loss.forward(engine, *logits_r, target); - NN_TRY_CHECK(loss_r); - auto grad_r = loss.backward(); - NN_TRY_CHECK(grad_r); - nn::TrainingStepOutput out; - out.grad = std::move(*grad_r); - out.gradient_units = logits_r->size(); - return out; - }); + return session.train_step(x, opts, + [&](nn::Model& model, nn::ComputeEngine& engine, + const nn::Tensor& input) -> nn::Result + { + auto logits_r = model.forward(input); + NN_TRY_CHECK(logits_r); + nn::MSELoss loss; + auto target = make_tensor(engine, logits_r->rows(), + logits_r->cols(), 0.f, 0.f); + auto loss_r = loss.forward(engine, *logits_r, target); + NN_TRY_CHECK(loss_r); + auto grad_r = loss.backward(); + NN_TRY_CHECK(grad_r); + nn::TrainingStepOutput out; + out.grad = std::move(*grad_r); + out.gradient_units = logits_r->size(); + return out; + }); } void test_accumulation_and_session_state() @@ -158,14 +161,16 @@ void test_global_batch_normalization() { auto engine = nn::CpuEngine::create(); - auto make_zero_model = [&]() { + auto make_zero_model = [&]() + { nn::Model model(*engine); NN_ASSERT_ALWAYS(model.add(1, 1), "norm: add"); for (auto& p : model.parameters()) NN_ASSERT_ALWAYS(engine->zero(p.get()), "norm: zero parameter"); return model; }; - auto make_sgd_session = [&](std::size_t accum) { + auto make_sgd_session = [&](std::size_t accum) + { nn::TrainingConfig cfg; cfg.optimizer = "sgd"; cfg.lr = 0.1f; @@ -230,23 +235,24 @@ void test_before_update_abort() nn::TrainingStepOptions opts; opts.normalization_units = 6; opts.before_update = []() -> nn::Result { return true; }; - auto r = s->train_step( - x, std::move(opts), - [&](nn::Model& m, nn::ComputeEngine& e, const nn::Tensor& input) - -> nn::Result { - auto logits = m.forward(input); - NN_TRY_CHECK(logits); - nn::MSELoss loss; - auto target = make_tensor(e, logits->rows(), logits->cols(), 0.f, 0.f); - auto l = loss.forward(e, *logits, target); - NN_TRY_CHECK(l); - auto g = loss.backward(); - NN_TRY_CHECK(g); - nn::TrainingStepOutput out; - out.grad = std::move(*g); - out.gradient_units = logits->size(); - return out; - }); + auto r = s->train_step(x, std::move(opts), + [&](nn::Model& m, nn::ComputeEngine& e, + const nn::Tensor& input) -> nn::Result + { + auto logits = m.forward(input); + NN_TRY_CHECK(logits); + nn::MSELoss loss; + auto target = + make_tensor(e, logits->rows(), logits->cols(), 0.f, 0.f); + auto l = loss.forward(e, *logits, target); + NN_TRY_CHECK(l); + auto g = loss.backward(); + NN_TRY_CHECK(g); + nn::TrainingStepOutput out; + out.grad = std::move(*g); + out.gradient_units = logits->size(); + return out; + }); CHECK(r && r->aborted && !r->parameters_updated, "abort: 回调可在 update 前中止"); CHECK(s->step_count() == 0, "abort: optimizer step 未推进"); CHECK(s->state().microsteps_in_update == 0 && !s->state().accumulation_dirty, @@ -315,8 +321,7 @@ void test_precision_passthrough() if (!session_r) return; nn::TrainingSession session = std::move(*session_r); - CHECK(session.config().precision.compute == nn::Precision::F16, - "config 保留调用方精度配方"); + CHECK(session.config().precision.compute == nn::Precision::F16, "config 保留调用方精度配方"); CHECK(session.model().parameters()[0].get().precision() == nn::Precision::F16, "模型参数按 f16 存储"); @@ -395,9 +400,7 @@ int main() test_errors(); test_destroy_order(); - if (g_fail == 0) - std::printf("training_session_test: ALL PASS\n"); - else - std::printf("training_session_test: %d failure(s)\n", g_fail); + if (g_fail == 0) std::printf("training_session_test: ALL PASS\n"); + else std::printf("training_session_test: %d failure(s)\n", g_fail); return g_fail == 0 ? 0 : 1; } diff --git a/src/write_epoch_test.cpp b/src/write_epoch_test.cpp index 9411b730..4bdad4ce 100644 --- a/src/write_epoch_test.cpp +++ b/src/write_epoch_test.cpp @@ -48,7 +48,8 @@ nn::Tensor zeros(nn::ComputeEngine& eng, std::size_t rows, std::size_t cols) nn::Tensor filled(nn::ComputeEngine& eng, std::size_t rows, std::size_t cols, float v) { nn::Matrix m(rows, cols); - for (auto& x : m.span()) x = v; + for (auto& x : m.span()) + x = v; return upload(eng, m); } @@ -58,8 +59,7 @@ void expect_bump(nn::ComputeEngine& eng, const char* what, nn::Tensor& t, Fn&& f { const std::uint64_t before = t.write_epoch(); auto r = fn(); - if (!r.has_value()) - std::printf("[FAIL] %s: 调用失败 —— %s\n", what, r.error().message.c_str()); + if (!r.has_value()) std::printf("[FAIL] %s: 调用失败 —— %s\n", what, r.error().message.c_str()); CHECK(r.has_value(), what); const std::uint64_t after = t.write_epoch(); if (!(after > before)) @@ -70,13 +70,11 @@ void expect_bump(nn::ComputeEngine& eng, const char* what, nn::Tensor& t, Fn&& f (void)eng; } -template -void expect_no_bump(const char* what, const nn::Tensor& t, Fn&& fn) +template void expect_no_bump(const char* what, const nn::Tensor& t, Fn&& fn) { const std::uint64_t before = t.write_epoch(); auto r = fn(); - if (!r.has_value()) - std::printf("[FAIL] %s: 调用失败 —— %s\n", what, r.error().message.c_str()); + if (!r.has_value()) std::printf("[FAIL] %s: 调用失败 —— %s\n", what, r.error().message.c_str()); CHECK(r.has_value(), what); if (t.write_epoch() != before) std::printf("[FAIL] %s: 只读路径推进了写版本（before=%llu after=%llu）\n", what, @@ -96,102 +94,62 @@ void test_write_entries(nn::ComputeEngine& eng, const char* tag) nn::Tensor src = filled(eng, 3, 2, 1.0f); nn::Tensor idx = filled(eng, 1, 1, 0.0f); // 行索引 0 nn::Matrix m(3, 2); - for (auto& x : m.span()) x = 0.5f; + for (auto& x : m.span()) + x = 0.5f; std::vector flat(6, 0.25f); - expect_bump(eng, "write 推进写版本", t, [&] - { - return eng.write(t, std::span(flat)); - }); - expect_bump(eng, "set_index 推进写版本", t, [&] - { - return eng.set_index(t, 0, 0, 1.0f); - }); - expect_bump(eng, "copy_from 推进写版本", t, [&] - { - return eng.copy_from(t, m); - }); - expect_bump(eng, "accumulate 推进写版本", t, [&] - { - return eng.accumulate(t, src); - }); - expect_bump(eng, "add_inplace 推进写版本", t, [&] - { - return eng.add_inplace(t, src); - }); - expect_bump(eng, "scale_inplace 推进写版本", t, [&] - { - return eng.scale_inplace(t, 2.0f); - }); + expect_bump(eng, "write 推进写版本", t, + [&] { return eng.write(t, std::span(flat)); }); + expect_bump(eng, "set_index 推进写版本", t, [&] { return eng.set_index(t, 0, 0, 1.0f); }); + expect_bump(eng, "copy_from 推进写版本", t, [&] { return eng.copy_from(t, m); }); + expect_bump(eng, "accumulate 推进写版本", t, [&] { return eng.accumulate(t, src); }); + expect_bump(eng, "add_inplace 推进写版本", t, [&] { return eng.add_inplace(t, src); }); + expect_bump(eng, "scale_inplace 推进写版本", t, [&] { return eng.scale_inplace(t, 2.0f); }); { // insert_rows 要求 dst 能容纳 src 的全部行 nn::Tensor tall = zeros(eng, 4, 2); - expect_bump(eng, "insert_rows 推进写版本", tall, [&] - { - return eng.insert_rows(tall, 1, src); - }); + expect_bump(eng, "insert_rows 推进写版本", tall, + [&] { return eng.insert_rows(tall, 1, src); }); } - expect_bump(eng, "scatter_add_rows 推进写版本", t, [&] - { - return eng.scatter_add_rows(t, idx, src); - }); - expect_bump(eng, "eval_expr_into（DSL）推进写版本", t, [&] - { - // 优化器更新参数走的就是这条（dsl::compute_into → eval_expr_into） - return nn::dsl::compute_into(eng, - nn::dsl::leaf(t) + nn::dsl::leaf(src), t); - }); + expect_bump(eng, "scatter_add_rows 推进写版本", t, + [&] { return eng.scatter_add_rows(t, idx, src); }); + expect_bump(eng, "eval_expr_into（DSL）推进写版本", t, + [&] + { + // 优化器更新参数走的就是这条（dsl::compute_into → eval_expr_into） + return nn::dsl::compute_into(eng, nn::dsl::leaf(t) + nn::dsl::leaf(src), t); + }); { // copy_into / cast_into：目标传递的其余两个出口 nn::Tensor dst = zeros(eng, 3, 2); - expect_bump(eng, "copy_into 推进写版本", dst, [&] - { - return eng.copy_into(dst, src); - }); + expect_bump(eng, "copy_into 推进写版本", dst, [&] { return eng.copy_into(dst, src); }); nn::Tensor f16src = upload(eng, m, nn::Precision::F16); nn::Tensor f32dst = zeros(eng, 3, 2); - expect_bump(eng, "cast_into 推进写版本", f32dst, [&] - { - return eng.cast_into(f16src, f32dst); - }); + expect_bump(eng, "cast_into 推进写版本", f32dst, + [&] { return eng.cast_into(f16src, f32dst); }); } - expect_bump(eng, "zero 推进写版本", t, [&] - { - return eng.zero(t); - }); + expect_bump(eng, "zero 推进写版本", t, [&] { return eng.zero(t); }); // ── [2] 只读路径不推进 ─────────────────────────────────────────────── std::printf("[2] 只读路径不推进写版本（%s）\n", tag); - expect_no_bump("read 不推进写版本", t, [&] - { - std::vector buf(6); - return eng.read(t, std::span(buf)); - }); - expect_no_bump("to_matrix 不推进写版本", t, [&] - { - return eng.to_matrix(t, nn::Precision::F32); - }); - expect_no_bump("clone 不推进源写版本", t, [&] - { - return eng.clone(t); - }); + expect_no_bump("read 不推进写版本", t, + [&] + { + std::vector buf(6); + return eng.read(t, std::span(buf)); + }); + expect_no_bump("to_matrix 不推进写版本", t, + [&] { return eng.to_matrix(t, nn::Precision::F32); }); + expect_no_bump("clone 不推进源写版本", t, [&] { return eng.clone(t); }); { // 合法形状：A (2,3) × B (3,2)（t/src 是 (3,2)，不足以相乘） nn::Tensor a = filled(eng, 2, 3, 1.0f); nn::Tensor b = filled(eng, 3, 2, 1.0f); - expect_no_bump("matmul 不推进 A 操作数写版本", a, [&] - { - return eng.matmul(a, b); - }); - expect_no_bump("matmul 不推进 B 操作数写版本", b, [&] - { - return eng.matmul(a, b); - }); + expect_no_bump("matmul 不推进 A 操作数写版本", a, [&] { return eng.matmul(a, b); }); + expect_no_bump("matmul 不推进 B 操作数写版本", b, [&] { return eng.matmul(a, b); }); } expect_no_bump("dsl::compute 不推进输入写版本", t, [&] - { - return nn::dsl::compute(eng, nn::dsl::leaf(t) * nn::dsl::leaf(src), 3, 2); - }); + { return nn::dsl::compute(eng, nn::dsl::leaf(t) * nn::dsl::leaf(src), 3, 2); }); // ── [3] 别名可见：拷贝共享存储 ─────────────────────────────────────── std::printf("[3] 别名（Tensor 拷贝）上的写入对原句柄可见（%s）\n", tag); @@ -227,13 +185,9 @@ nn::BitLinear make_bitlinear(nn::ComputeEngine& eng, std::size_t in, std::size_t // (τ, wq) 的宿主快照，用于逐位比较 nn::Matrix wq_of(nn::ComputeEngine& eng, nn::BitLinear& bl) -{ - return download(eng, bl.quantized_weights()); -} +{ return download(eng, bl.quantized_weights()); } nn::Matrix tau_of(nn::ComputeEngine& eng, nn::BitLinear& bl) -{ - return download(eng, bl.row_scales()); -} +{ return download(eng, bl.row_scales()); } bool mat_same(const nn::Matrix& a, const nn::Matrix& b) { @@ -273,8 +227,7 @@ void test_bitlinear_cache(nn::ComputeEngine& eng, const char* tag) // 同权重、同输入 → 命中缓存（且输出逐位不变） auto y2 = bl.forward(Xt); CHECK(y2.has_value(), "第二次 forward 成功"); - if (y2.has_value()) - CHECK(mat_same(y0, download(eng, *y2)), "缓存命中时输出逐位不变"); + if (y2.has_value()) CHECK(mat_same(y0, download(eng, *y2)), "缓存命中时输出逐位不变"); CHECK(bl.quantize_calls() == 1, "同权重下第二次 forward 不重算 τ/wq"); // 换输入（权重没变）→ 仍不重算 @@ -309,8 +262,7 @@ void test_bitlinear_cache(nn::ComputeEngine& eng, const char* tag) CHECK(bl.quantize_calls() == 2, "优化器更新权重后重算一次"); const nn::Matrix wq1 = wq_of(eng, bl); const nn::Matrix tau1 = tau_of(eng, bl); - CHECK(!mat_same(wq0, wq1) || !mat_same(tau0, tau1), - "权重变化后 τ/wq 至少之一变化"); + CHECK(!mat_same(wq0, wq1) || !mat_same(tau0, tau1), "权重变化后 τ/wq 至少之一变化"); // 另一条写入路径：engine.copy_from（序列化加载 / 宿主直写） { @@ -373,8 +325,7 @@ void test_bitlinear_cache(nn::ComputeEngine& eng, const char* tag) CHECK(a.has_value() && b.has_value(), "打包路径 forward ×2 成功"); CHECK(pk.quantize_calls() == 1, "打包路径：同权重两次 forward 只量化一次"); if (a.has_value() && b.has_value()) - CHECK(mat_same(download(eng, *a), download(eng, *b)), - "打包路径缓存命中时输出逐位相同"); + CHECK(mat_same(download(eng, *a), download(eng, *b)), "打包路径缓存命中时输出逐位相同"); } } @@ -447,7 +398,6 @@ int main(int argc, char* argv[]) test_lazy_grads(eng, "CPU"); } - if (g_failures == 0) - std::puts("write_epoch_test: 全部通过"); + if (g_failures == 0) std::puts("write_epoch_test: 全部通过"); return g_failures == 0 ? 0 : 1; } diff --git a/tools/fused_generate.hpp b/tools/fused_generate.hpp index 1a26317b..3d340041 100644 --- a/tools/fused_generate.hpp +++ b/tools/fused_generate.hpp @@ -49,8 +49,8 @@ namespace { -[[nodiscard]] bool run_glslc(const std::string& glslc, - const std::string& src, const std::string& dst) +[[nodiscard]] bool run_glslc(const std::string& glslc, const std::string& src, + const std::string& dst) { // 仅给 glslc 路径加引号；src/dst 为无空格路径（out_dir 由 CMake 控制）， // 正斜杠相对路径不加引号，避免 cmd/system 解析问题。 @@ -69,8 +69,7 @@ namespace const std::streamsize sz = f.tellg(); // tellg() 失败返回 -1：不判就 static_cast(-1) 请求 ~4GB 分配。 // SPIR-V 恒为 4 字节字的整数倍——不是即文件损坏（截断的 glslc 产物）。 - if (sz <= 0 || sz % static_cast(sizeof(std::uint32_t)) != 0) - return {}; + if (sz <= 0 || sz % static_cast(sizeof(std::uint32_t)) != 0) return {}; f.seekg(0, std::ios::beg); std::vector v(static_cast(sz) / sizeof(std::uint32_t)); if (!nn::read_pod_span(f, std::span(v))) @@ -90,8 +89,8 @@ namespace { const auto& in = seq[i]; if (i) o << ","; - o << " {" << static_cast(in.op) << ", " << static_cast(in.dst) - << ", {" << static_cast(in.a.kind) << ", " << static_cast(in.a.idx) << "}" + o << " {" << static_cast(in.op) << ", " << static_cast(in.dst) << ", {" + << static_cast(in.a.kind) << ", " << static_cast(in.a.idx) << "}" << ", {" << static_cast(in.b.kind) << ", " << static_cast(in.b.idx) << "}" << ", {" << static_cast(in.c.kind) << ", " << static_cast(in.c.idx) << "}}"; } @@ -102,8 +101,7 @@ namespace if (std::isinf(v)) o << (v < 0 ? " -std::numeric_limits::infinity()" : " std::numeric_limits::infinity()"); - else - o << " static_cast(" << v << ")"; + else o << " static_cast(" << v << ")"; }; o << "ExprSpec{ \n"; o << " std::vector{"; @@ -113,9 +111,8 @@ namespace { const auto& v = spec.views[i]; if (i) o << ","; - o << " {" << static_cast(v.kind) << ", " - << static_cast(v.negate_first_half) << ", " << v.param << "u, " - << v.param2 << "u}"; + o << " {" << static_cast(v.kind) << ", " << static_cast(v.negate_first_half) + << ", " << v.param << "u, " << v.param2 << "u}"; } o << "},\n std::vector{"; for (std::size_t i = 0; i < spec.consts.size(); ++i) @@ -133,11 +130,10 @@ namespace if (spec.matmul) { o << "MatmulSpec{" << static_cast(spec.matmul->a_input) << ", " - << static_cast(spec.matmul->b_input) << ", " - << static_cast(spec.matmul->transA) << ", " - << static_cast(spec.matmul->transB) << ", " - << spec.matmul->k << "u, " - << spec.matmul->batch << "u}"; // batch 必须显式生成（漏则退化为默认 1，registry spec 失真） + << static_cast(spec.matmul->b_input) << ", " << static_cast(spec.matmul->transA) + << ", " << static_cast(spec.matmul->transB) << ", " << spec.matmul->k << "u, " + << spec.matmul->batch + << "u}"; // batch 必须显式生成（漏则退化为默认 1，registry spec 失真） } o << "}, \n std::optional{"; if (spec.fold) @@ -159,10 +155,9 @@ namespace if (f.matmul) { o << "MatmulSpec{" << static_cast(f.matmul->a_input) << ", " - << static_cast(f.matmul->b_input) << ", " - << static_cast(f.matmul->transA) << ", " - << static_cast(f.matmul->transB) << ", " - << f.matmul->k << "u, " << f.matmul->batch << "u}"; + << static_cast(f.matmul->b_input) << ", " << static_cast(f.matmul->transA) + << ", " << static_cast(f.matmul->transB) << ", " << f.matmul->k << "u, " + << f.matmul->batch << "u}"; } o << "}, std::optional{"; if (f.vecacc) @@ -183,8 +178,7 @@ namespace for (std::size_t i = 0; i < spec.extras.size(); ++i) { if (i) o << ","; - o << " static_cast(" - << static_cast(spec.extras[i]) << ")"; + o << " static_cast(" << static_cast(spec.extras[i]) << ")"; } o << "} }"; return o.str(); @@ -200,7 +194,8 @@ namespace nn::tool [[nodiscard]] inline bool tool_env_flag(const char* name) { #if defined(_MSC_VER) - char* b = nullptr; std::size_t n = 0; + char* b = nullptr; + std::size_t n = 0; _dupenv_s(&b, &n, name); const bool v = (b != nullptr && b[0] != '\0' && b[0] != '0'); std::free(b); @@ -217,7 +212,8 @@ namespace nn::tool { int v = fallback; #if defined(_MSC_VER) - char* b = nullptr; std::size_t n = 0; + char* b = nullptr; + std::size_t n = 0; _dupenv_s(&b, &n, name); if (b != nullptr && b[0] != '\0') v = std::atoi(b); std::free(b); @@ -241,8 +237,7 @@ namespace nn::tool // 该结构/变体的每线程元素数（1=标量, 4=vec4）：vec4 资格 = 非 fold、非归约、 // 且逐元素形态可向量化；native16（ALU）变体恒标量（生成器同源判定）。 -[[nodiscard]] inline std::uint32_t fused_vec_width_for(const nn::ExprSpec& spec, - bool alu) +[[nodiscard]] inline std::uint32_t fused_vec_width_for(const nn::ExprSpec& spec, bool alu) { if (alu || spec.fold) return 1u; if (nn::expr_spec_reduce_axis(spec) >= 0) return 1u; @@ -315,16 +310,17 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) }; std::vector pool; pool.reserve(nthreads > 0 ? nthreads - 1 : 0); - for (unsigned t = 1; t < nthreads; ++t) pool.emplace_back(worker); + for (unsigned t = 1; t < nthreads; ++t) + pool.emplace_back(worker); worker(); // 调用线程也干活，省一个线程的启动/回收 - for (std::thread& th : pool) th.join(); + for (std::thread& th : pool) + th.join(); } // 磁盘上的文件是否与给定内容相同（不存在 = 不同）。 // **文本模式**（与旧的无条件覆写一致）：Windows 上写出 CRLF、读入时再翻译回 // LF，两边口径一致；改二进制模式会让生成头从 CRLF 变 LF，平白制造全文件 diff。 -[[nodiscard]] inline bool file_content_equal(const std::string& path, - const std::string& content) +[[nodiscard]] inline bool file_content_equal(const std::string& path, const std::string& content) { std::ifstream in(path); if (!in) return false; @@ -338,8 +334,7 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) // C++ TU 都 include 这份注册表）**不重建**。改一个 Layer 头会重新跑 // scan_exprs，但多数时候折叠出的结构逐字节不变——"无条件 ofstream 覆写" // 会把这种情况放大成全量重编。 -[[nodiscard]] inline bool write_if_different(const std::string& path, - const std::string& content) +[[nodiscard]] inline bool write_if_different(const std::string& path, const std::string& content) { if (file_content_equal(path, content)) return true; // 内容相同 → 不碰 mtime std::ofstream f(path); @@ -350,10 +345,9 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) // 生成融合 shader 注册表（见文件头说明）。reg 由调用方（scan_exprs）收集好后 // 传入。 -[[nodiscard]] inline bool generate_fused_registry( - const std::string& out_dir, - const std::string& glslc, - const nn::fused::ExprRegistry& reg) +[[nodiscard]] inline bool generate_fused_registry(const std::string& out_dir, + const std::string& glslc, + const nn::fused::ExprRegistry& reg) { std::error_code ec; std::filesystem::create_directories(out_dir, ec); @@ -371,7 +365,8 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) H << "// ═══════════════════════════════════════════════════════════════\n"; H << "// fused_registry.hpp — AUTO-GENERATED by tools/fused_generate.hpp，请勿手动编辑\n"; H << "// 生成者 = scan_exprs（收集 + 合成一体，见 tools/scan_exprs.cpp）\n"; - H << "// 每个条目：key(=expr_spec_key[#x|#s|#a|#b|#f|#ks|#kr]) → {ExprSpec 结构, 内联 SPIR-V}\n"; + H << "// 每个条目：key(=expr_spec_key[#x|#s|#a|#b|#f|#ks|#kr]) → {ExprSpec 结构, 内联 " + "SPIR-V}\n"; H << "// · key V0 全 f32\n"; H << "// · key#x V1 运行期精度分派（PC `prec`，覆盖任意签名）\n"; H << "// · key#s V1s storage-only 编译期特化（全输入 f16 + 输出 f16，" @@ -399,8 +394,10 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) H << " std::size_t spirv_words;\n"; H << " int reduce_axis; // -1=逐元素, 0=行归约, 1=列归约\n"; H << " int has_matmul; // 1=含前置 matmul 段（push constants 多 rows+mm_k）\n"; - H << " std::uint32_t view_param_count; // 运行时视图参数个数（RowMod/RotateHalf 的 push constant vp 槽）\n"; - H << " std::uint32_t rparam_count; // 运行时标量参数个数（优化器 lr/eps/β 等的 push constant rp 槽）\n"; + H << " std::uint32_t view_param_count; // 运行时视图参数个数（RowMod/RotateHalf 的 push " + "constant vp 槽）\n"; + H << " std::uint32_t rparam_count; // 运行时标量参数个数（优化器 lr/eps/β 等的 push " + "constant rp 槽）\n"; H << " std::uint32_t vec_width; // 每线程处理元素数（1=标量, 4=vec4）\n"; H << " std::uint32_t prec_sig; // 变体签名：V0=0；V1=EXPR_PREC_SIG_DISPATCH\n"; H << " // （运行期按真实签名填 PC `prec`）；\n"; @@ -419,8 +416,7 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) // 混合归约轴（行+列）无法单 kernel 融合 / 空指令表无段 = 结构损坏 // （上游变换丢掉 fold 段即此形态）。跳过而非让生成器对空表 back() // 越界（表现为 0xC00000FD 栈崩溃）。正常运行期到不了这里。 - std::fprintf(stderr, "[skip] 不可融合结构: %s\n", - nn::expr_spec_key(spec).c_str()); + std::fprintf(stderr, "[skip] 不可融合结构: %s\n", nn::expr_spec_key(spec).c_str()); continue; } // matmul+列归约（列方向，如 col_max(matmul)）：生成器按 @@ -440,18 +436,20 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) return false; } // 含归约（raxis >= 0）→ 归约 kernel（workgroup 级共享内存归约）；否则逐元素 - const std::string glsl = (raxis >= 0) - ? emitter->generate_reduce("fused_" + key, spec) - : emitter->generate("fused_" + key, spec); + const std::string glsl = (raxis >= 0) ? emitter->generate_reduce("fused_" + key, spec) + : emitter->generate("fused_" + key, spec); if (glsl.empty()) { std::fprintf(stderr, "[FAIL] 生成 %s 失败（归约结构不支持？）\n", key.c_str()); return false; } - jobs.push_back(GenJob{key, key, + jobs.push_back(GenJob{key, + key, out_dir + "/fused_" + key + ".comp", out_dir + "/fused_" + key + ".spv", - glsl, {}, false}); + glsl, + {}, + false}); } // ── 结构变体：V1 运行期精度分派（键 = key#x）与 V2 native16（键 = key#a）── @@ -476,18 +474,17 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) { const nn::ExprSpec* spec; nn::ExprPrecSig sig; // 变体签名：分派哨兵 / "全 f16" 位图 - std::string vkey; // "key#x" 或 "key#a"（注册表键） - std::string suffix; // "key_x" / "key_a"（标识符与文件名，不含 '#'） - bool alu = false; // native16 变体（vec_width 恒 1） + std::string vkey; // "key#x" 或 "key#a"（注册表键） + std::string suffix; // "key_x" / "key_a"（标识符与文件名，不含 '#'） + bool alu = false; // native16 变体（vec_width 恒 1） }; std::vector emitted_variants; // 去重：同一 vkey 只发射一次（重复会生成重复的 kSpirv_ 标识符 → 重定义错误） std::unordered_set emitted_vkeys; // 单个 (spec, sig) 的发射：alu = native16（V2）；否则分派（V1） - const auto emit_one = [&](const nn::ExprSpec& spec, nn::ExprPrecSig sig, - bool alu) -> bool { - if (!is_fusable(spec)) - return false; + const auto emit_one = [&](const nn::ExprSpec& spec, nn::ExprPrecSig sig, bool alu) -> bool + { + if (!is_fusable(spec)) return false; if (alu && !nn::expr_prec_sig_native16(spec, sig)) return false; // 谓词不通过 → 不发 ALU 变体（V1 仍覆盖） const std::string key = nn::expr_spec_key(spec); @@ -497,14 +494,12 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) // V2 native16（结构谓词判定）→ 键 key#a const bool dispv = nn::expr_prec_sig_is_dispatch(sig); const std::string vkey = dispv ? (key + nn::EXPR_PREC_DISPATCH_SUFFIX) - : alu ? (key + nn::EXPR_PREC_ALU_SUFFIX) + : alu ? (key + nn::EXPR_PREC_ALU_SUFFIX) : (key + nn::EXPR_PREC_STORAGE_SUFFIX); - const std::string suffix = dispv ? (key + "_x") : alu ? (key + "_a") - : (key + "_s"); - if (!emitted_vkeys.insert(vkey).second) - return false; // 已发射 → 幂等跳过 + const std::string suffix = dispv ? (key + "_x") : alu ? (key + "_a") : (key + "_s"); + if (!emitted_vkeys.insert(vkey).second) return false; // 已发射 → 幂等跳过 const std::string comp_path = out_dir + "/fused_" + suffix + ".comp"; - const std::string spv_path = out_dir + "/fused_" + suffix + ".spv"; + const std::string spv_path = out_dir + "/fused_" + suffix + ".spv"; auto emitter = nn::emitter_registry::make("glsl"); if (!emitter) { @@ -513,12 +508,11 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) } const int raxis = nn::expr_spec_reduce_axis(spec); const std::string glsl = (raxis >= 0) - ? emitter->generate_reduce("fused_" + suffix, spec, sig) - : emitter->generate("fused_" + suffix, spec, sig, alu); + ? emitter->generate_reduce("fused_" + suffix, spec, sig) + : emitter->generate("fused_" + suffix, spec, sig, alu); if (glsl.empty()) { - if (alu) - return false; // 结构谓词通过但发射器不支持该形态 → 静默跳过 + if (alu) return false; // 结构谓词通过但发射器不支持该形态 → 静默跳过 //（V0/V1 仍覆盖，只是拿不到原生 f16 算术） // 分派变体（key#x）是唯一的签名无关 shader：发不出来 = 该结构 // 没有 f16 路径 → 必须可见。**两类成因分开打**（否则"设计内跳过" @@ -528,10 +522,9 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) if (!spec.extras.empty()) std::fprintf(stderr, "[skip] 多输出(extras) 结构 %s 按设计不生成 #x/#s/#a 变体" - "（仅 V0；混精度走逐成员降级）\n", vkey.c_str()); - else - std::fprintf(stderr, "[skip] 分派变体 %s 的形态暂不支持生成\n", + "（仅 V0；混精度走逐成员降级）\n", vkey.c_str()); + else std::fprintf(stderr, "[skip] 分派变体 %s 的形态暂不支持生成\n", vkey.c_str()); return false; } // 只登记任务：glslc 统一在 run_gen_jobs 里并行执行。 @@ -542,10 +535,8 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) for (const auto& spec : reg.specs) { - if (!is_fusable(spec)) - continue; - if (!nn_no_dispatch && - !emit_one(spec, nn::EXPR_PREC_SIG_DISPATCH, /*alu=*/false)) + if (!is_fusable(spec)) continue; + if (!nn_no_dispatch && !emit_one(spec, nn::EXPR_PREC_SIG_DISPATCH, /*alu=*/false)) { // V1 发不出来只是少了 f16 路径（见 emit_one），V2/V0 仍可能成功； // 真正的失败（生成 GLSL 失败）已在 emit_one 内报 [FAIL]/[skip]。 @@ -554,13 +545,15 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) // （键 key#s）。它覆盖 `#a` 谓词排除的全部结构（fold / matmul / 归约 / // rparams / 微常量），因此无 shaderFloat16 的设备也能拿到去掉了 `prec` // 三元的 f16 存储路径。 - emit_one(spec, nn::expr_prec_sig_make(nn::EXPR_PREC_SIG_INPUT_MASK, - /*out_f16=*/true), + emit_one(spec, + nn::expr_prec_sig_make(nn::EXPR_PREC_SIG_INPUT_MASK, + /*out_f16=*/true), /*alu=*/false); // V2：结构性"全输入 f16 + 输出 f16"签名（输入位掩码取满， // 谓词内部再按 views 数量与其余条件判定） - emit_one(spec, nn::expr_prec_sig_make(nn::EXPR_PREC_SIG_INPUT_MASK, - /*out_f16=*/true), + emit_one(spec, + nn::expr_prec_sig_make(nn::EXPR_PREC_SIG_INPUT_MASK, + /*out_f16=*/true), /*alu=*/true); } @@ -586,8 +579,7 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) for (const auto& vr : reg.variants) { const auto it_spec = spec_by_key.find(vr.key); - if (it_spec == spec_by_key.end()) - continue; // 结构不在最终集合（理论上不会） + if (it_spec == spec_by_key.end()) continue; // 结构不在最终集合（理论上不会） const nn::ExprSpec& spec = *it_spec->second; if (nn::expr_prec_sig_all_f16(spec, vr.sig)) { @@ -604,34 +596,33 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) // 文件名/标识符后缀不能含 '#' → 换成 '_xxxx' std::string suffix = vkey; for (char& ch : suffix) - if (ch == '#') - ch = '_'; - if (!emitted_vkeys.insert(vkey).second) - continue; // 幂等（同 key+签名只发一次） + if (ch == '#') ch = '_'; + if (!emitted_vkeys.insert(vkey).second) continue; // 幂等（同 key+签名只发一次） const int raxis = nn::expr_spec_reduce_axis(spec); - const std::string glsl = (raxis >= 0) - ? emitter->generate_reduce("fused_" + suffix, spec, vr.sig) - : emitter->generate("fused_" + suffix, spec, vr.sig, /*alu=*/false); + const std::string glsl = + (raxis >= 0) ? emitter->generate_reduce("fused_" + suffix, spec, vr.sig) + : emitter->generate("fused_" + suffix, spec, vr.sig, /*alu=*/false); if (glsl.empty()) { ++n_var_fail; - std::fprintf(stderr, - "[skip] 特化变体 %s 的形态暂不支持生成（回退 `#x`）\n", + std::fprintf(stderr, "[skip] 特化变体 %s 的形态暂不支持生成（回退 `#x`）\n", vkey.c_str()); continue; } - jobs.push_back(GenJob{suffix, vkey, + jobs.push_back(GenJob{suffix, + vkey, out_dir + "/fused_" + suffix + ".comp", out_dir + "/fused_" + suffix + ".spv", - glsl, {}, false}); + glsl, + {}, + false}); emitted_variants.push_back(VariantEmit{&spec, vr.sig, vkey, suffix, /*alu=*/false}); } if (!reg.variants.empty()) std::printf("[gen] 扫描签名特化：观察 %zu 条，发射 %zu 条" "（跳过全 f16 %zu 条 = 与 #s 同 shader；形态不支持 %zu 条）\n", - reg.variants.size(), - reg.variants.size() - n_skipped_all_f16 - n_var_fail, + reg.variants.size(), reg.variants.size() - n_skipped_all_f16 - n_var_fail, n_skipped_all_f16, n_var_fail); } @@ -682,13 +673,11 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) for (const MmGridTarget& tgt : batch_targets) { const nn::ExprSpec& spec = *tgt.spec; - if (!is_fusable(spec) || !nn::expr_spec_batchable(spec)) - continue; + if (!is_fusable(spec) || !nn::expr_spec_batchable(spec)) continue; const std::string vkey = tgt.key + nn::EXPR_BATCH_SUFFIX; std::string suffix = vkey; for (char& ch : suffix) - if (ch == '#') - ch = '_'; + if (ch == '#') ch = '_'; if (!emitted_bkeys.insert(vkey).second) continue; // 同 vkey 幂等（重复会生成重复 kSpirv_ 标识符） auto emitter = nn::emitter_registry::make("glsl"); @@ -697,19 +686,20 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) std::fprintf(stderr, "[FAIL] 无法创建 GLSL emitter（IR-D 注册表异常）\n"); return false; } - const std::string glsl = - emitter->generate_batched("fused_" + suffix, spec, tgt.sig); + const std::string glsl = emitter->generate_batched("fused_" + suffix, spec, tgt.sig); if (glsl.empty()) { // 资格判据与生成器同一函数 → 正常不可达；保底可见（不静默）。 - std::fprintf(stderr, "[skip] 批量变体 %s 的形态暂不支持生成\n", - vkey.c_str()); + std::fprintf(stderr, "[skip] 批量变体 %s 的形态暂不支持生成\n", vkey.c_str()); continue; } - jobs.push_back(GenJob{suffix, vkey, + jobs.push_back(GenJob{suffix, + vkey, out_dir + "/fused_" + suffix + ".comp", out_dir + "/fused_" + suffix + ".spv", - glsl, {}, false}); + glsl, + {}, + false}); emitted_batched.push_back(BatchedEmit{&spec, vkey, suffix, tgt.sig}); } @@ -739,15 +729,13 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) for (const MmGridTarget& tgt : mm_grid_targets) { const nn::ExprSpec& spec = *tgt.spec; - if (!is_fusable(spec)) - continue; + if (!is_fusable(spec)) continue; if (!spec.matmul || spec.fold || nn::expr_spec_reduce_axis(spec) != -1) continue; // 非纯 matmul 形态：无细 tile 派发分支 → 不产出 const std::string vkey = tgt.key + nn::EXPR_MM_FINE_SUFFIX; std::string suffix = vkey; for (char& ch : suffix) - if (ch == '#') - ch = '_'; + if (ch == '#') ch = '_'; if (!emitted_fkeys.insert(vkey).second) continue; // 同 vkey 幂等（重复会生成重复 kSpirv_ 标识符） auto emitter = nn::emitter_registry::make("glsl"); @@ -756,19 +744,20 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) std::fprintf(stderr, "[FAIL] 无法创建 GLSL emitter（IR-D 注册表异常）\n"); return false; } - const std::string glsl = - emitter->generate_matmul_fine("fused_" + suffix, spec, tgt.sig); + const std::string glsl = emitter->generate_matmul_fine("fused_" + suffix, spec, tgt.sig); if (glsl.empty()) { // 资格判据与生成器同一分支 → 正常不可达；保底可见（不静默）。 - std::fprintf(stderr, "[skip] 细 tile 变体 %s 的形态暂不支持生成\n", - vkey.c_str()); + std::fprintf(stderr, "[skip] 细 tile 变体 %s 的形态暂不支持生成\n", vkey.c_str()); continue; } - jobs.push_back(GenJob{suffix, vkey, + jobs.push_back(GenJob{suffix, + vkey, out_dir + "/fused_" + suffix + ".comp", out_dir + "/fused_" + suffix + ".spv", - glsl, {}, false}); + glsl, + {}, + false}); emitted_fine.push_back(FineEmit{&spec, vkey, suffix, tgt.sig}); } @@ -798,8 +787,7 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) for (const MmGridTarget& tgt : mm_grid_targets) { const nn::ExprSpec& spec = *tgt.spec; - if (!is_fusable(spec)) - continue; + if (!is_fusable(spec)) continue; if (!spec.matmul || spec.fold || nn::expr_spec_reduce_axis(spec) != -1) continue; // 与细 tile 同判据：非纯 matmul 形态无 K 切分派发分支 const std::string kskey = tgt.key + nn::EXPR_MM_KSPLIT_SUFFIX; @@ -817,8 +805,7 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) if (glsl_ks.empty()) { // 资格判据与生成器同一函数 → 正常不可达；保底可见（不静默）。 - std::fprintf(stderr, "[skip] K 切分变体 %s 的形态暂不支持生成\n", - kskey.c_str()); + std::fprintf(stderr, "[skip] K 切分变体 %s 的形态暂不支持生成\n", kskey.c_str()); continue; } const std::string glsl_kr = @@ -826,26 +813,32 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) if (glsl_kr.empty()) { // 配对约束：#ks 单独存在时运行期永不派发 → 一并不产出（不留死条目）。 - std::fprintf(stderr, "[skip] K 切分伴生归约 %s 的形态暂不支持生成" - "（#ks/#kr 配对约束 → 双双不产出）\n", krkey.c_str()); + std::fprintf(stderr, + "[skip] K 切分伴生归约 %s 的形态暂不支持生成" + "（#ks/#kr 配对约束 → 双双不产出）\n", + krkey.c_str()); continue; } std::string ks_suffix = kskey; std::string kr_suffix = krkey; for (char& ch : ks_suffix) - if (ch == '#') - ch = '_'; + if (ch == '#') ch = '_'; for (char& ch : kr_suffix) - if (ch == '#') - ch = '_'; - jobs.push_back(GenJob{ks_suffix, kskey, + if (ch == '#') ch = '_'; + jobs.push_back(GenJob{ks_suffix, + kskey, out_dir + "/fused_" + ks_suffix + ".comp", out_dir + "/fused_" + ks_suffix + ".spv", - glsl_ks, {}, false}); - jobs.push_back(GenJob{kr_suffix, krkey, + glsl_ks, + {}, + false}); + jobs.push_back(GenJob{kr_suffix, + krkey, out_dir + "/fused_" + kr_suffix + ".comp", out_dir + "/fused_" + kr_suffix + ".spv", - glsl_kr, {}, false}); + glsl_kr, + {}, + false}); emitted_ks.push_back(KsEmit{&spec, kskey, ks_suffix, tgt.sig}); owned_krs.push_back(spec); owned_krs.back().views.push_back(nn::ExprView{}); // kind = Linear（占位槽） @@ -856,8 +849,7 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) run_gen_jobs(glslc, jobs); for (const GenJob& j : jobs) { - if (!j.ok) - return false; // [FAIL] 明细已在 run_gen_jobs 内逐条打印 + if (!j.ok) return false; // [FAIL] 明细已在 run_gen_jobs 内逐条打印 } // SPIR-V 数组按 jobs 顺序回填（= 串行版的「基础结构全部在前、变体全部在后」） @@ -875,30 +867,25 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) H << "inline const FusedShader kFusedShaders[] = {\n"; for (const auto& spec : reg.specs) { - if (!is_fusable(spec)) - continue; // 与上方生成循环同一判据（共享 is_fusable） + if (!is_fusable(spec)) continue; // 与上方生成循环同一判据（共享 is_fusable） const std::string key = nn::expr_spec_key(spec); const std::string spec_lit = emit_spec(spec); // 三行共用，只生成一次 H << " { \"" << key << "\",\n " << spec_lit << ",\n" - << " kSpirv_" << key - << ", sizeof(kSpirv_" << key << ")/sizeof(std::uint32_t), " + << " kSpirv_" << key << ", sizeof(kSpirv_" << key << ")/sizeof(std::uint32_t), " << nn::expr_spec_reduce_axis(spec) << ", " << (spec.matmul ? 1 : 0) << ", " << nn::expr_spec_runtime_view_param_count(spec) << ", " << nn::expr_spec_runtime_param_count(spec) << ", " - << fused_vec_width_for(spec, /*alu=*/false) - << ", 0u },\n"; + << fused_vec_width_for(spec, /*alu=*/false) << ", 0u },\n"; } // 变体行（键 = key#x 或 key#a；元数据与基础结构同源——精度不进结构） for (const auto& ve : emitted_variants) { const nn::ExprSpec& spec = *ve.spec; H << " { \"" << ve.vkey << "\",\n " << emit_spec(spec) << ",\n" - << " kSpirv_" << ve.suffix - << ", sizeof(kSpirv_" << ve.suffix << ")/sizeof(std::uint32_t), " - << nn::expr_spec_reduce_axis(spec) << ", " << (spec.matmul ? 1 : 0) << ", " - << nn::expr_spec_runtime_view_param_count(spec) << ", " - << nn::expr_spec_runtime_param_count(spec) << ", " - << fused_vec_width_for(spec, ve.alu) + << " kSpirv_" << ve.suffix << ", sizeof(kSpirv_" << ve.suffix + << ")/sizeof(std::uint32_t), " << nn::expr_spec_reduce_axis(spec) << ", " + << (spec.matmul ? 1 : 0) << ", " << nn::expr_spec_runtime_view_param_count(spec) << ", " + << nn::expr_spec_runtime_param_count(spec) << ", " << fused_vec_width_for(spec, ve.alu) << ", " << ve.sig << "u },\n"; } // P3 批量变体行（键 = key#b）：元数据 = 逐元素形态（reduce_axis=-1、 @@ -907,10 +894,9 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) { const nn::ExprSpec& spec = *be.spec; H << " { \"" << be.vkey << "\",\n " << emit_spec(spec) << ",\n" - << " kSpirv_" << be.suffix - << ", sizeof(kSpirv_" << be.suffix << ")/sizeof(std::uint32_t), " - << "-1, 0, " - << nn::expr_spec_runtime_view_param_count(spec) << ", " + << " kSpirv_" << be.suffix << ", sizeof(kSpirv_" << be.suffix + << ")/sizeof(std::uint32_t), " + << "-1, 0, " << nn::expr_spec_runtime_view_param_count(spec) << ", " << nn::expr_spec_runtime_param_count(spec) << ", " << "1u, " << be.sig << "u },\n"; } @@ -921,13 +907,11 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) { const nn::ExprSpec& spec = *fe.spec; H << " { \"" << fe.vkey << "\",\n " << emit_spec(spec) << ",\n" - << " kSpirv_" << fe.suffix - << ", sizeof(kSpirv_" << fe.suffix << ")/sizeof(std::uint32_t), " - << nn::expr_spec_reduce_axis(spec) << ", " << (spec.matmul ? 1 : 0) << ", " - << nn::expr_spec_runtime_view_param_count(spec) << ", " + << " kSpirv_" << fe.suffix << ", sizeof(kSpirv_" << fe.suffix + << ")/sizeof(std::uint32_t), " << nn::expr_spec_reduce_axis(spec) << ", " + << (spec.matmul ? 1 : 0) << ", " << nn::expr_spec_runtime_view_param_count(spec) << ", " << nn::expr_spec_runtime_param_count(spec) << ", " - << fused_vec_width_for(spec, /*alu=*/false) - << ", " << fe.sig << "u },\n"; + << fused_vec_width_for(spec, /*alu=*/false) << ", " << fe.sig << "u },\n"; } // K 切分变体行（键 = key#ks）：元数据与基础结构同源（reduce_axis / // has_matmul=1 / vp / rparam / vec_width / prec_sig=0）；**PC 固定头 7 槽** @@ -936,13 +920,11 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) { const nn::ExprSpec& spec = *ke.spec; H << " { \"" << ke.vkey << "\",\n " << emit_spec(spec) << ",\n" - << " kSpirv_" << ke.suffix - << ", sizeof(kSpirv_" << ke.suffix << ")/sizeof(std::uint32_t), " - << nn::expr_spec_reduce_axis(spec) << ", " << (spec.matmul ? 1 : 0) << ", " - << nn::expr_spec_runtime_view_param_count(spec) << ", " + << " kSpirv_" << ke.suffix << ", sizeof(kSpirv_" << ke.suffix + << ")/sizeof(std::uint32_t), " << nn::expr_spec_reduce_axis(spec) << ", " + << (spec.matmul ? 1 : 0) << ", " << nn::expr_spec_runtime_view_param_count(spec) << ", " << nn::expr_spec_runtime_param_count(spec) << ", " - << fused_vec_width_for(spec, /*alu=*/false) - << ", " << ke.sig << "u },\n"; + << fused_vec_width_for(spec, /*alu=*/false) << ", " << ke.sig << "u },\n"; } // K 切分伴生归约行（键 = key#kr）：spec 为合成 spec（+部分和视图）， // PC 固定头 6 槽与 matmul 同布局（has_matmul=1 → 注册侧 6，无需特判）。 @@ -950,13 +932,11 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) { const nn::ExprSpec& spec = *ke.spec; H << " { \"" << ke.vkey << "\",\n " << emit_spec(spec) << ",\n" - << " kSpirv_" << ke.suffix - << ", sizeof(kSpirv_" << ke.suffix << ")/sizeof(std::uint32_t), " - << nn::expr_spec_reduce_axis(spec) << ", " << (spec.matmul ? 1 : 0) << ", " - << nn::expr_spec_runtime_view_param_count(spec) << ", " + << " kSpirv_" << ke.suffix << ", sizeof(kSpirv_" << ke.suffix + << ")/sizeof(std::uint32_t), " << nn::expr_spec_reduce_axis(spec) << ", " + << (spec.matmul ? 1 : 0) << ", " << nn::expr_spec_runtime_view_param_count(spec) << ", " << nn::expr_spec_runtime_param_count(spec) << ", " - << fused_vec_width_for(spec, /*alu=*/false) - << ", " << ke.sig << "u },\n"; + << fused_vec_width_for(spec, /*alu=*/false) << ", " << ke.sig << "u },\n"; } H << "};\n"; H << "inline constexpr std::size_t kFusedShaderCount =\n" @@ -980,7 +960,8 @@ inline void run_gen_jobs(const std::string& glslc, std::vector& jobs) std::fprintf(stderr, "[FAIL] 无法写入 %s\n", reg_path.c_str()); return false; } - std::printf("[gen] %zu 条融合表达式 + %zu 条精度变体 + %zu 条批量变体 + %zu 条细 tile 变体 + %zu 条 K 切分变体(#ks+#kr) -> %s%s\n", + std::printf("[gen] %zu 条融合表达式 + %zu 条精度变体 + %zu 条批量变体 + %zu 条细 tile 变体 + " + "%zu 条 K 切分变体(#ks+#kr) -> %s%s\n", reg.specs.size(), emitted_variants.size(), emitted_batched.size(), emitted_fine.size(), emitted_ks.size(), reg_path.c_str(), reg_changed ? "" : "（内容未变化，保持 mtime）"); diff --git a/tools/scan_exprs.cpp b/tools/scan_exprs.cpp index 9419bb73..005d84d0 100644 --- a/tools/scan_exprs.cpp +++ b/tools/scan_exprs.cpp @@ -51,8 +51,8 @@ #include #if __has_include() - #include - #define NN_SCAN_HAS_STACKTRACE 1 +#include +#define NN_SCAN_HAS_STACKTRACE 1 #endif // abort（NN_ASSERT）时打印调用栈：NN_ASSERT 只有断言点行号，缺"谁调的"—— @@ -61,7 +61,8 @@ // 老编译器无 时降级为无栈消息（__has_include 守卫）。注： // handler 内 to_string(stacktrace) 会分配、严格说非异步信号安全——abort // 路径本已终止进程，这里只求尽力打印，不保证死锁免疫。 -namespace { +namespace +{ void on_abort(int) { #ifdef NN_SCAN_HAS_STACKTRACE @@ -105,8 +106,8 @@ int main(int argc, char* argv[]) " 精度分派 shader（键 key#x），运行时按真实精度填 PC `prec`。\n"); return 2; } - const std::string out_dir = argv[1]; - const std::string glslc = argv[2]; + const std::string out_dir = argv[1]; + const std::string glslc = argv[2]; std::signal(SIGABRT, &on_abort); // NN_ASSERT → abort 带栈（见文件头） auto raw_engine = nn::CpuEngine::create(); @@ -116,584 +117,552 @@ int main(int argc, char* argv[]) // ── dry-run 主体：跑一遍 Layer 的 forward/backward（只收集**结构**）── // engine 是**形参**（遮蔽外层 raw_engine）：f16 边界 cast 由基类 NVI 入口 // 统一处理（原 PrecisionEngine 适配层已删除，见 docs/history/history.md「15 号归档」 §4.1）。 - [[maybe_unused]] const auto dry_run = [&](nn::ComputeEngine& engine, const nn::PrecisionProfile& prof) + [[maybe_unused]] const auto dry_run = + [&](nn::ComputeEngine& engine, const nn::PrecisionProfile& prof) { scan_prof = prof; // M1（docs/internal/17 §4.1）：库外不再直构 Tensor——scan 输入经 // 引擎创建（出生绑定与本段 dry-run 使用的 engine 一致）。 const auto scan_tensor = [&](std::size_t rows, std::size_t cols) + { return engine.create_tensor(rows, cols, scan_prof.compute); }; + // ── RoPE：forward + backward ───────────────────────────────────────── + // （apply 与 apply_step 折叠出的结构相同，会自动去重） + // + // 形状无关融合：RowMod/RotateHalf 的周期/块大小是**运行时视图参数** + // （不进 expr_spec_key），故 RoPE 结构不依赖 d_k——这里任意 d_k 收集到的 + // 结构完全相同（去重成一个），**任何 d_k（含非 2 的幂）都能命中该融合 + // shader**。下面保留多个 d_k 仅作覆盖验证；若未来表达式重新引入形状 + // 常量，需再补对应 dry-run。 + for (const std::size_t dk : + {std::size_t{16}, std::size_t{32}, std::size_t{64}, std::size_t{128}}) { - return engine.create_tensor(rows, cols, scan_prof.compute); - }; - // ── RoPE：forward + backward ───────────────────────────────────────── - // （apply 与 apply_step 折叠出的结构相同，会自动去重） - // - // 形状无关融合：RowMod/RotateHalf 的周期/块大小是**运行时视图参数** - // （不进 expr_spec_key），故 RoPE 结构不依赖 d_k——这里任意 d_k 收集到的 - // 结构完全相同（去重成一个），**任何 d_k（含非 2 的幂）都能命中该融合 - // shader**。下面保留多个 d_k 仅作覆盖验证；若未来表达式重新引入形状 - // 常量，需再补对应 dry-run。 - 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(dk); - rope.set_precision_profile(scan_prof); // 扫描期精度（f16 pass 收集变体） - (void)rope.init(engine); // M6 段 C：层由 init 绑定引擎 - nn::Tensor q = scan_tensor(2 * dk, 8); // rows 为 dk 的整数倍 - (void)rope.apply(engine, q, /*seq=*/8, /*backward=*/false); - (void)rope.apply(engine, q, /*seq=*/8, /*backward=*/true); - nn::Tensor q1 = scan_tensor(dk, 1); // 增量推理（单位置） - (void)rope.apply_step(engine, q1, /*pos=*/3, /*backward=*/false); - (void)rope.apply_step(engine, q1, /*pos=*/3, /*backward=*/true); - } + nn::RotaryEmbedding rope(dk); + rope.set_precision_profile(scan_prof); // 扫描期精度（f16 pass 收集变体） + (void)rope.init(engine); // M6 段 C：层由 init 绑定引擎 + nn::Tensor q = scan_tensor(2 * dk, 8); // rows 为 dk 的整数倍 + (void)rope.apply(engine, q, /*seq=*/8, /*backward=*/false); + (void)rope.apply(engine, q, /*seq=*/8, /*backward=*/true); + nn::Tensor q1 = scan_tensor(dk, 1); // 增量推理（单位置） + (void)rope.apply_step(engine, q1, /*pos=*/3, /*backward=*/false); + (void)rope.apply_step(engine, q1, /*pos=*/3, /*backward=*/true); + } - // ── MSELoss forward（diff = pred-target；diff_sq = diff*diff）──────── - { - const std::size_t R = 8, C = 5; - nn::MSELoss mse; - mse.set_precision_profile(scan_prof); // 扫描期精度（f16 pass 收集变体） - nn::Tensor pred = scan_tensor(R, C); - nn::Tensor target = scan_tensor(R, C); - (void)mse.forward(engine, pred, target); - } + // ── MSELoss forward（diff = pred-target；diff_sq = diff*diff）──────── + { + const std::size_t R = 8, C = 5; + nn::MSELoss mse; + mse.set_precision_profile(scan_prof); // 扫描期精度（f16 pass 收集变体） + nn::Tensor pred = scan_tensor(R, C); + nn::Tensor target = scan_tensor(R, C); + (void)mse.forward(engine, pred, target); + } - // ── CrossEntropyLoss 稠密 forward（grad=softmax-target；target*log_sm； - // log_col_sum = log(col_sum)）───────────────────────────────────── - { - const std::size_t C = 8, B = 5; - nn::CrossEntropyLoss ce; - ce.set_precision_profile(scan_prof); // 扫描期精度（f16 pass 收集变体） - nn::Tensor logits = scan_tensor(C, B); - nn::Tensor target = scan_tensor(C, B); - (void)ce.forward(engine, logits, target); - } + // ── CrossEntropyLoss 稠密 forward（grad=softmax-target；target*log_sm； + // log_col_sum = log(col_sum)）───────────────────────────────────── + { + const std::size_t C = 8, B = 5; + nn::CrossEntropyLoss ce; + ce.set_precision_profile(scan_prof); // 扫描期精度（f16 pass 收集变体） + nn::Tensor logits = scan_tensor(C, B); + nn::Tensor target = scan_tensor(C, B); + (void)ce.forward(engine, logits, target); + } - // ── 优化器 step + 梯度裁剪（全部变体）──────────────────────────────── - // 各优化器均为 DSL 融合表达式 / 目标传递（dsl::compute_into）： - // sgd : p += -lr*g - // sgd_momentum : v = β*v + (1-β)*g ; p += -lr*v - // adam : m/v 更新 + p += delta - // adamw : p *= (1-lr*wd) + adam - // muon : v = μ*v + g；Nesterov；Newton–Schulz（含 A=bA+cA² 等） - // clip_grad_norm: g *= scale - // 必须逐个 dry-run：GPU 运行时同一结构才能命中 AOT 融合 shader（闭合世界）。 - // 结构不依赖形状，任取小 R×C 即可（Muon 的 NS 只对 ≥2D 参数生效）。 - { - const std::size_t R = 8, C = 5; - // RV-04：名单来自 nn::optimizer_table()（唯一真源）——新增优化器自动覆盖 - for (const auto &meta : nn::optimizer_table()) + // ── 优化器 step + 梯度裁剪（全部变体）──────────────────────────────── + // 各优化器均为 DSL 融合表达式 / 目标传递（dsl::compute_into）： + // sgd : p += -lr*g + // sgd_momentum : v = β*v + (1-β)*g ; p += -lr*v + // adam : m/v 更新 + p += delta + // adamw : p *= (1-lr*wd) + adam + // muon : v = μ*v + g；Nesterov；Newton–Schulz（含 A=bA+cA² 等） + // clip_grad_norm: g *= scale + // 必须逐个 dry-run：GPU 运行时同一结构才能命中 AOT 融合 shader（闭合世界）。 + // 结构不依赖形状，任取小 R×C 即可（Muon 的 NS 只对 ≥2D 参数生效）。 { - nn::Tensor p = scan_tensor(R, C); - nn::Tensor g = scan_tensor(R, C); - // 精度配置与运行期同源（`text_train` 把 `cfg.precision` 透传给 - // create_optimizer）：优化器的状态/参数槽精度决定 K3 更新量 - // （`upd_prec = p_.param`）与写目标的存储精度 → 直接决定该结构的 - // 运行期精度签名。漏传 = scan 只按默认 f32 profile 登记优化器， - // `--f16` 下 K3 的 `#10000`（F32 输入 + F16 输出）签名不被观察 - // ⇒ 无 `key#10000` 特化变体 ⇒ 逐成员派发回退 `#x` 且拿不到 - // `key#10000#b` 批量变体（P2-12 残差的结构性根因之一）。 - const nn::PrecisionProfile opt_prof = scan_prof; - auto opt = nn::create_optimizer(meta.name, engine, - std::vector{p}, - std::vector{g}, - nn::Scalar{1e-3f}, - /*weight_decay=*/nn::Scalar{0}, - /*precision=*/opt_prof); - if (opt) + const std::size_t R = 8, C = 5; + // RV-04：名单来自 nn::optimizer_table()（唯一真源）——新增优化器自动覆盖 + for (const auto& meta : nn::optimizer_table()) { - (void)opt->step(); - (void)opt->clip_grad_norm(nn::Scalar{1e3f}); + nn::Tensor p = scan_tensor(R, C); + nn::Tensor g = scan_tensor(R, C); + // 精度配置与运行期同源（`text_train` 把 `cfg.precision` 透传给 + // create_optimizer）：优化器的状态/参数槽精度决定 K3 更新量 + // （`upd_prec = p_.param`）与写目标的存储精度 → 直接决定该结构的 + // 运行期精度签名。漏传 = scan 只按默认 f32 profile 登记优化器， + // `--f16` 下 K3 的 `#10000`（F32 输入 + F16 输出）签名不被观察 + // ⇒ 无 `key#10000` 特化变体 ⇒ 逐成员派发回退 `#x` 且拿不到 + // `key#10000#b` 批量变体（P2-12 残差的结构性根因之一）。 + const nn::PrecisionProfile opt_prof = scan_prof; + auto opt = nn::create_optimizer(meta.name, engine, std::vector{p}, + std::vector{g}, nn::Scalar{1e-3f}, + /*weight_decay=*/nn::Scalar{0}, + /*precision=*/opt_prof); + if (opt) + { + (void)opt->step(); + (void)opt->clip_grad_norm(nn::Scalar{1e3f}); + } } } - } - // ── 算子融合二期（docs/internal/02-operator-fusion.md）：matmul 参与 IR 融合 ── - // 结构 = Layer 内 dsl::matmul(A,B)+bias+relu 折叠后的派生物： - // 前置 matmul 段（MatmulSpec）+ 尾逐元素链（Add + Max）。 - // 这里直接构造折叠后的结构并登记（与 dsl::compute 登记 spec 是同一机制）， - // 保证 GPU 运行时同一结构命中 AOT 融合 shader（闭合世界）。 - // - transA/transB 是**运行期 operand layout**（PC `mm_trans`）→ **不进 key**： - // 4 种转置组合折叠成同一个 key，登记一次即可（历史上此处登记 4 份， - // key 去掉转置后它们全被去重合并）。 - // - k（求和维度）是形状参数 → 不进 key：任取一个 K 登记，运行时任何 K - // 都命中同一融合 shader（同 RowMod/RotateHalf 的视图参数处理：不进 key， - // 运行时填充）。 - { - nn::ExprSpec s; - s.views = {nn::expr::linear(), nn::expr::linear(), nn::expr::linear()}; - s.num_regs = 2; - // transA/transB 取 (0,0) 作代表值：它们不进 key - s.matmul = nn::MatmulSpec{0, 1, /*transA=*/0, /*transB=*/0, /*k=*/8}; - // Add r0 = matmul(A,B) + bias（Input 2） - s.instrs.push_back({static_cast(nn::ExprOp::Add), 0, - nn::expr::matmul_op(), nn::expr::input(2), {}}); - // Max r1 = max(r0, 0)（relu） - s.consts.push_back(nn::Scalar{0}); - s.instrs.push_back({static_cast(nn::ExprOp::Max), 1, - nn::expr::reg(0), nn::expr::cst(0), {}}); - nn::fused::global_registry().add(s); - // 纯 matmul（无逐元素链）：输出 = matmul 结果 + // ── 算子融合二期（docs/internal/02-operator-fusion.md）：matmul 参与 IR 融合 ── + // 结构 = Layer 内 dsl::matmul(A,B)+bias+relu 折叠后的派生物： + // 前置 matmul 段（MatmulSpec）+ 尾逐元素链（Add + Max）。 + // 这里直接构造折叠后的结构并登记（与 dsl::compute 登记 spec 是同一机制）， + // 保证 GPU 运行时同一结构命中 AOT 融合 shader（闭合世界）。 + // - transA/transB 是**运行期 operand layout**（PC `mm_trans`）→ **不进 key**： + // 4 种转置组合折叠成同一个 key，登记一次即可（历史上此处登记 4 份， + // key 去掉转置后它们全被去重合并）。 + // - k（求和维度）是形状参数 → 不进 key：任取一个 K 登记，运行时任何 K + // 都命中同一融合 shader（同 RowMod/RotateHalf 的视图参数处理：不进 key， + // 运行时填充）。 { - nn::ExprSpec pure; - pure.views = {nn::expr::linear(), nn::expr::linear()}; - pure.num_regs = 0; - pure.matmul = nn::MatmulSpec{0, 1, 0, 0, /*k=*/8}; - nn::fused::global_registry().add(pure); + nn::ExprSpec s; + s.views = {nn::expr::linear(), nn::expr::linear(), nn::expr::linear()}; + s.num_regs = 2; + // transA/transB 取 (0,0) 作代表值：它们不进 key + s.matmul = nn::MatmulSpec{0, 1, /*transA=*/0, /*transB=*/0, /*k=*/8}; + // Add r0 = matmul(A,B) + bias（Input 2） + s.instrs.push_back({static_cast(nn::ExprOp::Add), + 0, + nn::expr::matmul_op(), + nn::expr::input(2), + {}}); + // Max r1 = max(r0, 0)（relu） + s.consts.push_back(nn::Scalar{0}); + s.instrs.push_back({static_cast(nn::ExprOp::Max), + 1, + nn::expr::reg(0), + nn::expr::cst(0), + {}}); + nn::fused::global_registry().add(s); + // 纯 matmul（无逐元素链）：输出 = matmul 结果 + { + nn::ExprSpec pure; + pure.views = {nn::expr::linear(), nn::expr::linear()}; + pure.num_regs = 0; + pure.matmul = nn::MatmulSpec{0, 1, 0, 0, /*k=*/8}; + nn::fused::global_registry().add(pure); + } } - } - // ── 算子融合二期（docs/internal/02-operator-fusion.md S5）：matmul+归约组合（注意力结构） - // bmm_reduce / bmm_denom 的 IR 等价物：matmul 段被归约指令消费， - // 不物化 (M,N) 得分矩阵（kernel 内联点积重算）。结构不依赖形状。 - // Q (M,K)，K 存储 (N,K)（transB=1）：QK^T = matmul(Q, K, transB) - { - // row_max(QK^T)：行归约 max（bmm_reduce ReduceOp::Max 等价） - nn::ExprSpec s; - s.views = {nn::expr::linear(), nn::expr::linear()}; - s.num_regs = 1; - s.matmul = nn::MatmulSpec{0, 1, 0, 1, /*k=*/8}; - s.instrs.push_back({static_cast(nn::ExprOp::RowMax), 0, - nn::expr::matmul_op(), {}, {}}); - nn::fused::global_registry().add(s); - // row_sum(matmul)：行归约 sum（bmm_reduce ReduceOp::Sum 等价） - nn::ExprSpec s2; - s2.views = {nn::expr::linear(), nn::expr::linear()}; - s2.num_regs = 1; - s2.matmul = nn::MatmulSpec{0, 1, 0, 1, /*k=*/8}; - s2.instrs.push_back({static_cast(nn::ExprOp::RowSum), 0, - nn::expr::matmul_op(), {}, {}}); - nn::fused::global_registry().add(s2); - // 列方向：col_max(matmul)（bmm_reduce reduce_cols=false 等价） - nn::ExprSpec s3; - s3.views = {nn::expr::linear(), nn::expr::linear()}; - s3.num_regs = 1; - s3.matmul = nn::MatmulSpec{0, 1, 0, 1, /*k=*/8}; - s3.instrs.push_back({static_cast(nn::ExprOp::ColMax), 0, - nn::expr::matmul_op(), {}, {}}); - nn::fused::global_registry().add(s3); - // denom：row_sum(exp(QK^T - rb(row_max)))（bmm_denom 等价；row_max - // 经 RowBroadcast 视图 (M,1)，不物化得分矩阵） - nn::ExprSpec sd; - sd.views = {nn::expr::linear(), nn::expr::linear(), nn::expr::row_broadcast()}; - sd.num_regs = 3; - sd.matmul = nn::MatmulSpec{0, 1, 0, 1, /*k=*/8}; - sd.instrs.push_back({static_cast(nn::ExprOp::Sub), 0, - nn::expr::matmul_op(), nn::expr::input(2), {}}); - sd.instrs.push_back({static_cast(nn::ExprOp::Exp), 1, - nn::expr::reg(0), {}, {}}); - sd.instrs.push_back({static_cast(nn::ExprOp::RowSum), 2, - nn::expr::reg(1), {}, {}}); - nn::fused::global_registry().add(sd); - } + // ── 算子融合二期（docs/internal/02-operator-fusion.md S5）：matmul+归约组合（注意力结构） + // bmm_reduce / bmm_denom 的 IR 等价物：matmul 段被归约指令消费， + // 不物化 (M,N) 得分矩阵（kernel 内联点积重算）。结构不依赖形状。 + // Q (M,K)，K 存储 (N,K)（transB=1）：QK^T = matmul(Q, K, transB) + { + // row_max(QK^T)：行归约 max（bmm_reduce ReduceOp::Max 等价） + nn::ExprSpec s; + s.views = {nn::expr::linear(), nn::expr::linear()}; + s.num_regs = 1; + s.matmul = nn::MatmulSpec{0, 1, 0, 1, /*k=*/8}; + s.instrs.push_back( + {static_cast(nn::ExprOp::RowMax), 0, nn::expr::matmul_op(), {}, {}}); + nn::fused::global_registry().add(s); + // row_sum(matmul)：行归约 sum（bmm_reduce ReduceOp::Sum 等价） + nn::ExprSpec s2; + s2.views = {nn::expr::linear(), nn::expr::linear()}; + s2.num_regs = 1; + s2.matmul = nn::MatmulSpec{0, 1, 0, 1, /*k=*/8}; + s2.instrs.push_back( + {static_cast(nn::ExprOp::RowSum), 0, nn::expr::matmul_op(), {}, {}}); + nn::fused::global_registry().add(s2); + // 列方向：col_max(matmul)（bmm_reduce reduce_cols=false 等价） + nn::ExprSpec s3; + s3.views = {nn::expr::linear(), nn::expr::linear()}; + s3.num_regs = 1; + s3.matmul = nn::MatmulSpec{0, 1, 0, 1, /*k=*/8}; + s3.instrs.push_back( + {static_cast(nn::ExprOp::ColMax), 0, nn::expr::matmul_op(), {}, {}}); + nn::fused::global_registry().add(s3); + // denom：row_sum(exp(QK^T - rb(row_max)))（bmm_denom 等价；row_max + // 经 RowBroadcast 视图 (M,1)，不物化得分矩阵） + nn::ExprSpec sd; + sd.views = {nn::expr::linear(), nn::expr::linear(), nn::expr::row_broadcast()}; + sd.num_regs = 3; + sd.matmul = nn::MatmulSpec{0, 1, 0, 1, /*k=*/8}; + sd.instrs.push_back({static_cast(nn::ExprOp::Sub), + 0, + nn::expr::matmul_op(), + nn::expr::input(2), + {}}); + sd.instrs.push_back( + {static_cast(nn::ExprOp::Exp), 1, nn::expr::reg(0), {}, {}}); + sd.instrs.push_back( + {static_cast(nn::ExprOp::RowSum), 2, nn::expr::reg(1), {}, {}}); + nn::fused::global_registry().add(sd); + } - // ── CausalSelfAttention（IR 掩码组合 forward/backward）────────── - // 掩码表达式 4 配置（causal / causal+alibi / causal+doc / causal+alibi+doc） - // 各产生确定结构：m/l/W（matmul 段 + Row/Col/Batch 操作数 + 视图）与 - // backward 的 R/X（纯逐元素）。必须全部 dry-run 覆盖（闭合世界）。 - { - const std::size_t d_model = 16, heads = 2, seq = 4, batch = 2; - const auto run_csa = [&](nn::CausalSelfAttention& attn) { + // ── CausalSelfAttention（IR 掩码组合 forward/backward）────────── + // 掩码表达式 4 配置（causal / causal+alibi / causal+doc / causal+alibi+doc） + // 各产生确定结构：m/l/W（matmul 段 + Row/Col/Batch 操作数 + 视图）与 + // backward 的 R/X（纯逐元素）。必须全部 dry-run 覆盖（闭合世界）。 + { + const std::size_t d_model = 16, heads = 2, seq = 4, batch = 2; + const auto run_csa = [&](nn::CausalSelfAttention& attn) + { + (void)attn.init(engine); + nn::Tensor x = scan_tensor(d_model, batch * seq); + // 禁止 (void) 吞错：forward 失败会让缓存为空，backward 直接 + // 在 batched_matmul 读空张量上 NN_ASSERT（栈无上下文难定位） + auto fr = attn.forward(x); + if (!fr) + { + std::fprintf(stderr, "[scan] CSA forward FAILED: %s\n", + fr.error().message.c_str()); + std::fflush(stderr); + std::abort(); // 带栈停在真凶处 + } + nn::Tensor grad = scan_tensor(d_model, batch * seq); + auto br = attn.backward(grad); + if (!br) + { + std::fprintf(stderr, "[scan] CSA backward FAILED: %s\n", + br.error().message.c_str()); + std::fflush(stderr); + std::abort(); + } + }; + for (const auto enc : {nn::PosEncodingType::Learned, nn::PosEncodingType::ALiBi}) + { + nn::CausalSelfAttention attn(d_model, heads, /*max_len=*/1024, + /*seq_len=*/seq, enc); + attn.set_precision_profile(scan_prof); // 扫描期精度（f16 pass 收集变体） + (void)attn.init(engine); // M6 段 C：层由 init 绑定引擎 + run_csa(attn); + } + // doc 变体（doc_ids 每位置文档 id，长度 = batch*seq） + { + const std::size_t doc_ids[8] = {0, 0, 1, 1, 0, 0, 1, 1}; + nn::CausalSelfAttention attn_d(d_model, heads, 1024, seq, + nn::PosEncodingType::Learned); + attn_d.set_precision_profile(scan_prof); // 扫描期精度（f16 pass 收集变体） + (void)attn_d.init(engine); // M6 段 C：层由 init 绑定引擎 + attn_d.set_doc_ids(doc_ids); + run_csa(attn_d); + } + { + const std::size_t doc_ids[8] = {0, 0, 1, 1, 0, 0, 1, 1}; + nn::CausalSelfAttention attn_ad(d_model, heads, 1024, seq, + nn::PosEncodingType::ALiBi); + attn_ad.set_precision_profile(scan_prof); // 扫描期精度（f16 pass 收集变体） + (void)attn_ad.init(engine); // M6 段 C：层由 init 绑定引擎 + attn_ad.set_doc_ids(doc_ids); + run_csa(attn_ad); + } + // GQA（docs/internal/22 §3.2 的 P2）：K/V_t 走 **HeadGroup 视图** + // （头分组比是运行期 vp 参数、**不进 key** → 登记一次覆盖任意 n_rep， + // 但"视图种类"本身是结构 → 必须是独立一条，否则 GPU 闭合世界对 + // GQA 的 fold 直接硬报错）。heads=2 / n_head_kv=1 → ratio 2。 + { + nn::CausalSelfAttention attn_gqa(d_model, heads, 1024, seq, + nn::PosEncodingType::RoPE, + /*subln=*/false, nn::NormType::LayerNorm, + /*n_head_kv=*/1); + attn_gqa.set_precision_profile(scan_prof); + (void)attn_gqa.init(engine); + run_csa(attn_gqa); + } + } + + // ── MultiHeadAttention（MHA=Plain 双向无掩码）─────────────── + // forward fold 的 Plain 变体 + **backward recompute 的裸 S 表达式** + // （无掩码 dsl::compute(matmul)，MHA dry-run 是该结构唯一的 + // 闭合世界注册来源）；R/X 与 grad 累加与 + // CSA 同构同 key。 + { + const std::size_t d_model = 16, heads = 2, seq = 4, batch = 2; + nn::MultiHeadAttention attn(d_model, heads, /*seq_len=*/seq); + attn.set_precision_profile(scan_prof); // 扫描期精度（f16 pass 收集变体） (void)attn.init(engine); nn::Tensor x = scan_tensor(d_model, batch * seq); - // 禁止 (void) 吞错：forward 失败会让缓存为空，backward 直接 - // 在 batched_matmul 读空张量上 NN_ASSERT（栈无上下文难定位） auto fr = attn.forward(x); if (!fr) { - std::fprintf(stderr, "[scan] CSA forward FAILED: %s\n", - fr.error().message.c_str()); + std::fprintf(stderr, "[scan] MHA forward FAILED: %s\n", fr.error().message.c_str()); std::fflush(stderr); - std::abort(); // 带栈停在真凶处 + std::abort(); } nn::Tensor grad = scan_tensor(d_model, batch * seq); auto br = attn.backward(grad); if (!br) { - std::fprintf(stderr, "[scan] CSA backward FAILED: %s\n", + std::fprintf(stderr, "[scan] MHA backward FAILED: %s\n", br.error().message.c_str()); std::fflush(stderr); std::abort(); } - }; - for (const auto enc : {nn::PosEncodingType::Learned, - nn::PosEncodingType::ALiBi}) - { - nn::CausalSelfAttention attn(d_model, heads, /*max_len=*/1024, - /*seq_len=*/seq, enc); - attn.set_precision_profile(scan_prof); // 扫描期精度（f16 pass 收集变体） - (void)attn.init(engine); // M6 段 C：层由 init 绑定引擎 - run_csa(attn); - } - // doc 变体（doc_ids 每位置文档 id，长度 = batch*seq） - { - const std::size_t doc_ids[8] = {0, 0, 1, 1, 0, 0, 1, 1}; - nn::CausalSelfAttention attn_d(d_model, heads, 1024, seq, - nn::PosEncodingType::Learned); - attn_d.set_precision_profile(scan_prof); // 扫描期精度（f16 pass 收集变体） - (void)attn_d.init(engine); // M6 段 C：层由 init 绑定引擎 - attn_d.set_doc_ids(doc_ids); - run_csa(attn_d); - } - { - const std::size_t doc_ids[8] = {0, 0, 1, 1, 0, 0, 1, 1}; - nn::CausalSelfAttention attn_ad(d_model, heads, 1024, seq, - nn::PosEncodingType::ALiBi); - attn_ad.set_precision_profile(scan_prof); // 扫描期精度（f16 pass 收集变体） - (void)attn_ad.init(engine); // M6 段 C：层由 init 绑定引擎 - attn_ad.set_doc_ids(doc_ids); - run_csa(attn_ad); } - // GQA（docs/internal/22 §3.2 的 P2）：K/V_t 走 **HeadGroup 视图** - // （头分组比是运行期 vp 参数、**不进 key** → 登记一次覆盖任意 n_rep， - // 但"视图种类"本身是结构 → 必须是独立一条，否则 GPU 闭合世界对 - // GQA 的 fold 直接硬报错）。heads=2 / n_head_kv=1 → ratio 2。 - { - nn::CausalSelfAttention attn_gqa(d_model, heads, 1024, seq, - nn::PosEncodingType::RoPE, - /*subln=*/false, nn::NormType::LayerNorm, - /*n_head_kv=*/1); - attn_gqa.set_precision_profile(scan_prof); - (void)attn_gqa.init(engine); - run_csa(attn_gqa); - } - } - // ── MultiHeadAttention（MHA=Plain 双向无掩码）─────────────── - // forward fold 的 Plain 变体 + **backward recompute 的裸 S 表达式** - // （无掩码 dsl::compute(matmul)，MHA dry-run 是该结构唯一的 - // 闭合世界注册来源）；R/X 与 grad 累加与 - // CSA 同构同 key。 - { - const std::size_t d_model = 16, heads = 2, seq = 4, batch = 2; - nn::MultiHeadAttention attn(d_model, heads, /*seq_len=*/seq); - attn.set_precision_profile(scan_prof); // 扫描期精度（f16 pass 收集变体） - (void)attn.init(engine); - nn::Tensor x = scan_tensor(d_model, batch * seq); - auto fr = attn.forward(x); - if (!fr) + // ── CrossAttention（issue #33：矩形 cross 注意力）────────────── + // 表达式结构与 MHA 同 key（Plain fold / 裸 S / R / X 均同构；seq_q/ + // seq_kv 是形状参数不进 key）——本 dry-run 的职责是**显式登记** + // seq_q ≠ seq_kv 的矩形路径（闭合世界按"注册过才存在"口径），并在 + // 构建期验证 forward_pair/backward_pair 在矩形形状下可跑通。 { - std::fprintf(stderr, "[scan] MHA forward FAILED: %s\n", - fr.error().message.c_str()); - std::fflush(stderr); - std::abort(); - } - nn::Tensor grad = scan_tensor(d_model, batch * seq); - auto br = attn.backward(grad); - if (!br) - { - std::fprintf(stderr, "[scan] MHA backward FAILED: %s\n", - br.error().message.c_str()); - std::fflush(stderr); - std::abort(); + const std::size_t d_model = 16, heads = 2, seq_q = 4, seq_kv = 6, batch = 2; + nn::CrossAttention attn(d_model, heads, /*seq_len=*/seq_q, + /*seq_kv_len=*/seq_kv); + attn.set_precision_profile(scan_prof); // 扫描期精度（f16 pass 收集变体） + (void)attn.init(engine); + nn::Tensor q = scan_tensor(d_model, batch * seq_q); + nn::Tensor mem = scan_tensor(d_model, batch * seq_kv); + auto fr = attn.forward_pair(q, mem); + if (!fr) + { + std::fprintf(stderr, "[scan] CrossAttention forward FAILED: %s\n", + fr.error().message.c_str()); + std::fflush(stderr); + std::abort(); + } + nn::Tensor grad = scan_tensor(d_model, batch * seq_q); + auto br = attn.backward_pair(grad); + if (!br) + { + std::fprintf(stderr, "[scan] CrossAttention backward FAILED: %s\n", + br.error().message.c_str()); + std::fflush(stderr); + std::abort(); + } } - } - // ── CrossAttention（issue #33：矩形 cross 注意力）────────────── - // 表达式结构与 MHA 同 key（Plain fold / 裸 S / R / X 均同构；seq_q/ - // seq_kv 是形状参数不进 key）——本 dry-run 的职责是**显式登记** - // seq_q ≠ seq_kv 的矩形路径（闭合世界按"注册过才存在"口径），并在 - // 构建期验证 forward_pair/backward_pair 在矩形形状下可跑通。 - { - const std::size_t d_model = 16, heads = 2, seq_q = 4, seq_kv = 6, - batch = 2; - nn::CrossAttention attn(d_model, heads, /*seq_len=*/seq_q, - /*seq_kv_len=*/seq_kv); - attn.set_precision_profile(scan_prof); // 扫描期精度（f16 pass 收集变体） - (void)attn.init(engine); - nn::Tensor q = scan_tensor(d_model, batch * seq_q); - nn::Tensor mem = scan_tensor(d_model, batch * seq_kv); - auto fr = attn.forward_pair(q, mem); - if (!fr) + // ── CrossEntropyLoss 稀疏 forward（IR 组合 denom/loss_vec/grad）── + // 结构不依赖形状；含 mask 与 label 越界修正的统一 mask 构造。 { - std::fprintf(stderr, "[scan] CrossAttention forward FAILED: %s\n", - fr.error().message.c_str()); - std::fflush(stderr); - std::abort(); + const std::size_t C = 8, B = 5; + nn::CrossEntropyLoss ce; + ce.set_precision_profile(scan_prof); // 扫描期精度（f16 pass 收集变体） + nn::Tensor logits = scan_tensor(C, B); + std::vector labels(B, 1); + std::vector mask(B, 1.0f); + (void)ce.forward_sparse(engine, logits, labels, mask, C); + // 无 mask（全 1 修正）与带越界 label 的变体 + std::vector labels2(B, 0); + labels2[0] = 99; // 越界 → mask 修正为 0 + (void)ce.forward_sparse(engine, logits, labels2, {}, C); } - nn::Tensor grad = scan_tensor(d_model, batch * seq_q); - auto br = attn.backward_pair(grad); - if (!br) + + // ── Softmax forward/backward（独立层；各模型主干走 fold 不经过它）── + // 输入槽去重（SpecBuilder::add_input）后 leaf(S)/row_reduce_max(S) 各由 + // 4 槽并为 2 槽 → 结构 key 变化；锚点符号实例（占位张量）**不参与去重** + // 只登记旧结构，故此类"锚点覆盖但 dry-run 未覆盖"的表达式必须在此以 + // 真实张量 dry-run 登记（否则运行期去重后 key 未命中 → 闭合世界硬报错）。 { - std::fprintf(stderr, "[scan] CrossAttention backward FAILED: %s\n", - br.error().message.c_str()); - std::fflush(stderr); - std::abort(); + const std::size_t R = 6, C = 9; // 与 fused_gpu_test run_softmax 同形 + nn::Softmax sm; + sm.set_precision_profile(scan_prof); + (void)sm.init(engine); + nn::Tensor x = scan_tensor(R, C); + (void)sm.forward(x); + nn::Tensor g = scan_tensor(R, C); + (void)sm.backward(g); } - } - - // ── CrossEntropyLoss 稀疏 forward（IR 组合 denom/loss_vec/grad）── - // 结构不依赖形状；含 mask 与 label 越界修正的统一 mask 构造。 - { - const std::size_t C = 8, B = 5; - nn::CrossEntropyLoss ce; - ce.set_precision_profile(scan_prof); // 扫描期精度（f16 pass 收集变体） - nn::Tensor logits = scan_tensor(C, B); - std::vector labels(B, 1); - std::vector mask(B, 1.0f); - (void)ce.forward_sparse(engine, logits, labels, mask, C); - // 无 mask（全 1 修正）与带越界 label 的变体 - std::vector labels2(B, 0); - labels2[0] = 99; // 越界 → mask 修正为 0 - (void)ce.forward_sparse(engine, logits, labels2, {}, C); - } - - // ── Softmax forward/backward（独立层；各模型主干走 fold 不经过它）── - // 输入槽去重（SpecBuilder::add_input）后 leaf(S)/row_reduce_max(S) 各由 - // 4 槽并为 2 槽 → 结构 key 变化；锚点符号实例（占位张量）**不参与去重** - // 只登记旧结构，故此类"锚点覆盖但 dry-run 未覆盖"的表达式必须在此以 - // 真实张量 dry-run 登记（否则运行期去重后 key 未命中 → 闭合世界硬报错）。 - { - const std::size_t R = 6, C = 9; // 与 fused_gpu_test run_softmax 同形 - nn::Softmax sm; - sm.set_precision_profile(scan_prof); - (void)sm.init(engine); - nn::Tensor x = scan_tensor(R, C); - (void)sm.forward(x); - nn::Tensor g = scan_tensor(R, C); - (void)sm.backward(g); - } - - // ── 归约表达式内联常量（push-constant 头长度回归）──────────────────── - // `col_reduce_sum(select(cond, 1, 0))` 是**带常量池的归约**结构：GPU 侧 - // push-constant 固定头长度必须按形态算（归约且无 matmul = 4 个 uint）。 - // ⚠ 固定头长度必须与生成器 PC 声明**逐形态一致**：错一个 uint 即常量池 - // 整体后移 → GPU 静默错值而 CPU 正常。fused_gpu_test 的 run_reduce_consts - // 做 CPU/GPU 对比覆盖。 - { - const std::size_t kk = 4, cols = 3; - nn::Tensor x = scan_tensor(kk, cols); - nn::Tensor mx = scan_tensor(1, cols); - (void)nn::dsl::compute_reduce(engine, - nn::dsl::col_reduce_sum(nn::dsl::select( - nn::dsl::leaf(x) == nn::dsl::col_broadcast(mx), - nn::Scalar{1}, nn::Scalar{0})), - kk, cols); - } - // ── BitLinear 的**两条量化缓冲路径**（F3 / P2）──────────────────────── - // 模型 pass 只产出"该引擎默认走的那条"（CPU 引擎 supports_ternary_gemm() - // = true → 打包路径），于是稠密路径的 `select → wq_` 与 matmul 尾链**不再 - // 被任何用例登记**；而未接打包的后端（GPU 未接时）运行期走稠密路径 → - // 闭合世界硬报错（t1_58_test_gpu 就是这么暴露的）。两条路径的结构都必须 - // 登记，故这里各 dry-run 一份（结构不依赖精度）。 - for (const bool packed : {true, false}) - { - nn::BitLinear bl(4, 3); - nn::PrecisionProfile t1 = scan_prof; - t1.param = nn::Precision::T1_58; - bl.set_precision_profile(t1); - (void)bl.init(engine); - bl.set_ternary_pack_enabled(packed); - nn::Tensor x = scan_tensor(4, 2); - auto fr = bl.forward(x); - if (!fr) + // ── 归约表达式内联常量（push-constant 头长度回归）──────────────────── + // `col_reduce_sum(select(cond, 1, 0))` 是**带常量池的归约**结构：GPU 侧 + // push-constant 固定头长度必须按形态算（归约且无 matmul = 4 个 uint）。 + // ⚠ 固定头长度必须与生成器 PC 声明**逐形态一致**：错一个 uint 即常量池 + // 整体后移 → GPU 静默错值而 CPU 正常。fused_gpu_test 的 run_reduce_consts + // 做 CPU/GPU 对比覆盖。 { - std::fprintf(stderr, "[scan] BitLinear(%s) forward FAILED: %s\n", - packed ? "packed" : "dense", fr.error().message.c_str()); - std::fflush(stderr); - std::abort(); + const std::size_t kk = 4, cols = 3; + nn::Tensor x = scan_tensor(kk, cols); + nn::Tensor mx = scan_tensor(1, cols); + (void)nn::dsl::compute_reduce( + engine, + nn::dsl::col_reduce_sum(nn::dsl::select( + nn::dsl::leaf(x) == nn::dsl::col_broadcast(mx), nn::Scalar{1}, nn::Scalar{0})), + kk, cols); } - nn::Tensor g = scan_tensor(3, 2); - auto br = bl.backward(g); - if (!br) + + // ── BitLinear 的**两条量化缓冲路径**（F3 / P2）──────────────────────── + // 模型 pass 只产出"该引擎默认走的那条"（CPU 引擎 supports_ternary_gemm() + // = true → 打包路径），于是稠密路径的 `select → wq_` 与 matmul 尾链**不再 + // 被任何用例登记**；而未接打包的后端（GPU 未接时）运行期走稠密路径 → + // 闭合世界硬报错（t1_58_test_gpu 就是这么暴露的）。两条路径的结构都必须 + // 登记，故这里各 dry-run 一份（结构不依赖精度）。 + for (const bool packed : {true, false}) { - std::fprintf(stderr, "[scan] BitLinear(%s) backward FAILED: %s\n", - packed ? "packed" : "dense", br.error().message.c_str()); - std::fflush(stderr); - std::abort(); + nn::BitLinear bl(4, 3); + nn::PrecisionProfile t1 = scan_prof; + t1.param = nn::Precision::T1_58; + bl.set_precision_profile(t1); + (void)bl.init(engine); + bl.set_ternary_pack_enabled(packed); + nn::Tensor x = scan_tensor(4, 2); + auto fr = bl.forward(x); + if (!fr) + { + std::fprintf(stderr, "[scan] BitLinear(%s) forward FAILED: %s\n", + packed ? "packed" : "dense", fr.error().message.c_str()); + std::fflush(stderr); + std::abort(); + } + nn::Tensor g = scan_tensor(3, 2); + auto br = bl.backward(g); + if (!br) + { + std::fprintf(stderr, "[scan] BitLinear(%s) backward FAILED: %s\n", + packed ? "packed" : "dense", br.error().message.c_str()); + std::fflush(stderr); + std::abort(); + } } - } - // ── RL：experience transforms / categorical policy / objectives ──────── - // RL 头不经过 Layer 模型 pass；这里显式触达每个 public evaluate 路径， - // 让 dsl::compute* 的 FusedAnchor 在构建期登记结构。scan 模式的算子是 - // 占位返回，具体数值分支不重要；所有表达式调用点都会被实例化。 - { - constexpr std::size_t R = 3, C = 2; - nn::Tensor logits = scan_tensor(R, C); - nn::Tensor action = scan_tensor(1, C); - nn::Tensor one = scan_tensor(1, C); - nn::Tensor value = scan_tensor(1, C); - nn::Tensor reward = scan_tensor(R, C); - nn::Tensor done = scan_tensor(R, C); - nn::Tensor mask_rc = scan_tensor(R, C); - nn::Tensor old_log_prob = scan_tensor(1, C); - - nn::CategoricalDistribution dist(scan_prof.stable); - (void)dist.log_prob(engine, logits, action); - (void)dist.entropy(engine, logits); - (void)dist.log_prob_with_gradient(engine, logits, action, one); - (void)dist.entropy_with_gradient(engine, logits, one); + // ── RL：experience transforms / categorical policy / objectives ──────── + // RL 头不经过 Layer 模型 pass；这里显式触达每个 public evaluate 路径， + // 让 dsl::compute* 的 FusedAnchor 在构建期登记结构。scan 模式的算子是 + // 占位返回，具体数值分支不重要；所有表达式调用点都会被实例化。 + { + constexpr std::size_t R = 3, C = 2; + nn::Tensor logits = scan_tensor(R, C); + nn::Tensor action = scan_tensor(1, C); + nn::Tensor one = scan_tensor(1, C); + nn::Tensor value = scan_tensor(1, C); + nn::Tensor reward = scan_tensor(R, C); + nn::Tensor done = scan_tensor(R, C); + nn::Tensor mask_rc = scan_tensor(R, C); + nn::Tensor old_log_prob = scan_tensor(1, C); - // PPO/value clip 的 select 结构：目标方法在 dry-run 中可能因全零 - // mask 提前返回；这里直接触达同一表达式类型，避免闭合世界漏项。 - nn::Tensor ratio = scan_tensor(1, C); - nn::Tensor old_value_scan = scan_tensor(1, C); - nn::Tensor value_scan = scan_tensor(1, C); - nn::Tensor returns_scan = scan_tensor(1, C); - const nn::Scalar lower = 0.8f; - const nn::Scalar upper = 1.2f; - auto clipped_ratio_expr = nn::dsl::select( - nn::dsl::leaf(ratio) < nn::dsl::rparam(lower), - nn::dsl::rparam(lower), - nn::dsl::select( - nn::dsl::leaf(ratio) > nn::dsl::rparam(upper), - nn::dsl::rparam(upper), nn::dsl::leaf(ratio))); - nn::Tensor clipped_ratio = NN_UNWRAP(nn::dsl::compute( - engine, clipped_ratio_expr, 1, C, scan_prof.stable)); - nn::Tensor clip_low = NN_UNWRAP(nn::dsl::compute( - engine, - nn::dsl::leaf(old_value_scan) - nn::dsl::rparam(0.2f), - 1, C, scan_prof.stable)); - nn::Tensor clip_high = NN_UNWRAP(nn::dsl::compute( - engine, - nn::dsl::leaf(old_value_scan) + nn::dsl::rparam(0.2f), - 1, C, scan_prof.stable)); - auto clipped_value_expr = nn::dsl::select( - nn::dsl::leaf(value_scan) < nn::dsl::leaf(clip_low), - nn::dsl::leaf(clip_low), - nn::dsl::select( - nn::dsl::leaf(value_scan) > nn::dsl::leaf(clip_high), - nn::dsl::leaf(clip_high), nn::dsl::leaf(value_scan))); - nn::Tensor clipped_value = NN_UNWRAP(nn::dsl::compute( - engine, clipped_value_expr, 1, C, scan_prof.stable)); - nn::Tensor clipped_diff = NN_UNWRAP(nn::dsl::compute( - engine, - nn::dsl::leaf(clipped_value) - nn::dsl::leaf(returns_scan), - 1, C, scan_prof.stable)); - nn::Tensor value_diff = NN_UNWRAP(nn::dsl::compute( - engine, - nn::dsl::leaf(value_scan) - nn::dsl::leaf(returns_scan), - 1, C, scan_prof.stable)); - nn::Tensor value_diff_sq = NN_UNWRAP(nn::dsl::compute( - engine, nn::dsl::leaf(value_diff) * nn::dsl::leaf(value_diff), - 1, C, scan_prof.stable)); - nn::Tensor value_clipped_sq = NN_UNWRAP(nn::dsl::compute( - engine, - nn::dsl::leaf(clipped_diff) * nn::dsl::leaf(clipped_diff), - 1, C, scan_prof.stable)); - auto value_max_expr = nn::dsl::max( - nn::dsl::leaf(value_diff_sq), - nn::dsl::leaf(value_clipped_sq)); - (void)nn::dsl::compute(engine, value_max_expr, 1, C, - scan_prof.stable); + nn::CategoricalDistribution dist(scan_prof.stable); + (void)dist.log_prob(engine, logits, action); + (void)dist.entropy(engine, logits); + (void)dist.log_prob_with_gradient(engine, logits, action, one); + (void)dist.entropy_with_gradient(engine, logits, one); - nn::Tensor ppo_objective = NN_UNWRAP(nn::dsl::compute( - engine, - nn::dsl::leaf(ratio) * nn::dsl::leaf(one), - 1, C, scan_prof.stable)); - nn::Tensor ppo_clipped_objective = NN_UNWRAP(nn::dsl::compute( - engine, - nn::dsl::leaf(clipped_ratio) * nn::dsl::leaf(one), - 1, C, scan_prof.stable)); - auto ppo_min_expr = nn::dsl::min( - nn::dsl::leaf(ppo_objective), - nn::dsl::leaf(ppo_clipped_objective)); - (void)nn::dsl::compute(engine, ppo_min_expr, 1, C, - scan_prof.stable); - const auto ratio_leaf = nn::dsl::leaf(ratio); - const auto adv_leaf = nn::dsl::leaf(one); - const auto positive_factor = nn::dsl::select(ratio_leaf > nn::dsl::rparam(upper), - nn::Scalar{0}, ratio_leaf * adv_leaf); - const auto negative_factor = nn::dsl::select(ratio_leaf < nn::dsl::rparam(lower), - nn::Scalar{0}, ratio_leaf * adv_leaf); - auto ppo_factor_expr = nn::dsl::select( - adv_leaf > nn::Scalar{0}, positive_factor, negative_factor); - nn::Tensor ppo_factor_scan = NN_UNWRAP(nn::dsl::compute( - engine, ppo_factor_expr, 1, C, scan_prof.stable)); - (void)nn::dsl::compute(engine, - nn::dsl::leaf(ppo_factor_scan) * - nn::dsl::leaf(one) * - nn::dsl::rparam(-1.0f), - 1, C, scan_prof.stable); - auto ppo_outside_expr = nn::dsl::select( - ratio_leaf < nn::dsl::rparam(lower), nn::Scalar{1}, - nn::dsl::select(ratio_leaf > nn::dsl::rparam(upper), - nn::Scalar{1}, nn::Scalar{0})) * - nn::dsl::leaf(one); - (void)nn::dsl::compute(engine, ppo_outside_expr, 1, C, - scan_prof.stable); - const auto value_outside = nn::dsl::select( - nn::dsl::leaf(value_scan) < nn::dsl::leaf(clip_low), - nn::Scalar{1}, - nn::dsl::select( - nn::dsl::leaf(value_scan) > nn::dsl::leaf(clip_high), - nn::Scalar{1}, nn::Scalar{0})); - auto value_outside_mask_expr = - value_outside * nn::dsl::leaf(one); - (void)nn::dsl::compute(engine, value_outside_mask_expr, 1, C, - scan_prof.stable); - nn::Tensor value_outside_tensor = NN_UNWRAP(nn::dsl::compute( - engine, value_outside, 1, C, scan_prof.stable)); - nn::Tensor value_clipped_grad = NN_UNWRAP(nn::dsl::compute( - engine, - nn::dsl::select( - nn::dsl::leaf(value_outside_tensor) > nn::Scalar{0}, - nn::Scalar{0}, - nn::dsl::leaf(value_diff) * nn::dsl::rparam(2.0f)), - 1, C, scan_prof.stable)); - nn::Tensor value_unclipped_grad = NN_UNWRAP(nn::dsl::compute( - engine, nn::dsl::leaf(value_diff) * nn::dsl::rparam(2.0f), - 1, C, scan_prof.stable)); - auto value_choose_clipped = - nn::dsl::leaf(value_clipped_sq) > - nn::dsl::leaf(value_diff_sq); - auto value_grad_expr = nn::dsl::select( - value_choose_clipped, nn::dsl::leaf(value_clipped_grad), - nn::dsl::leaf(value_unclipped_grad)) * - nn::dsl::leaf(one) * nn::dsl::rparam(1.0f); - (void)nn::dsl::compute(engine, value_grad_expr, 1, C, - scan_prof.stable); + // PPO/value clip 的 select 结构：目标方法在 dry-run 中可能因全零 + // mask 提前返回；这里直接触达同一表达式类型，避免闭合世界漏项。 + nn::Tensor ratio = scan_tensor(1, C); + nn::Tensor old_value_scan = scan_tensor(1, C); + nn::Tensor value_scan = scan_tensor(1, C); + nn::Tensor returns_scan = scan_tensor(1, C); + const nn::Scalar lower = 0.8f; + const nn::Scalar upper = 1.2f; + auto clipped_ratio_expr = nn::dsl::select( + nn::dsl::leaf(ratio) < nn::dsl::rparam(lower), nn::dsl::rparam(lower), + nn::dsl::select(nn::dsl::leaf(ratio) > nn::dsl::rparam(upper), + nn::dsl::rparam(upper), nn::dsl::leaf(ratio))); + nn::Tensor clipped_ratio = + NN_UNWRAP(nn::dsl::compute(engine, clipped_ratio_expr, 1, C, scan_prof.stable)); + nn::Tensor clip_low = NN_UNWRAP( + nn::dsl::compute(engine, nn::dsl::leaf(old_value_scan) - nn::dsl::rparam(0.2f), 1, + C, scan_prof.stable)); + nn::Tensor clip_high = NN_UNWRAP( + nn::dsl::compute(engine, nn::dsl::leaf(old_value_scan) + nn::dsl::rparam(0.2f), 1, + C, scan_prof.stable)); + auto clipped_value_expr = nn::dsl::select( + nn::dsl::leaf(value_scan) < nn::dsl::leaf(clip_low), nn::dsl::leaf(clip_low), + nn::dsl::select(nn::dsl::leaf(value_scan) > nn::dsl::leaf(clip_high), + nn::dsl::leaf(clip_high), nn::dsl::leaf(value_scan))); + nn::Tensor clipped_value = + NN_UNWRAP(nn::dsl::compute(engine, clipped_value_expr, 1, C, scan_prof.stable)); + nn::Tensor clipped_diff = NN_UNWRAP( + nn::dsl::compute(engine, nn::dsl::leaf(clipped_value) - nn::dsl::leaf(returns_scan), + 1, C, scan_prof.stable)); + nn::Tensor value_diff = NN_UNWRAP( + nn::dsl::compute(engine, nn::dsl::leaf(value_scan) - nn::dsl::leaf(returns_scan), 1, + C, scan_prof.stable)); + nn::Tensor value_diff_sq = NN_UNWRAP( + nn::dsl::compute(engine, nn::dsl::leaf(value_diff) * nn::dsl::leaf(value_diff), 1, + C, scan_prof.stable)); + nn::Tensor value_clipped_sq = NN_UNWRAP( + nn::dsl::compute(engine, nn::dsl::leaf(clipped_diff) * nn::dsl::leaf(clipped_diff), + 1, C, scan_prof.stable)); + auto value_max_expr = + nn::dsl::max(nn::dsl::leaf(value_diff_sq), nn::dsl::leaf(value_clipped_sq)); + (void)nn::dsl::compute(engine, value_max_expr, 1, C, scan_prof.stable); - nn::GaeAdvantage gae; - (void)gae.apply(engine, reward, reward, &done, &mask_rc, &one); - nn::MonteCarloReturn returns; - (void)returns.apply(engine, reward, &done, &mask_rc); - nn::MaskedNormalizer normalizer; - (void)normalizer.apply(engine, reward, mask_rc); - const nn::GroupSpan groups[1]{{0, R}}; - nn::GroupAdvantage group; - (void)group.apply(engine, reward, groups); + nn::Tensor ppo_objective = NN_UNWRAP(nn::dsl::compute( + engine, nn::dsl::leaf(ratio) * nn::dsl::leaf(one), 1, C, scan_prof.stable)); + nn::Tensor ppo_clipped_objective = NN_UNWRAP(nn::dsl::compute( + engine, nn::dsl::leaf(clipped_ratio) * nn::dsl::leaf(one), 1, C, scan_prof.stable)); + auto ppo_min_expr = + nn::dsl::min(nn::dsl::leaf(ppo_objective), nn::dsl::leaf(ppo_clipped_objective)); + (void)nn::dsl::compute(engine, ppo_min_expr, 1, C, scan_prof.stable); + const auto ratio_leaf = nn::dsl::leaf(ratio); + const auto adv_leaf = nn::dsl::leaf(one); + const auto positive_factor = nn::dsl::select(ratio_leaf > nn::dsl::rparam(upper), + nn::Scalar{0}, ratio_leaf * adv_leaf); + const auto negative_factor = nn::dsl::select(ratio_leaf < nn::dsl::rparam(lower), + nn::Scalar{0}, ratio_leaf * adv_leaf); + auto ppo_factor_expr = + nn::dsl::select(adv_leaf > nn::Scalar{0}, positive_factor, negative_factor); + nn::Tensor ppo_factor_scan = + NN_UNWRAP(nn::dsl::compute(engine, ppo_factor_expr, 1, C, scan_prof.stable)); + (void)nn::dsl::compute(engine, + nn::dsl::leaf(ppo_factor_scan) * nn::dsl::leaf(one) * + nn::dsl::rparam(-1.0f), + 1, C, scan_prof.stable); + auto ppo_outside_expr = + nn::dsl::select(ratio_leaf < nn::dsl::rparam(lower), nn::Scalar{1}, + nn::dsl::select(ratio_leaf > nn::dsl::rparam(upper), nn::Scalar{1}, + nn::Scalar{0})) * + nn::dsl::leaf(one); + (void)nn::dsl::compute(engine, ppo_outside_expr, 1, C, scan_prof.stable); + const auto value_outside = nn::dsl::select( + nn::dsl::leaf(value_scan) < nn::dsl::leaf(clip_low), nn::Scalar{1}, + nn::dsl::select(nn::dsl::leaf(value_scan) > nn::dsl::leaf(clip_high), nn::Scalar{1}, + nn::Scalar{0})); + auto value_outside_mask_expr = value_outside * nn::dsl::leaf(one); + (void)nn::dsl::compute(engine, value_outside_mask_expr, 1, C, scan_prof.stable); + nn::Tensor value_outside_tensor = + NN_UNWRAP(nn::dsl::compute(engine, value_outside, 1, C, scan_prof.stable)); + nn::Tensor value_clipped_grad = NN_UNWRAP(nn::dsl::compute( + engine, + nn::dsl::select(nn::dsl::leaf(value_outside_tensor) > nn::Scalar{0}, nn::Scalar{0}, + nn::dsl::leaf(value_diff) * nn::dsl::rparam(2.0f)), + 1, C, scan_prof.stable)); + nn::Tensor value_unclipped_grad = NN_UNWRAP(nn::dsl::compute( + engine, nn::dsl::leaf(value_diff) * nn::dsl::rparam(2.0f), 1, C, scan_prof.stable)); + auto value_choose_clipped = + nn::dsl::leaf(value_clipped_sq) > nn::dsl::leaf(value_diff_sq); + auto value_grad_expr = + nn::dsl::select(value_choose_clipped, nn::dsl::leaf(value_clipped_grad), + nn::dsl::leaf(value_unclipped_grad)) * + nn::dsl::leaf(one) * nn::dsl::rparam(1.0f); + (void)nn::dsl::compute(engine, value_grad_expr, 1, C, scan_prof.stable); - nn::ExperienceSchema exp = NN_UNWRAP(nn::ExperienceSchema::create( - {"action", "action_mask", "advantage", "return", "old_value", - "old_log_prob"})); - nn::TensorRecord record = NN_UNWRAP(nn::TensorRecord::create( - exp.schema(), - {nn::NamedTensor{.name = "action", .value = action}, - nn::NamedTensor{.name = "action_mask", .value = one}, - nn::NamedTensor{.name = "value_mask", .value = one}, - nn::NamedTensor{.name = "advantage", .value = one}, - nn::NamedTensor{.name = "return", .value = one}, - nn::NamedTensor{.name = "old_value", .value = one}, - nn::NamedTensor{.name = "old_log_prob", - .value = old_log_prob}})); - nn::PolicyBatchView view = NN_UNWRAP( - nn::PolicyBatchView::create(record)); + nn::GaeAdvantage gae; + (void)gae.apply(engine, reward, reward, &done, &mask_rc, &one); + nn::MonteCarloReturn returns; + (void)returns.apply(engine, reward, &done, &mask_rc); + nn::MaskedNormalizer normalizer; + (void)normalizer.apply(engine, reward, mask_rc); + const nn::GroupSpan groups[1]{{0, R}}; + nn::GroupAdvantage group; + (void)group.apply(engine, reward, groups); - nn::ReinforceObjective reinforce; - (void)reinforce.evaluate_context(engine, logits, view); - nn::ValueRegressionObjective value_reg; - (void)value_reg.evaluate_context(engine, value, view); - nn::EntropyBonusObjective entropy; - (void)entropy.evaluate_context(engine, logits, view); - nn::SampledKlObjective kl; - (void)kl.evaluate_context(engine, logits, view); - nn::ClippedPolicyObjective ppo_policy; - (void)ppo_policy.evaluate_context(engine, logits, view); - nn::ClippedValueObjective ppo_value; - (void)ppo_value.evaluate_context(engine, value, view); + nn::ExperienceSchema exp = NN_UNWRAP(nn::ExperienceSchema::create( + {"action", "action_mask", "advantage", "return", "old_value", "old_log_prob"})); + nn::TensorRecord record = NN_UNWRAP(nn::TensorRecord::create( + exp.schema(), {nn::NamedTensor{.name = "action", .value = action}, + nn::NamedTensor{.name = "action_mask", .value = one}, + nn::NamedTensor{.name = "value_mask", .value = one}, + nn::NamedTensor{.name = "advantage", .value = one}, + nn::NamedTensor{.name = "return", .value = one}, + nn::NamedTensor{.name = "old_value", .value = one}, + nn::NamedTensor{.name = "old_log_prob", .value = old_log_prob}})); + nn::PolicyBatchView view = NN_UNWRAP(nn::PolicyBatchView::create(record)); - nn::TensorSchema output = NN_UNWRAP(nn::TensorSchema::create( - "rl_output", - {nn::tensor_field("policy", - nn::TensorShapeSpec::exact(R, C)), - nn::tensor_field("value", - nn::TensorShapeSpec::exact(1, C))})); - nn::PpoObjective ppo; - (void)ppo.evaluate(engine, output, logits, value, view); - } + nn::ReinforceObjective reinforce; + (void)reinforce.evaluate_context(engine, logits, view); + nn::ValueRegressionObjective value_reg; + (void)value_reg.evaluate_context(engine, value, view); + nn::EntropyBonusObjective entropy; + (void)entropy.evaluate_context(engine, logits, view); + nn::SampledKlObjective kl; + (void)kl.evaluate_context(engine, logits, view); + nn::ClippedPolicyObjective ppo_policy; + (void)ppo_policy.evaluate_context(engine, logits, view); + nn::ClippedValueObjective ppo_value; + (void)ppo_value.evaluate_context(engine, value, view); + nn::TensorSchema output = NN_UNWRAP(nn::TensorSchema::create( + "rl_output", {nn::tensor_field("policy", nn::TensorShapeSpec::exact(R, C)), + nn::tensor_field("value", nn::TensorShapeSpec::exact(1, C))})); + nn::PpoObjective ppo; + (void)ppo.evaluate(engine, output, logits, value, view); + } }; // dry_run 结束 #if !defined(NN_SCAN_NO_DRYRUN) // 跑两遍：f32 定**结构**，f16 定**精度签名**。 @@ -717,8 +686,7 @@ int main(int argc, char* argv[]) // 本 pass 产出，故**失败即构建失败**：静默降级会产出不完整的注册表，只在 // GPU 运行期以闭合世界硬报错暴露（离根因很远）。与 dry-run 的 fail-fast // （见 run_csa 的 abort）一致。 - const auto model_pass = [&](nn::ComputeEngine& engine, - const nn::PrecisionProfile& prof) -> bool + const auto model_pass = [&](nn::ComputeEngine& engine, const nn::PrecisionProfile& prof) -> bool { scan_prof = prof; std::size_t n_fail = 0; @@ -756,8 +724,8 @@ int main(int argc, char* argv[]) auto eout = model.forward(input); if (!eout) { - std::fprintf(stderr, "[scan][FAIL] %s 推理态 forward 失败：%s\n", - name, eout.error().message.c_str()); + std::fprintf(stderr, "[scan][FAIL] %s 推理态 forward 失败：%s\n", name, + eout.error().message.c_str()); ++n_fail; } else @@ -766,8 +734,8 @@ int main(int argc, char* argv[]) auto ebr = model.backward(egrad); if (!ebr) { - std::fprintf(stderr, "[scan][FAIL] %s 推理态 backward 失败：%s\n", - name, ebr.error().message.c_str()); + std::fprintf(stderr, "[scan][FAIL] %s 推理态 backward 失败：%s\n", name, + ebr.error().message.c_str()); ++n_fail; } } @@ -803,16 +771,16 @@ int main(int argc, char* argv[]) // ── 1) MNIST MLP：归一化维（LayerNorm/RMSNorm/BatchNorm 是三条不同 // 融合表达式：归约链形状与原地更新都不同）── - for (const auto nrm : {nn::NormType::LayerNorm, nn::NormType::RMSNorm, - nn::NormType::BatchNorm}) + for (const auto nrm : + {nn::NormType::LayerNorm, nn::NormType::RMSNorm, nn::NormType::BatchNorm}) { const char* nm = (nrm == nn::NormType::LayerNorm) ? "mnist_mlp_ln" - : (nrm == nn::NormType::RMSNorm) ? "mnist_mlp_rms" + : (nrm == nn::NormType::RMSNorm) ? "mnist_mlp_rms" : "mnist_mlp_bn"; - try_build(nm, - [&] { return nn::build_mnist_mlp_model( - engine, nn::MNIST_LAYER_DIMS, nrm, prof); }, - scan_tensor(nn::MNIST_LAYER_DIMS.front(), B)); + try_build( + nm, + [&] { return nn::build_mnist_mlp_model(engine, nn::MNIST_LAYER_DIMS, nrm, prof); }, + scan_tensor(nn::MNIST_LAYER_DIMS.front(), B)); } // ── 1b) 三值（T1_58）MLP：BitLinear 的量化 / 去量化 / STE 三段结构 ── // BitLinear 的表达式（τ 的 absmean 归约、wq 的阈值 select、 @@ -824,42 +792,57 @@ int main(int argc, char* argv[]) { nn::PrecisionProfile t1 = prof; t1.param = nn::Precision::T1_58; - try_build("mnist_mlp_t1_58", - [&] { return nn::build_mnist_mlp_model( - engine, nn::MNIST_LAYER_DIMS, - nn::NormType::LayerNorm, t1); }, - scan_tensor(nn::MNIST_LAYER_DIMS.front(), B)); + try_build( + "mnist_mlp_t1_58", + [&] + { + return nn::build_mnist_mlp_model(engine, nn::MNIST_LAYER_DIMS, + nn::NormType::LayerNorm, t1); + }, + scan_tensor(nn::MNIST_LAYER_DIMS.front(), B)); } // ── 2) MNIST Transformer（ViT 风格；输入 img²）── - try_build("mnist_transformer", - [&] { return nn::build_mnist_transformer_model( - engine, nn::MNIST_IMG_SIZE, nn::MNIST_PATCH_SIZE, - nn::MNIST_TF_D_MODEL, nn::MNIST_TF_NUM_HEADS, - nn::MNIST_TF_D_FF, nn::MNIST_TF_NUM_LAYERS, prof); }, - scan_tensor(nn::MNIST_IMG_SIZE * nn::MNIST_IMG_SIZE, B)); + try_build( + "mnist_transformer", + [&] + { + return nn::build_mnist_transformer_model( + engine, nn::MNIST_IMG_SIZE, nn::MNIST_PATCH_SIZE, nn::MNIST_TF_D_MODEL, + nn::MNIST_TF_NUM_HEADS, nn::MNIST_TF_D_FF, nn::MNIST_TF_NUM_LAYERS, prof); + }, + scan_tensor(nn::MNIST_IMG_SIZE * nn::MNIST_IMG_SIZE, B)); // ── 2b) 归一化维（NormType / NormPlace 是运行期配置决定的结构，类型 // 层面推不出 → 与 GPT 的 pe/act/norm 同理，必须在这里补 dry-run）── - try_build("mnist_transformer_final_ln", - [&] { return nn::build_mnist_transformer_model( - engine, nn::MNIST_IMG_SIZE, nn::MNIST_PATCH_SIZE, - nn::MNIST_TF_D_MODEL, nn::MNIST_TF_NUM_HEADS, - nn::MNIST_TF_D_FF, nn::MNIST_TF_NUM_LAYERS, prof, - nn::NormType::LayerNorm, nn::NormPlace::Final); }, - scan_tensor(nn::MNIST_IMG_SIZE * nn::MNIST_IMG_SIZE, B)); - try_build("mnist_transformer_final_rms", - [&] { return nn::build_mnist_transformer_model( - engine, nn::MNIST_IMG_SIZE, nn::MNIST_PATCH_SIZE, - nn::MNIST_TF_D_MODEL, nn::MNIST_TF_NUM_HEADS, - nn::MNIST_TF_D_FF, nn::MNIST_TF_NUM_LAYERS, prof, - nn::NormType::RMSNorm, nn::NormPlace::Final); }, - scan_tensor(nn::MNIST_IMG_SIZE * nn::MNIST_IMG_SIZE, B)); - try_build("mnist_transformer_prenorm_bn", - [&] { return nn::build_mnist_transformer_model( - engine, nn::MNIST_IMG_SIZE, nn::MNIST_PATCH_SIZE, - nn::MNIST_TF_D_MODEL, nn::MNIST_TF_NUM_HEADS, - nn::MNIST_TF_D_FF, nn::MNIST_TF_NUM_LAYERS, prof, - nn::NormType::BatchNorm, nn::NormPlace::None); }, - scan_tensor(nn::MNIST_IMG_SIZE * nn::MNIST_IMG_SIZE, B)); + try_build( + "mnist_transformer_final_ln", + [&] + { + return nn::build_mnist_transformer_model( + engine, nn::MNIST_IMG_SIZE, nn::MNIST_PATCH_SIZE, nn::MNIST_TF_D_MODEL, + nn::MNIST_TF_NUM_HEADS, nn::MNIST_TF_D_FF, nn::MNIST_TF_NUM_LAYERS, prof, + nn::NormType::LayerNorm, nn::NormPlace::Final); + }, + scan_tensor(nn::MNIST_IMG_SIZE * nn::MNIST_IMG_SIZE, B)); + try_build( + "mnist_transformer_final_rms", + [&] + { + return nn::build_mnist_transformer_model( + engine, nn::MNIST_IMG_SIZE, nn::MNIST_PATCH_SIZE, nn::MNIST_TF_D_MODEL, + nn::MNIST_TF_NUM_HEADS, nn::MNIST_TF_D_FF, nn::MNIST_TF_NUM_LAYERS, prof, + nn::NormType::RMSNorm, nn::NormPlace::Final); + }, + scan_tensor(nn::MNIST_IMG_SIZE * nn::MNIST_IMG_SIZE, B)); + try_build( + "mnist_transformer_prenorm_bn", + [&] + { + return nn::build_mnist_transformer_model( + engine, nn::MNIST_IMG_SIZE, nn::MNIST_PATCH_SIZE, nn::MNIST_TF_D_MODEL, + nn::MNIST_TF_NUM_HEADS, nn::MNIST_TF_D_FF, nn::MNIST_TF_NUM_LAYERS, prof, + nn::NormType::BatchNorm, nn::NormPlace::None); + }, + scan_tensor(nn::MNIST_IMG_SIZE * nn::MNIST_IMG_SIZE, B)); // ── 3) CNN：池化窗口维（R 是运行期视图参数，pool=3 → R=9）── for (const std::size_t pool : {std::size_t{2}, std::size_t{3}}) { @@ -868,20 +851,22 @@ int main(int argc, char* argv[]) cfg.fc_dims = nn::MNIST_CNN_FC; cfg.pool = pool; // 注：CnnConfig 无精制度字段（CNN 走默认 f32 profile） - try_build(pool == 2 ? "cnn_pool2" : "cnn_pool3", - [&] { return nn::build_cnn_model(engine, cfg); }, - scan_tensor(1 * 28 * 28, B)); + try_build( + pool == 2 ? "cnn_pool2" : "cnn_pool3", + [&] { return nn::build_cnn_model(engine, cfg); }, scan_tensor(1 * 28 * 28, B)); } // ── 3b) CNN 归一化维（NormPlace × NormType 同上：运行期配置决定）── { - struct CnnNormCase { const char* name; nn::NormType nt; nn::NormPlace np; }; + struct CnnNormCase + { + const char* name; + nn::NormType nt; + nn::NormPlace np; + }; for (const CnnNormCase c : { - CnnNormCase{"cnn_norm_both_ln", nn::NormType::LayerNorm, - nn::NormPlace::Both}, - CnnNormCase{"cnn_norm_conv_rms", nn::NormType::RMSNorm, - nn::NormPlace::Conv}, - CnnNormCase{"cnn_norm_head_bn", nn::NormType::BatchNorm, - nn::NormPlace::Head}, + CnnNormCase{"cnn_norm_both_ln", nn::NormType::LayerNorm, nn::NormPlace::Both}, + CnnNormCase{"cnn_norm_conv_rms", nn::NormType::RMSNorm, nn::NormPlace::Conv}, + CnnNormCase{"cnn_norm_head_bn", nn::NormType::BatchNorm, nn::NormPlace::Head}, }) { nn::CnnConfig cfg; @@ -890,9 +875,9 @@ int main(int argc, char* argv[]) cfg.pool = 2; cfg.norm_type = c.nt; cfg.norm_place = c.np; - try_build(c.name, - [&] { return nn::build_cnn_model(engine, cfg); }, - scan_tensor(1 * 28 * 28, B)); + try_build( + c.name, [&] { return nn::build_cnn_model(engine, cfg); }, + scan_tensor(1 * 28 * 28, B)); } } // ── 4) GPT：位置编码 × 激活 × 归一化（每一维都换一组融合表达式： @@ -900,57 +885,62 @@ int main(int argc, char* argv[]) // SwiGLU ≠ GeLU、RMSNorm ≠ LayerNorm）── struct GptCase { - const char* name; - nn::PosEncodingType pe; - nn::ActivationType act; - nn::NormType norm; - std::size_t n_head_kv = 0; // 0 = MHA；>0 = GQA（头映射结构） - bool subln = false; - bool tie = false; - bool ternary = false; // P1.5：param=T1_58 → BitLinear + const char* name; + nn::PosEncodingType pe; + nn::ActivationType act; + nn::NormType norm; + std::size_t n_head_kv = 0; // 0 = MHA；>0 = GQA（头映射结构） + bool subln = false; + bool tie = false; + bool ternary = false; // P1.5：param=T1_58 → BitLinear }; - for (const GptCase& c : { - GptCase{"gpt_learned_gelu_ln", nn::PosEncodingType::Learned, - nn::ActivationType::GeLU, nn::NormType::LayerNorm}, - GptCase{"gpt_alibi_gelu_ln", nn::PosEncodingType::ALiBi, - nn::ActivationType::GeLU, nn::NormType::LayerNorm}, - GptCase{"gpt_sinusoidal_gelu_ln", nn::PosEncodingType::Sinusoidal, - nn::ActivationType::GeLU, nn::NormType::LayerNorm}, - GptCase{"gpt_rope_gelu_ln", nn::PosEncodingType::RoPE, - nn::ActivationType::GeLU, nn::NormType::LayerNorm}, - GptCase{"gpt_learned_swiglu_rms", nn::PosEncodingType::Learned, - nn::ActivationType::SwiGLU, nn::NormType::RMSNorm}, - // BitNet 2B4T 形态：门控平方 ReLU（relu2 ≠ swiglu ≠ gelu） - GptCase{"gpt_rope_relu2_rms", nn::PosEncodingType::RoPE, - nn::ActivationType::ReLU2, nn::NormType::RMSNorm}, - // 2B4T 完整组合：GQA（HeadGroup 视图）+ SubLN + tied head。 - // GQA 的 fold / masked_scores / grad_A / grad_Q 是**新结构** - // （K/V 走 HeadGroup 视图）→ 必须在此登记（闭合世界）。 - GptCase{"gpt_rope_relu2_rms_gqa_subln_tied", nn::PosEncodingType::RoPE, - nn::ActivationType::ReLU2, nn::NormType::RMSNorm, - // 模型 pass 的小配置 H=2 → n_head_kv=1 才得到 ratio>1 - // （比率是**视图参数**、不进 key → 任意 n_rep 共用一份 shader） - /*n_head_kv=*/1, /*subln=*/true, /*tie=*/true}, - // ── 三值 GPT（P1.5-W2）：权重侧接线后，注意力 4 个投影、 - // fc1/fc2、LM head 全部是 BitLinear。BitLinear 的三段结构 - // （τ 的 absmean 归约 / wq 的阈值 select / 去量化点积与 STE - // 累加）已由 1b) 的 MLP 用例登记；这里再按 GPT 组合登记一遍， - // 把"三值 + GQA/SubLN/tied + ReLU²/RoPE"这一整条链路钉进 - // 闭合世界（漏登记 = GPU 硬报错，见 AGENTS §7）。 - GptCase{"gpt_learned_gelu_ln_t1_58", nn::PosEncodingType::Learned, - nn::ActivationType::GeLU, nn::NormType::LayerNorm, - /*n_head_kv=*/0, /*subln=*/false, /*tie=*/false, - /*ternary=*/true}, - GptCase{"gpt_rope_relu2_rms_gqa_subln_tied_t1_58", - nn::PosEncodingType::RoPE, nn::ActivationType::ReLU2, - nn::NormType::RMSNorm, - /*n_head_kv=*/1, /*subln=*/true, /*tie=*/true, - /*ternary=*/true}}) + for (const GptCase& c : + {GptCase{"gpt_learned_gelu_ln", nn::PosEncodingType::Learned, nn::ActivationType::GeLU, + nn::NormType::LayerNorm}, + GptCase{"gpt_alibi_gelu_ln", nn::PosEncodingType::ALiBi, nn::ActivationType::GeLU, + nn::NormType::LayerNorm}, + GptCase{"gpt_sinusoidal_gelu_ln", nn::PosEncodingType::Sinusoidal, + nn::ActivationType::GeLU, nn::NormType::LayerNorm}, + GptCase{"gpt_rope_gelu_ln", nn::PosEncodingType::RoPE, nn::ActivationType::GeLU, + nn::NormType::LayerNorm}, + GptCase{"gpt_learned_swiglu_rms", nn::PosEncodingType::Learned, + nn::ActivationType::SwiGLU, nn::NormType::RMSNorm}, + // BitNet 2B4T 形态：门控平方 ReLU（relu2 ≠ swiglu ≠ gelu） + GptCase{"gpt_rope_relu2_rms", nn::PosEncodingType::RoPE, nn::ActivationType::ReLU2, + nn::NormType::RMSNorm}, + // 2B4T 完整组合：GQA（HeadGroup 视图）+ SubLN + tied head。 + // GQA 的 fold / masked_scores / grad_A / grad_Q 是**新结构** + // （K/V 走 HeadGroup 视图）→ 必须在此登记（闭合世界）。 + GptCase{"gpt_rope_relu2_rms_gqa_subln_tied", nn::PosEncodingType::RoPE, + nn::ActivationType::ReLU2, nn::NormType::RMSNorm, + // 模型 pass 的小配置 H=2 → n_head_kv=1 才得到 ratio>1 + // （比率是**视图参数**、不进 key → 任意 n_rep 共用一份 shader） + /*n_head_kv=*/1, /*subln=*/true, /*tie=*/true}, + // ── 三值 GPT（P1.5-W2）：权重侧接线后，注意力 4 个投影、 + // fc1/fc2、LM head 全部是 BitLinear。BitLinear 的三段结构 + // （τ 的 absmean 归约 / wq 的阈值 select / 去量化点积与 STE + // 累加）已由 1b) 的 MLP 用例登记；这里再按 GPT 组合登记一遍， + // 把"三值 + GQA/SubLN/tied + ReLU²/RoPE"这一整条链路钉进 + // 闭合世界（漏登记 = GPU 硬报错，见 AGENTS §7）。 + GptCase{"gpt_learned_gelu_ln_t1_58", nn::PosEncodingType::Learned, + nn::ActivationType::GeLU, nn::NormType::LayerNorm, + /*n_head_kv=*/0, /*subln=*/false, /*tie=*/false, + /*ternary=*/true}, + GptCase{"gpt_rope_relu2_rms_gqa_subln_tied_t1_58", nn::PosEncodingType::RoPE, + nn::ActivationType::ReLU2, nn::NormType::RMSNorm, + /*n_head_kv=*/1, /*subln=*/true, /*tie=*/true, + /*ternary=*/true}}) { nn::GptConfig cfg; - cfg.vocab_size = V; cfg.d_model = D; cfg.seq_len = S; - cfg.num_heads = H; cfg.d_ff = F; cfg.num_layers = L; - cfg.pos_enc = c.pe; cfg.activation = c.act; cfg.norm_type = c.norm; + cfg.vocab_size = V; + cfg.d_model = D; + cfg.seq_len = S; + cfg.num_heads = H; + cfg.d_ff = F; + cfg.num_layers = L; + cfg.pos_enc = c.pe; + cfg.activation = c.act; + cfg.norm_type = c.norm; cfg.n_head_kv = c.n_head_kv; cfg.subln = c.subln; cfg.tie_embeddings = c.tie; @@ -962,13 +952,18 @@ int main(int argc, char* argv[]) for (const bool causal : {true, false}) { nn::RAPTConfig cfg; - cfg.vocab_size = V; cfg.d_model = D; cfg.seq_len = S; - cfg.num_heads = H; cfg.d_ff = F; cfg.num_layers = L; + cfg.vocab_size = V; + cfg.d_model = D; + cfg.seq_len = S; + cfg.num_heads = H; + cfg.d_ff = F; + cfg.num_layers = L; cfg.pos_enc = nn::PosEncodingType::RoPE; // RAPT v1 仅支持 RoPE cfg.causal = causal; cfg.precision = prof; - try_build(causal ? "rapt_causal" : "rapt_bidir", - [&] { return nn::build_rapt_model(engine, cfg); }, ids); + try_build( + causal ? "rapt_causal" : "rapt_bidir", + [&] { return nn::build_rapt_model(engine, cfg); }, ids); } return n_fail == 0; }; @@ -978,9 +973,8 @@ int main(int argc, char* argv[]) #if !defined(NN_SCAN_NO_DRYRUN) if (!model_pass(*raw_engine, nn::profile_f32())) { - std::fprintf(stderr, - "[scan][FAIL] 模型 pass 失败 → 注册表不完整，中止构建" - "（闭合世界：缺结构只在 GPU 运行期报错）\n"); + std::fprintf(stderr, "[scan][FAIL] 模型 pass 失败 → 注册表不完整，中止构建" + "（闭合世界：缺结构只在 GPU 运行期报错）\n"); return 1; } // f16 模型 pass：同为**签名**来源（整模型路径的精度组合与 per-layer @@ -988,8 +982,7 @@ int main(int argc, char* argv[]) // 结构对结构是第一遍的重复（注册表按 key 去重），只有签名是新信息。 if (!model_pass(*raw_engine, nn::profile_f16())) { - std::fprintf(stderr, - "[scan][FAIL] f16 模型 pass 失败 → 特化签名不完整，中止构建\n"); + std::fprintf(stderr, "[scan][FAIL] f16 模型 pass 失败 → 特化签名不完整，中止构建\n"); return 1; } #endif @@ -1029,39 +1022,33 @@ int main(int argc, char* argv[]) for (const auto& i : s.instrs) { if (!ins.empty()) ins += ";"; - ins += std::to_string(static_cast(i.op)) + "(" - + std::to_string(static_cast(i.a.kind)) + "/" - + std::to_string(static_cast(i.a.idx)) + "," - + std::to_string(static_cast(i.b.kind)) + "/" - + std::to_string(static_cast(i.b.idx)) + "," - + std::to_string(static_cast(i.c.kind)) + "/" - + std::to_string(static_cast(i.c.idx)) + ")"; + ins += std::to_string(static_cast(i.op)) + "(" + + std::to_string(static_cast(i.a.kind)) + "/" + + std::to_string(static_cast(i.a.idx)) + "," + + std::to_string(static_cast(i.b.kind)) + "/" + + std::to_string(static_cast(i.b.idx)) + "," + + std::to_string(static_cast(i.c.kind)) + "/" + + std::to_string(static_cast(i.c.idx)) + ")"; } - return "raxis=" + std::to_string(raxis) - + " mm=" + (s.matmul ? "1" : "0") - + " fold=" + (s.fold ? "1" : "0") - + " nreg=" + std::to_string(s.num_regs) - + " nconst=" + std::to_string(s.consts.size()) - + " nrparam=" + std::to_string(s.rparams.size()) - + " nview=" + std::to_string(s.views.size()) - + " viewkinds=[" + v + "]" - + " instrs=[" + ins + "]"; + return "raxis=" + std::to_string(raxis) + " mm=" + (s.matmul ? "1" : "0") + + " fold=" + (s.fold ? "1" : "0") + " nreg=" + std::to_string(s.num_regs) + + " nconst=" + std::to_string(s.consts.size()) + + " nrparam=" + std::to_string(s.rparams.size()) + + " nview=" + std::to_string(s.views.size()) + " viewkinds=[" + v + "]" + + " instrs=[" + ins + "]"; }; std::size_t n_dry_only = 0; for (const auto& s : reg.specs) { const std::string k = nn::expr_spec_key(s); - if (anchor_keys.count(k)) - continue; + if (anchor_keys.count(k)) continue; ++n_dry_only; - if (dump_src) - std::printf("[scan][dry-only] %s %s\n", k.c_str(), summarize(s).c_str()); + if (dump_src) std::printf("[scan][dry-only] %s %s\n", k.c_str(), summarize(s).c_str()); } for (const auto& s : anchor.specs) { const std::string k = nn::expr_spec_key(s); - if (dry_keys.count(k)) - continue; + if (dry_keys.count(k)) continue; if (dump_src) std::printf("[scan][anchor-only] %s %s\n", k.c_str(), summarize(s).c_str()); } @@ -1071,8 +1058,8 @@ int main(int argc, char* argv[]) if (dump_src) for (const auto& s : reg.specs) if (anchor_keys.count(nn::expr_spec_key(s))) - std::printf("[scan][both] %s %s\n", - nn::expr_spec_key(s).c_str(), summarize(s).c_str()); + std::printf("[scan][both] %s %s\n", nn::expr_spec_key(s).c_str(), + summarize(s).c_str()); for (const auto& s : anchor.specs) reg.add(s); // 分段口径（**互不重叠**）：dry-run 侧 = ①per-layer dry-run + matmul/bmm @@ -1081,19 +1068,15 @@ int main(int argc, char* argv[]) std::printf("[scan] 结构来源：dry-run 侧 %zu（per-layer+显式 %zu、模型 pass +%zu）" " + 锚点 %zu（新增 %zu）= 合计 %zu" "（锚点独有 %zu、dry-run 独有 %zu）\n", - n_before_anchor, - n_after_dry, - n_after_model - n_after_dry, - anchor.specs.size(), - reg.specs.size() - n_before_anchor, - reg.specs.size(), + n_before_anchor, n_after_dry, n_after_model - n_after_dry, anchor.specs.size(), + reg.specs.size() - n_before_anchor, reg.specs.size(), reg.specs.size() - n_before_anchor, n_dry_only); // 锚点独有 = 手写 dry-run 覆盖不到的调用点（该数突然变大 → 有新层路径 // 没纳入 dry-run）。 if (anchor.specs.empty() || n_before_anchor == reg.specs.size()) std::fprintf(stderr, - "[scan][warn] 锚点未贡献任何结构：可能被 -DNN_SCAN_NO_ANCHOR 关闭，" - "或所有调用点都已被 dry-run 覆盖\n"); + "[scan][warn] 锚点未贡献任何结构：可能被 -DNN_SCAN_NO_ANCHOR 关闭，" + "或所有调用点都已被 dry-run 覆盖\n"); } std::printf("[scan] 收集到 %zu 条融合表达式\n", reg.specs.size()); // ── 签名覆盖：**构建期不再需要任何背书** ────────────────────────────── @@ -1108,7 +1091,6 @@ int main(int argc, char* argv[]) // NN_SCAN_NO_DISPATCH=1、或设备无 16bit 存储时不注册 → 落到边界 cast。 // ── 生成阶段（原 gen_fused）：同一进程内直接消费注册表 → fused_registry.hpp // 构建期因此只有一步；不再经 expr_specs.bin 中间序列化。 - if (!nn::tool::generate_fused_registry(out_dir, glslc, reg)) - return 1; + if (!nn::tool::generate_fused_registry(out_dir, glslc, reg)) return 1; return 0; } From c94f315144d8feb8d50543ed198bfd0038ba516e Mon Sep 17 00:00:00 2001 From: EthanPeng-2048 Date: Sat, 31 Oct 2026 02:58:59 +0800 Subject: [PATCH 4/8] =?UTF-8?q?chore:=20bump=20=E5=AD=90=E4=BB=93=E6=8C=87?= =?UTF-8?q?=E9=92=88=E8=87=B3=E5=90=84=E8=87=AA=E6=9C=80=E6=96=B0=20tag?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - extern/base → v1.0.2（1926464） - extern/dataset → v1.0.1（a9c3829） - extern/kvrecord → v1.1.1（276c666） - extern/threadpool → v1.0.1（17a0253） - extern/tokenizer → v1.0.1（4751259） 各 tag 均含统一 clang-format 格式化与嵌套子仓指针更新（递归同步，无 API/行为变化）。 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- extern/dataset | 2 +- extern/kvrecord | 2 +- extern/threadpool | 2 +- extern/tokenizer | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/extern/dataset b/extern/dataset index e58c91f1..a9c38290 160000 --- a/extern/dataset +++ b/extern/dataset @@ -1 +1 @@ -Subproject commit e58c91f1b648b6667eb765be9b8a808a6364b28f +Subproject commit a9c38290e40abd4d1b3862db1998d37fef163f86 diff --git a/extern/kvrecord b/extern/kvrecord index d043582f..276c6663 160000 --- a/extern/kvrecord +++ b/extern/kvrecord @@ -1 +1 @@ -Subproject commit d043582f0441ec09db3b2618aa3af53ddda45d14 +Subproject commit 276c6663ea70a036ac0443e1d63a7c52b13c047a diff --git a/extern/threadpool b/extern/threadpool index 657d266a..17a02530 160000 --- a/extern/threadpool +++ b/extern/threadpool @@ -1 +1 @@ -Subproject commit 657d266af0551e0e9504988bfbadfdb7f229707d +Subproject commit 17a02530a8989538d864144735f73fcd0fa26ccc diff --git a/extern/tokenizer b/extern/tokenizer index eefcc14d..47512594 160000 --- a/extern/tokenizer +++ b/extern/tokenizer @@ -1 +1 @@ -Subproject commit eefcc14d3650e6c0dba8521fec903d75fc58fa46 +Subproject commit 47512594f7427b582a7eb78ee4a8ec229fd1fcc1 From 1654d2198ee20786604a856d56cbd87b6428f7ec Mon Sep 17 00:00:00 2001 From: EthanPeng-2048 Date: Sat, 31 Oct 2026 10:57:31 +0800 Subject: [PATCH 5/8] =?UTF-8?q?docs/tool:=20=E6=94=B6=E5=8F=A3=E6=A0=BC?= =?UTF-8?q?=E5=BC=8F=E5=8C=96=E5=90=8E=E7=9A=84=E9=A3=8E=E6=A0=BC=E5=8F=A3?= =?UTF-8?q?=E5=BE=84=E4=B8=8E=20preset=20=E6=9E=84=E5=BB=BA=E7=9A=84?= =?UTF-8?q?=E4=BA=A7=E7=89=A9=E5=AE=9A=E4=BD=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - .clang-format / CONTRIBUTING / AGENTS / cpp-standards §8：把「整仓重排不在维护范围、 约 4 成行会变」改为「2026-10-31 已全仓统一格式化，现有代码即基线，改动走 git clang-format」， cpp-standards §8.2 的 K&R 示例改为与配置一致的 Allman（此前文档与工具互相矛盾） - history.md 增「全仓 clang-format 基线化」条，显式推翻 2026-10-11 的「整仓重排被否决」结论 - bench/gui_cli_audit.py、cli_controllers.py、gui.py：可执行文件定位再遍历一层 build//， 此前 CMakePresets 构建下门禁与 GUI 取不到产物（FileNotFoundError / 静默退回兜底名单） --- .clang-format | 15 ++++++--------- AGENTS.md | 2 +- CONTRIBUTING.md | 8 ++++---- bench/gui_cli_audit.py | 9 ++++++++- cli_controllers.py | 7 ++++++- docs/contributors/cpp-standards.md | 19 +++++++++++++------ docs/history/history.md | 15 +++++++++++++++ gui.py | 8 ++++++-- 8 files changed, 59 insertions(+), 24 deletions(-) diff --git a/.clang-format b/.clang-format index ce36e32d..267b328d 100644 --- a/.clang-format +++ b/.clang-format @@ -1,14 +1,11 @@ -# 本文件是「目标风格」的机器可读定义，与 docs/todo/10-development-standards.md §8 -# 「代码风格规范」一一对应。它是**增量**工具，不是整仓重排脚本： +# 本文件是代码风格的机器可读定义，与 docs/contributors/cpp-standards.md §8 +# 「代码风格规范」一一对应。**全仓已于 2026-10-31 按本配置统一格式化**（提交 5854b98， +# 5 个子仓同步），因此现有代码即基线，新增/改动代码的增量格式化不应再产生大 diff： # -# ✅ 推荐：对新文件、或你正在改的代码用 `git clang-format`（只格式化改动行） -# ⚠ 存量文件保留手工排版（对齐的中文注释块、按语义分组的续行），对**整个**文件 -# 跑 clang-format 会产生大 diff —— 实测**即便用本文件**，全仓 150 个源文件整体 -# 格式化仍会改约 **41%** 的行（`include/` 单独看约 44%）；主因是续行折行策略与 -# `if (c) stmt;` 单行/两行之争（两种写法仓内各约一半）。 -# 整仓重排**不在本仓的维护范围**；提交时请勿顺手格式化未改动的大段代码。 +# ✅ 只用 `git clang-format` 格式化**你这次改动**的行； +# ⚠ 不要在功能 PR 里顺手格式化无关文件（会让 review 与 git blame 失真）。 # -# 与仓内实际风格的取舍（实测依据，见 docs/history/history.md 的「代码风格」条）： +# 与仓内风格的取舍（实测依据，见 docs/history/history.md 的「代码风格」条）： # · 大括号：函数/类型/控制语句一律另起一行（Allman）——仓内主流写法； # · 顶层 namespace 内容不缩进（NamespaceIndentation: None）——仓内主流； # 少数历史文件（algebra_matrix.hpp / core_threadpool.hpp / cli/*.hpp）带缩进， diff --git a/AGENTS.md b/AGENTS.md index ba8962c6..2d1467af 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -69,7 +69,7 @@ cmake -B build -G Ninja -DNN_ENABLE_TESTS=ON && cmake --build build && ctest --t | 加/改 GUI 参数行（P2-6 声明式表） | `gui.py` 各 Tab 的 `fields()` —— **一个 `Field` = 一行 UI + 一条 `collect_args` 语义**（`Field`/`collect_fields`/`ParamsMixin` 在文件中部）；表表达不了的条件语义（engine→`--gpu`、精度预设→`--f16`/`--t1_58`、输入方式→`input`）留在该 Tab 的 `collect_args` 后处理里，并用 `_v()`/`ui_only=True` 声明给门禁 | | 参与/贡献入口（**给非作者**：构建/测试/规范/PR/CI 口径） | `CONTRIBUTING.md` + `CODE_OF_CONDUCT.md`；预设构建 `CMakePresets.json`（`cmake --preset dev/portable/cpu/probes`） | | 可引用的实测结论（结论/证据/重跑命令/边界） | `research/`（索引 `research/README.md`；一次研究一个子目录）。代码/设计稿引用这些结论时按四段式引用 | -| 代码风格（机器可读） | `.clang-format`（**目标风格**，用 `git clang-format` 只格式化自己的改动行；**不要**整仓重排）+ `.editorconfig` + `.gitattributes`（全仓 **LF**）；规则文本见 `docs/contributors/cpp-standards.md` §8 | +| 代码风格（机器可读） | `.clang-format`（全仓已于 2026-10-31 按本配置统一格式化，现有代码即基线；改动时用 `git clang-format` 只格式化自己的改动行）+ `.editorconfig` + `.gitattributes`（全仓 **LF**）；规则文本见 `docs/contributors/cpp-standards.md` §8 | ## 4. 分层架构（L0→L5，严格单向依赖，上层只依赖下层公有接口） diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 68fb93b4..9c9c6340 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -109,16 +109,16 @@ NN_VULKAN_DEVICE=850M ./build/test/bitnet_struct_test --gpu # 显式选卡 （命名 / 缩进 4 空格 / include 顺序 / 注释语言），机器可读版本是同目录的 [`.clang-format`](.clang-format)。 ```bash -# ✅ 推荐：只格式化**你这次改动**的行（历史文件是手工排版的，整文件格式化会产生大 diff） +# ✅ 推荐：只格式化**你这次改动**的行 git clang-format # 对暂存区相对 HEAD 的改动 git clang-format --diff # 只看会改什么 -# 单文件（新文件随意；旧文件请确认 diff 可接受） +# 单文件（新文件随意；旧文件整体跑通常已无 diff） clang-format -i path/to/new_file.hpp ``` -> `.clang-format` 是**目标风格**，不是整仓重排脚本：对现有文件整体跑它约 4 成行会变 -> （仓内 `if (c) stmt;` 单行与两行两种写法都大量存在）。**请不要在功能 PR 里顺手格式化无关文件。** +> 全仓已于 2026-10-31 按 `.clang-format` **统一格式化**（提交 `5854b98`，子仓同步）， +> 现有代码即基线；新增/改动代码用上面的增量方式保持一致即可。**请不要在功能 PR 里顺手格式化无关文件。** > 缩进/行尾/编码这类字符级约定由 [`.editorconfig`](.editorconfig) 与 [`.gitattributes`](.gitattributes) 统一（**全仓 LF**）。 ### 3.3 提交信息 diff --git a/bench/gui_cli_audit.py b/bench/gui_cli_audit.py index dcefeaba..ddb4aad3 100644 --- a/bench/gui_cli_audit.py +++ b/bench/gui_cli_audit.py @@ -71,8 +71,15 @@ def exe_path(exe: str) -> Path: 历史上这里写死 `build/.exe`，在 Linux 上直接 FileNotFoundError —— 本门禁因此只能在 Windows 跑。按"先 .exe 后无后缀"的顺序探测。 + `CMakePresets.json` 的构建把产物放在 `build//` 下，故再遍历一层 preset 目录 + （`--help` 文案与构建类型无关，取哪个 preset 都不影响本门禁的判定）。 """ - for cand in (ROOT / "build" / f"{exe}.exe", ROOT / "build" / exe): + names = (f"{exe}.exe", exe) + cands = [ROOT / "build" / n for n in names] + for sub in sorted((ROOT / "build").glob("*")): + if sub.is_dir(): + cands += [sub / n for n in names] + for cand in cands: if cand.exists(): return cand return ROOT / "build" / f"{exe}.exe" diff --git a/cli_controllers.py b/cli_controllers.py index 5d5bb50f..66aa01b2 100644 --- a/cli_controllers.py +++ b/cli_controllers.py @@ -160,7 +160,12 @@ def _get_executable_path(self) -> str: exe_path = build_dir / exe_name if exe_path.exists(): return str(exe_path) - + + # preset 构建（CMakePresets.json）把产物放在 build// 下 + for sub in sorted(build_dir.glob("*")): + if sub.is_dir() and (sub / exe_name).exists(): + return str(sub / exe_name) + # 尝试在当前目录查找 current_dir = Path(__file__).parent exe_path = current_dir / exe_name diff --git a/docs/contributors/cpp-standards.md b/docs/contributors/cpp-standards.md index 7e89aee9..2fc8281f 100644 --- a/docs/contributors/cpp-standards.md +++ b/docs/contributors/cpp-standards.md @@ -684,20 +684,27 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON) ### 8.2 格式规范 +> 机器可读真源是仓库根的 [`.clang-format`](../../.clang-format)（全仓已于 2026-10-31 +> 按该配置统一格式化，现有代码即基线）；下面只列最常用的几条。 + ```cpp -// 缩进：4 个空格 -if (condition) { - do_something(); +// 缩进：4 个空格；行宽：100 列 +// 大括号：Allman（函数 / 类型 / 控制语句一律另起一行） +void function() +{ + // ... } -// 大括号：K&R 风格 -void function() { - // ... +if (condition) +{ + do_something(); } // 指针/引用：靠近类型名 int* ptr; // 推荐 const std::vector& vec; // 推荐 + +// 注释不重排、行尾注释不对齐；include 顺序按 §8.3（工具不排序） ``` ### 8.3 包含顺序 diff --git a/docs/history/history.md b/docs/history/history.md index 86ffcc59..3b97aab6 100644 --- a/docs/history/history.md +++ b/docs/history/history.md @@ -2178,6 +2178,21 @@ GPU 后端 / GPU 引擎头文件「历史状态类注释」摘录。整改原则 - 验收：`.clang-format` 能被 clang-format 23 解析（`--style=file` 退出码 0）；`.gitattributes` 落地后 `git status` 无重规范化噪声（工作树本来就全 LF）。 +## 全仓 clang-format 基线化（2026-10-31） + +- 类型：演进记录（**推翻**上文 2026-10-11「整仓 clang-format 重排被否决」的结论） +- 内容：`5854b98` 对全仓 187 个 `*.cpp` / `*.hpp` / `*.inc` 跑 `clang-format -i`（配置即 + `.clang-format`，未改任何取值），净改动约 1.9 万行；随后 5 个子仓各自同配置格式化并打 + tag（base `v1.0.2`、dataset `v1.0.1`、kvrecord `v1.1.1`、threadpool `v1.0.1`、 + tokenizer `v1.0.1`），门面 `c94f315` bump 指针。 +- 为什么推翻旧结论：旧结论的前提是「存量手工排版是风格基线，整仓重排会制造 41% 的行级 + 噪声」；维护者决定把噪声**一次性付清**，此后新改动不再需要与两种写法共存，`git clang-format` + 的增量格式化也不再产生大 diff。 +- 代价（已知，不再重复申诉）：`git blame` 在该提交处断代，历史行归属需 `-w --follow` 或 + `git blame <5854b98^> -- path` 回看；`src/` 测试与 4 个 `compute_vk_backend_*.inc` 同批被重排。 +- 验收：Release 全量 ctest **73 注册 → 72 pass + 1 skip + 0 fail**（skip = `scan_bitwise_test`， + Windows 平台口径），证明格式化未改语义（宏续行、字符串拼接、4 个 `.inc` 表均编译通过）。 + --- # AGENTS.md diff --git a/gui.py b/gui.py index 89c5e0d4..9784a749 100644 --- a/gui.py +++ b/gui.py @@ -66,9 +66,13 @@ def cli_help_text(exe_name: str) -> str: """跑一次 CLI 的 --help 取帮助文本；未构建 / 超时 / 异常 → 空串。""" - exe = BUILD_DIR / (exe_name + (".exe" if os.name == "nt" else "")) + name = exe_name + (".exe" if os.name == "nt" else "") + exe = BUILD_DIR / name if not exe.exists(): - return "" + # preset 构建（CMakePresets.json）把产物放在 build// 下 + exe = next((p / name for p in sorted(BUILD_DIR.glob("*")) if (p / name).exists()), None) + if exe is None: + return "" try: proc = subprocess.run([str(exe), "--help"], capture_output=True, text=True, errors="replace", timeout=15) From a49311d337c92cd1ecf9b8f01b37d89bcf775752 Mon Sep 17 00:00:00 2001 From: EthanPeng-2048 Date: Sat, 31 Oct 2026 12:57:43 +0800 Subject: [PATCH 6/8] =?UTF-8?q?build/test:=20=E4=BF=AE=E5=A4=8D=E7=BA=AF?= =?UTF-8?q?=20CPU=20=E6=9E=84=E5=BB=BA=E7=9A=84=20-Werror=20=E5=81=87?= =?UTF-8?q?=E9=98=B3=E6=80=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - gpu_batch_scope_test 的 make_conv_model 只在 NN_HAS_VULKAN 内使用， 无 Vulkan 构建下触发 -Wunused-function -Werror（WSL clang/GPU=OFF 实测）， 用 #ifdef 收口。 - GCC 15.2 -O3 + PCH 对 std::optional/std::unique_ptr 内联析构报 -Wmaybe-uninitialized 假阳性（layer_lifetime_test 的 Model 析构路径， 诊断位置均为 ），仅对 GNU 关闭该子告警，-Wuninitialized 保留。 验证：WSL clang 24+libc++ 140/140 步、71/71 ctest 全过； WSL g++ 15.2+libstdc++ 71/71 ctest 全过。 --- CMakeLists.txt | 8 ++++++++ src/gpu_batch_scope_test.cpp | 2 ++ 2 files changed, 10 insertions(+) diff --git a/CMakeLists.txt b/CMakeLists.txt index 7aeb3da3..a7896999 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -111,6 +111,14 @@ if(MSVC) add_link_options(/STACK:8388608) else() add_compile_options(-fno-exceptions -Wall -Wextra -Wpedantic -Werror) + # GCC 15.2 在 -O3 + PCH 下对 std::optional / std::unique_ptr 的**内联析构** + # 报 -Wmaybe-uninitialized 假阳性（诊断里的内存位置全是 ；实测 + # src/layer_lifetime_test.cpp 的 `unique_ptr` 析构路径，成员已由 + # make_unique/默认初始化，无真实未初始化读取）。Clang 不产生该告警，因此 + # 只对 GNU 关闭此子告警；-Wuninitialized 保持开启。 + if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") + add_compile_options(-Wno-maybe-uninitialized) + endif() endif() # Windows 上的 clang++（GNU-like 驱动）不是 CMake 的 MSVC 分支，但链接仍走 diff --git a/src/gpu_batch_scope_test.cpp b/src/gpu_batch_scope_test.cpp index 1f9f667a..aca1a4c0 100644 --- a/src/gpu_batch_scope_test.cpp +++ b/src/gpu_batch_scope_test.cpp @@ -28,6 +28,7 @@ int g_fail = 0; } \ } while (0) +#ifdef NN_HAS_VULKAN namespace { @@ -42,6 +43,7 @@ nn::Result make_conv_model(nn::ComputeEngine& eng) } } // namespace +#endif // NN_HAS_VULKAN int main() { From d7d34c88a7e4b75a5449603bbabceddebf1eb12f Mon Sep 17 00:00:00 2001 From: EthanPeng-2048 Date: Sat, 31 Oct 2026 13:48:51 +0800 Subject: [PATCH 7/8] =?UTF-8?q?release:=20v1.14.0=20=E5=AE=9A=E6=A1=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - project(VERSION) → 1.14.0 - 新增 docs/release-notes/v1.14.0.md（主题/新增/迁移/验收基线/已知限制） - release-notes 索引加入 v1.14.0 - release-requirements 新增 §10 v1.14.0 发布实跑记录 （Release 72 pass+1 skip / Debug 73 pass / WSL clang+libc++ 与 g+++libstdc++ 全过） --- CMakeLists.txt | 2 +- docs/contributors/release-requirements.md | 63 +++++++ docs/release-notes/README.md | 1 + docs/release-notes/v1.14.0.md | 211 ++++++++++++++++++++++ 4 files changed, 276 insertions(+), 1 deletion(-) create mode 100644 docs/release-notes/v1.14.0.md diff --git a/CMakeLists.txt b/CMakeLists.txt index a7896999..803ce5e5 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -10,7 +10,7 @@ cmake_minimum_required(VERSION 3.30) # VERSION 以 git tag 为准（本仓无 CONFIGURE_FILE 消费该字段， # 它只是元数据——发版时同步到最新 tag，实际版本以 git tag / gh release 为准） -project(neuralnet.cpp VERSION 1.13.0 LANGUAGES CXX) +project(neuralnet.cpp VERSION 1.14.0 LANGUAGES CXX) # ══════════════════════════════════════════════════════════════════════════ # 作为库被消费（add_subdirectory / FetchContent 嵌入 3rd_party）的门控 diff --git a/docs/contributors/release-requirements.md b/docs/contributors/release-requirements.md index 6a187c90..3a56f1f7 100644 --- a/docs/contributors/release-requirements.md +++ b/docs/contributors/release-requirements.md @@ -357,3 +357,66 @@ f32 `.nnmodel` warm start 1 step：loss **7.6060**；文件 SHA-256 - 审查修复：PPO off-policy 负 advantage 梯度、GAE mask 值泄漏、policy batch 形状、 forward 失败后的具名反向状态、不可信 schema 解码上限、pipeline post-hook 状态提交； 全部加入对应 RL 用例。 + +--- + +## 10. 2026-10-31 `v1.14.0` 发布实跑记录 + +> 测试对象 = `main` 的发布定档提交。口径：GPU = **NVIDIA CMP 40HX / Vulkan SDK 1.4.350**； +> Windows 两树用本机 preset（`NN_ENABLE_NATIVE=ON`），WSL 两树 `NN_ENABLE_NATIVE=OFF` + `NN_ENABLE_GPU=OFF`； +> 构建一律 `-j32`（本机 32 核）。 + +### 10.1 完整 ctest（Release / Debug） + +| 树 | 构建 | 完整 ctest | +|---|---|---| +| Windows clang 23.1 + MSVC STL / **Release** / Vulkan | 160/160 步 | **73 注册 → 72 pass + 1 skip + 0 fail**（84.3 s，`ctest -j16`）；skip = `scan_bitwise_test` | +| Windows clang 23.1 + MSVC STL / **Debug** / Vulkan | 160/160 步 | **73 注册 → 73 pass + 0 skip + 0 fail**（434.6 s，顺序） | + +> Debug 必须顺序执行：`-j16` 并行时 `f16_precision_test` 被 GPU 争用拖到 40 min+（单跑 50.6 s 通过）； +> Release 并行无此现象。v1.13.0 记录的既有失败（`t1_58_test` / `write_epoch_test`）本版全部 Passed。 + +### 10.2 真实语料端到端训练（40HX） + +TinyStories 前 20,000 条 → **20,000 docs / 4,241,421 tokens / 8,208 词 / 66,273 个 seq64 窗口**； +GPT `d_model=64 / heads=4 / layers=2 / d_ff=128 / seq_len=64`，`--batch-size 8 --accum-steps 2`，20 step： + +| 配方 | loss（step 20） | 备注 | +|---|---|---| +| f32 | 8.7566 → **7.6110** | 1.2 s，无 NaN/冻结 | +| `--f16` | 8.7620 → **7.6741** | 0.4 s | +| `--t1_58` | 8.9024 → **8.1851** | 0.5 s | + +- 断点续训：`--save-interval 5` 的 step20 `.nnckpt` → `--resume` 续 3 步，loss + **7.6525 → 7.5733 → 7.5231**；日志确认优化器 76 张量、`t=20` 恢复； +- warm start：f32 `.nnmodel` 续 1 步 loss 7.6361，SHA-256 `19A26875…` → `77968371…`（参数确实移动）。 + +### 10.3 四组「编译器 × 标准库」 + +| 配置 | 结果 | +|---|---| +| Windows clang 23.1 + **MSVC STL** / Release / Vulkan | 见 §10.1（72 pass + 1 skip） | +| Windows clang 23.1 + **MSVC STL** / Debug / Vulkan | 见 §10.1（73 pass） | +| WSL g++ 15.2 + **libstdc++** / Release / CPU-only | 140/140 步；**71 注册 → 59 pass + 12 skip + 0 fail**（11.5 s） | +| WSL clang 24.0 + **libc++** / Release / CPU-only | 140/140 步；**71 注册 → 59 pass + 12 skip + 0 fail**（10.3 s） | + +> `cl.exe + MSVC STL` 一组本机无 MSVC 工具链（`cl.exe` 不在 PATH、VS 未安装），**未在本机复跑**， +> 沿用 CI 口径，记为已知覆盖缺口。WSL 的 12 个 GPU/offload 用例按 77 skip（无 Vulkan）。 +> +> 本轮修掉的两个**纯 CPU 构建**真实缺陷（`dev` 提交 `a49311d`）： +> ① 无 Vulkan 时 `gpu_batch_scope_test` 的 `make_conv_model` 触发 `-Wunused-function -Werror`； +> ② GCC 15.2 在 `-O3` + PCH 下对 `std::optional/std::unique_ptr` 内联析构报 +> `-Wmaybe-uninitialized` 假阳性（诊断位置均为 ``）→ 仅对 GNU 关闭该子告警。 + +### 10.4 门禁与独立发布审查 + +- `bench/doc_inventory.ps1`：**PASS**（`L2-VIOLATIONS: 0`、`HANDWRITTEN-RESIDUE: 0`、`TERMINATING-CALLS: 0`）； +- `bench/audit_dangling.ps1`：exit 0；`bench/gui_cli_audit.py`：**PASS，可行动 0** + （修复了 preset 构建产物定位：`build//`）； +- 独立只读审查（`v1.13.0..dev`）：**P0 无**；格式化提交按 token 流比对 182 个 C++ 文件 **0 语义差异**； + `8406b39` 的 `NN_TRY/NN_FAIL` 改写与宏展开逐字等价； +- **P1（2，记入 release note §5 已知限制）**： + ① 二次续训的 Adam `t` 退化（`make_train_state` 用 `session_.step_count()` 写 `optimizer_step`）； + ② `TrainingSession::resume` 先覆盖权重后 `verify_resume`（错误路径半提交）； +- **P2（4）**：步进驱动 `max_steps==0` 空跑成功返回、`save_checkpoint` 覆盖 `global_step` 的口径分歧、 + `InferenceScope::error()` 可解引用空 optional、`5854b98` 夹带 `extern/base` 指针且提交说明与实际不符。 diff --git a/docs/release-notes/README.md b/docs/release-notes/README.md index c26a3afc..699213ea 100644 --- a/docs/release-notes/README.md +++ b/docs/release-notes/README.md @@ -27,6 +27,7 @@ | 版本 | 主题 | |---|---| +| [v1.14.0](v1.14.0.md) | 训练骨架（`TrainState`/`BatchSource`/`RunDriver`）+ 开放规格（`ModelSpecBase`）+ `Embedding`/`CrossAttention` + 宿主桥公开 API + 全仓格式化基线 | | [v1.13.0](v1.13.0.md) | 子仓拆分（core / threadpool / kvrecord / dataset / tokenizer → `extern/` submodule） | | [v1.12.0](v1.12.0.md) | RL 与可组合训练基础设施（具名多输出 / experience / objectives / `TrainingGroupSession`） | | [v1.11.0](v1.11.0.md) | batch/step 计数语义统一 + 梯度累积归一化修复 | diff --git a/docs/release-notes/v1.14.0.md b/docs/release-notes/v1.14.0.md new file mode 100644 index 00000000..a41ba4dc --- /dev/null +++ b/docs/release-notes/v1.14.0.md @@ -0,0 +1,211 @@ +# neuralnet.cpp v1.14.0 Release Notes + +> **发布范围**：`v1.13.0`（tag `v1.13.0`，2026-10-26）→ `v1.14.0`（tag `v1.14.0`）。 +> **主题**：**训练骨架与开放规格**（RFC 32 阶段 A–D）+ **层与引擎接口补齐** +> （`Embedding`／`CrossAttention`／宿主桥公开 API／推理 batch 作用域）+ **全仓风格基线**。 +> **兼容性**：模型/断点文件格式仍为 **v8**，权重与词表段落读写不变；旧断点缺 `state_kind` +> 键一律回退 `TextTrainState`，旧文件数值 `type`（0–7）仍可读。**破坏性变更集中在 +> C++ API**（`TrainState` 虚基类化、`ModelSpec` 多态化、复合层默认广播），见 §3 迁移指南。 + +--- + +## 1. 主题与定位 + +v1.14.0 把「训练循环怎么写」从各领域各自实现收进库层的**领域无关骨架**，同时把 +**模型规格**开放给第三方自定义类型；并在层/引擎层补齐此前只能靠调用方拼装的三个能力 +（可学习查表、交叉注意力、宿主↔设备标量桥）。本版不改 CLI 语义、不改 `.nnmodel` +/ `.nnckpt` / `.nndataset` 的字节布局。 + +- 设计真源：[`../internal/32-trainer-skeleton-and-openspec.md`](../internal/32-trainer-skeleton-and-openspec.md)、 + [`../internal/30-training-session-step.md`](../internal/30-training-session-step.md)、 + [`../internal/29-session-architecture.md`](../internal/29-session-architecture.md)； +- 调用指南：[`../usage/02-quickstart-train-infer.md`](../usage/02-quickstart-train-infer.md)、 + [`../usage/03-compute-engine-usage.md`](../usage/03-compute-engine-usage.md)。 + +## 2. 新增能力 + +### 2.1 领域无关训练骨架（RFC 32 阶段 A–C） + +- **`TrainState` 虚基类化**：领域无关字段（计数器 / 配方 / 优化器指针）与编解码骨架 + （`train_state_to_kv` / `train_state_from_kv`）上收为基类；域钩子 + `state_kind()` / `clone()` / `write_extra_kv()` / `read_extra_kv()` / `verify_extra()`。 + 文本域派生 `TextTrainState`（`stride` + 数据集身份下沉）。 +- **`BatchSource` + `IndexedBatchSource`**（新头 `batch_source.hpp`）：数据供给契约与 + 预置索引适配器；shuffle 快照可重放（单流重放同序，`e248b08`）。 +- **`RunDriver`**（新头 `run_driver.hpp`）：epoch/步进两种驱动、lr 调度接线、 + `RunHooks{do_update, evaluate, on_epoch_end, on_checkpoint, on_phase}` 领域注入点、 + `RunReport` 统一统计。`global_step` 定义为**绝对步号**（跨续训累计）。 +- **`TextTrainer::fit()` 改写为 RunDriver 实例**（dogfood）：CLI 行为与日志口径不变， + 配方 ↔ 会话逐项一致性校验保留。 + +### 2.2 一站式断点编排（issue #37） + +`TrainingSession::save_checkpoint(path, spec, state, tokenizer_data)` 内部同步 +`global_step` / `lr` / 优化器 `t` 后以 Checkpoint 角色落盘； +`TrainingSession::resume(path, current, policy)` 一次完成载权重 → 读状态段 → +交叉校验（`ResumePolicy::Strict|Relaxed`）→ 恢复优化器与计数器，返回 +`unique_ptr`。调用方不再手工同步两侧计数器。 + +### 2.3 开放模型规格（issue #26，RFC 32 阶段 D） + +- `ModelSpecBase` 携带类型参数与四个分发虚方法 + （`arch_compatible` / `summary` / `write_kv` / `read_kv`）；`ModelSpec` 退化为公共字段外壳 + 加 `unique_ptr`（`clone()` 深拷贝）。 +- 五个内建子类：`MlpSpec` / `TransformerSpec` / `GptSpec` / `CnnSpec` / `RaptSpec`； + 原先四处 `switch(type)` 的 `spec_matches` / `spec_summary` / `spec_to_kv` / `spec_from_kv` + 全部收敛为虚方法。 +- **`ModelRegistry::register_type("MyNet", factory)`**：自定义类型按字符串类型名注册， + 文件 `type` 键写类型名；内建直通 `make_spec_by_name`。 + +### 2.4 新层 + +| 层 | 头文件 | 要点 | +|---|---|---| +| `Embedding` | `compute_layer_embedding.hpp` | 可学习查表；`padding_idx` 行梯度恒零（覆盖写而非累加）；梯度缓冲惰性分配（推理路径不付显存） | +| `CrossAttention` | `compute_layer_attention.hpp` | `AttentionBase` 家族成员，复用 `AttnScoreMask`（`Plain` / `Causal` / `CausalDoc`）；encoder-decoder 交叉注意力，Q 来自解码侧、K/V 来自编码侧 | + +两个层都有独立单测（`embedding_test` / `cross_attention_test`，梯度对拍 + 形状/掩码语义）。 + +### 2.5 引擎与宿主桥 + +- **宿主桥升公开 API**（issue #30）：`nn::upload_span` / `nn::download_span` / + `nn::download_vector`（f16 目标自动 cast，`fresh_dst` 语义收口）；`detail::` 包装保留。 +- **GPU 单次推理 batch 作用域**（issue #28）：`InferenceScope` RAII + `Model::forward` + 自动作用域；池底材在作用域结束时整块归还，OOM 报错指向根因（不再是模糊的分配失败）。 + +### 2.6 调度与采样扩展点 + +- **lr 调度**（issue #31）：新增 `multistep`（milestone × gamma，ImageNet 风格）与 + `step_cosine`（预热 + 分段余弦 + 段末重启，SGDR 风格）；**warmup 对所有 schedule 统一生效** + （退火进度按扣掉 warmup 的有效区间计算）。 +- **`SamplingConfig::logits_transform`**（issue #34）：采样前对 logits 的就地变换钩子 + （`std::function)>`），用于温度外置、禁词、约束解码等定制策略。 + +### 2.7 复合层默认广播（issue #29） + +`Layer` 基类的 `set_training` / `clear_cache` / `set_precision_profile` 默认实现从 no-op +改为按 `param_children()` **逐子层转发**；复合层只列子层即可，漏写转发不再可能让子层 +停在 f32 / eval 行为错误 / 显存残留。需特判者仍可 override（先调基类再覆盖）。 + +### 2.8 子仓与风格基线 + +- 5 个子仓统一 `v1..<漏洞修复>` 编号并各自出 tag： + base `v1.0.2`（sha256 栈缓冲溢出修复 + 格式化）、dataset `v1.0.1`、kvrecord `v1.1.1` + （格式规范 + Python 参考实现）、threadpool `v1.0.1`、tokenizer `v1.0.1`。 +- **全仓（含 5 子仓）按 `.clang-format` 统一格式化**（`5854b98`，约 1.9 万行）； + 此后新增/改动代码走 `git clang-format` 增量，不再有「两种写法共存」的噪声。 +- CI：checkout 拉取 submodule、`.gitmodules` 改 HTTPS（拆仓后 CI 双平台构建修复）。 + +## 3. 破坏性变更与迁移指南 + +### 3.1 `TrainState` 虚基类化（C++ API） + +- 状态段读回改为多态：`peek/load_train_state` 返回 `unique_ptr`，快照走 `clone()`。 +- 自定义训练状态需派生 `TrainState` 并实现 + `state_kind()`（落盘键值）/ `clone()` / `write_extra_kv()` / `read_extra_kv()` / `verify_extra()`。 +- 迁移：只在 `TextTrainer` / `TrainingSession` 上跑的调用方**无需改动**（`TextTrainState` + 是旧断点的权威类型，缺 `state_kind` 键自动回退）。 + +### 3.2 `ModelSpec` 多态化（C++ API + 文件 `type` 键含义） + +- 文件 `type` 键由数值枚举改为**字符串类型名**；`spec_matches` / `spec_summary` / + `spec_to_kv` / `spec_from_kv` 四个自由函数被 `ModelSpecBase` 虚方法取代。 +- 读侧保留向下兼容：旧文件数值 `type` 0–7 经 `legacy_spec_type_name` 映射 + （4 = ALiBi_GPT 归 `GPT` 同族；6 = ZiPT 仍明确报错）。 +- 迁移：直接调上述四个自由函数的代码改调虚方法（`spec.arch->arch_compatible(...)` 等）； + 自定义架构按 §2.3 注册。 + +### 3.3 复合层默认广播（行为变更） + +- 以前「复合层不 override 就不转发」的隐式 no-op 行为**不再存在**：只列 `param_children()` + 即自动广播三开关。若某个复合层**故意**要求子层保持默认（例如三值分流的特判）， + 必须显式 override 并在 override 里只对目标子层生效。 +- 该变更会让此前被静默忽略的配置错误**显式生效**（这是预期行为，不是回归）。 + +### 3.4 工具与文档路径 + +- `bench/doc_align_audit.ps1` 与 `.template/` 图标原件已删除（`62affb8`）； +- 文档目录重构为 `docs/{introduction,development,usage,contributors,internal,history,release-notes,benchmarks}`， + release note 从仓库根 `release-notes/` 迁入 `docs/release-notes/`——**旧外链失效**； +- `bench/gui_cli_audit.py`、`cli_controllers.py`、`gui.py` 的产物定位改为兼容 + `build//`（此前 `cmake --preset` 构建下门禁与 GUI 取不到可执行文件）。 + +### 3.5 测试目标 + +- 主仓移除 7 个已下沉到子仓的测试目标（对应能力由子仓各自测试覆盖）。 + +## 4. 验收基线 + +> 口径：`NN_ENABLE_NATIVE=OFF` 可移植基线不可用时的本机 `dev` preset 用 `NN_ENABLE_NATIVE=ON`； +> GPU = **NVIDIA CMP 40HX / Vulkan SDK 1.4.350**；构建目录全部从空目录开始（`build/dev`、`build/debug`）。 + +### 4.1 完整 ctest（Release / Debug） + +| 树 | 构建 | 完整 ctest | +|---|---|---| +| Windows clang 23.1 + MSVC STL / **Release** / Vulkan | `cmake --build --preset dev -- -j32`（160 步） | **73 注册 → 72 pass + 1 skip + 0 fail**（84.3 s，`ctest -j16`）；skip = `scan_bitwise_test`（Windows 平台口径） | +| Windows clang 23.1 + MSVC STL / **Debug** / Vulkan | `cmake --build build/debug -- -j32`（160 步） | **73 注册 → 73 pass + 0 skip + 0 fail**（434.6 s，顺序执行）；`rapt_test_f16` / `write_epoch_test` / `t1_58_test`（含 `_gpu`）全部 Passed | + +> Debug 树按**顺序**跑：`-j16` 并行会让 `f16_precision_test`（Debug 单跑 50.6 s）在多测试争抢 GPU 时 +> 拖到 40 分钟以上；Release 树并行无此问题（`-j16` 仅 84 s）。 + +> v1.13.0 记录的既有失败（`t1_58_test` / `write_epoch_test`）在本版已转为通过。 + +### 4.2 真实语料端到端训练（40HX） + +数据集 = TinyStories 前 20,000 条（`scripts/download_tinystories.py`）→ +**20,000 docs / 4,241,421 tokens / 8,208 词表 / 66,273 个 seq64 窗口**（`tokenizer_train` + +`dataset_gen` 生成 `datasets/tiny.nndataset`）。GPT `d_model=64 / heads=4 / layers=2 / +d_ff=128 / seq_len=64`，`--batch-size 8 --accum-steps 2`，**20 update step**： + +| 配方 | loss 轨迹（step 5/10/15/20） | 结果 | +|---|---|---| +| f32 | 8.7566 → 8.4056 → 8.0662 → **7.6110** | 1.2 s，无 NaN/冻结 | +| `--f16` | 8.7620 → 8.4353 → 8.0173 → **7.6741** | 0.4 s，无 NaN/冻结 | +| `--t1_58` | 8.9024 → 8.6983 → 8.4716 → **8.1851** | 0.5 s，无 NaN/冻结 | + +- **断点续训**：f32 `--save-interval 5` 产出 step20 `.nnckpt` → `--resume` 续 3 步， + step21–23 loss **7.6525 → 7.5733 → 7.5231**；日志确认优化器状态 76 张量 / `t=20`（adam）恢复。 +- **warm start**：f32 `.nnmodel` 续 1 步 loss 7.6361，文件 SHA-256 + `19A26875…D22952` → `77968371…C151B663`，确认参数确实移动。 + +### 4.3 四组「编译器 × 标准库」 + +| 配置 | 结果 | +|---|---| +| Windows clang 23.1 + **MSVC STL** / Release / Vulkan | 见 §4.1 第 1 行（73：72 pass + 1 skip） | +| Windows clang 23.1 + **MSVC STL** / Debug / Vulkan | 见 §4.1 第 2 行（73：72 pass + 1 skip） | +| WSL g++ 15.2 + **libstdc++** / Release / CPU-only | 见 [`../contributors/release-requirements.md`](../contributors/release-requirements.md) §10.3 | +| WSL clang 24.0 + **libc++** / Release / CPU-only | 见 [`../contributors/release-requirements.md`](../contributors/release-requirements.md) §10.3 | + +> `cl.exe + MSVC STL` 一组本机无 MSVC 工具链（`cl.exe` 不在 PATH，VS 未安装），**未在本机复跑**； +> 记为已知覆盖缺口，沿用 CI 的 MSVC 口径。WSL 的 GPU 用例按 77 skip（无 Vulkan）。 + +### 4.4 门禁 + +- `bench/doc_inventory.ps1`：**PASS**（`L2-VIOLATIONS: 0`、`HANDWRITTEN-RESIDUE: 0`、 + `TERMINATING-CALLS: 0`）； +- `bench/gui_cli_audit.py`：**PASS，可行动 0**（七节全绿，含断点续训禁列 32 项与多轮对话闭环）； +- `bench/audit_dangling.ps1`：exit 0。 + +## 5. 已知限制 + +1. `--t1_58` 仅 GPT 支持；`--model rapt` 选三值会被配置期拒绝（既有口径）。 +2. WSL CPU-only 树中 GPU/offload 用例按 77 skip（无 Vulkan 设备），非源码问题。 +3. 四编译器矩阵的 `cl.exe` 一组未在本机复跑（见 §4.3 脚注）。 +4. 全仓格式化使 `git blame` 在 `5854b98` 处断代；查历史行归属需 + `git blame -w` 或回看父提交。 +5. `TextTrainer` 仍只服务文本/GPT 域；通用 Step 协议由 `RunHooks::do_update` 自行实现 + （本版未抽象出通用 step 接口）。 +6. **二次续训的 Adam `t` 退化**（P1）：CLI 续训只恢复优化器状态，不恢复 + `TrainingSession::step_count()`，而 `TextTrainer::make_train_state` 用 `step_count()` 写 + `TrainState.optimizer_step`。因此「续训 → 再存档 → 再续训」链路中的 `t` 会退化为本段步数 + （首次续训不受影响：断点内的 `t` 取自优化器权威值）。修复方向：改用 + `session.optimizer().state_step()`。 +7. **`TrainingSession::resume` 错误路径半提交**（P1）：当前先 `load_model` 覆盖权重、后 + `verify_resume`，Strict 校验失败时模型已被改写却返回错误。修复方向：先读取并校验训练状态段， + 通过后再载入权重。 +8. 其他 P2：步进驱动在 `steps_per_epoch()==nullopt && max_steps==0` 时空跑并成功返回； + `save_checkpoint` 用 `optimizer_step` 覆盖 `counters.global_step`（与 RunDriver 的 + `epoch·steps_per_epoch+step` 定义不一致）；`InferenceScope::error()` 在未 `ok()` 时可解引用 + 空 optional（现有调用均先判 `ok()`）。 From 08cf4453973d93b9cad661eb9fa87fed415a4839 Mon Sep 17 00:00:00 2001 From: EthanPeng-2048 Date: Sat, 31 Oct 2026 14:54:04 +0800 Subject: [PATCH 8/8] =?UTF-8?q?fix(train):=20=E6=96=AD=E7=82=B9=E7=BB=AD?= =?UTF-8?q?=E8=AE=AD=E6=97=A0=E5=8D=8A=E6=8F=90=E4=BA=A4=20+=20=E4=BC=98?= =?UTF-8?q?=E5=8C=96=E5=99=A8=20t=20=E5=8F=96=E6=9D=83=E5=A8=81=E5=80=BC?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - TrainingSession::resume 先读状态段并交叉校验，通过后才载入权重， Strict 失败不再改写模型参数（issue #43） - 会话镜像 state_.optimizer_step 改为优化器权威 t（state_step()）， 不再用 TrainState.counters.global_step（issue #42） - TextTrainer::make_train_state 的 optimizer_step 改用 optimizer().state_step()， 修复「续训 → 再存档 → 再续训」的 Adam t 退化（issue #42） - session_checkpoint_test 新增二次续训 t 连续累计与失败路径逐位不变用例 - 版本号定档 v1.14.1 --- CMakeLists.txt | 2 +- include/neuralnet.cpp/text_trainer.hpp | 5 ++- include/neuralnet.cpp/training_session.hpp | 19 ++++++---- src/session_checkpoint_test.cpp | 44 +++++++++++++++++++++- src/text_trainer_test.cpp | 4 +- 5 files changed, 62 insertions(+), 12 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 803ce5e5..351bac4a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -10,7 +10,7 @@ cmake_minimum_required(VERSION 3.30) # VERSION 以 git tag 为准（本仓无 CONFIGURE_FILE 消费该字段， # 它只是元数据——发版时同步到最新 tag，实际版本以 git tag / gh release 为准） -project(neuralnet.cpp VERSION 1.14.0 LANGUAGES CXX) +project(neuralnet.cpp VERSION 1.14.1 LANGUAGES CXX) # ══════════════════════════════════════════════════════════════════════════ # 作为库被消费（add_subdirectory / FetchContent 嵌入 3rd_party）的门控 diff --git a/include/neuralnet.cpp/text_trainer.hpp b/include/neuralnet.cpp/text_trainer.hpp index 9de038c2..faaa1c9a 100644 --- a/include/neuralnet.cpp/text_trainer.hpp +++ b/include/neuralnet.cpp/text_trainer.hpp @@ -404,11 +404,14 @@ class TextTrainer } // 断点状态的权威装配（计数器 = "下一步的位置"，epoch 边界归一化）。 + // `optimizer_step` 取优化器权威 t（`state_step()`），**不是**会话镜像 + // `step_count()`：CLI 续训只恢复优化器、不同步会话计数，用镜像会让 + // 「续训 → 再存档」把 t 退化为本段步数（issue #42）。 [[nodiscard]] TextTrainState make_train_state(std::size_t epoch, std::size_t step_done, Scalar current_lr) const { return make_text_train_state(cfg_, identity_, steps_per_epoch_, epoch, step_done, - current_lr, session_.step_count()); + current_lr, session_.optimizer().state_step()); } TrainingSession session_; diff --git a/include/neuralnet.cpp/training_session.hpp b/include/neuralnet.cpp/training_session.hpp index 4eaa4522..0c33dd63 100644 --- a/include/neuralnet.cpp/training_session.hpp +++ b/include/neuralnet.cpp/training_session.hpp @@ -543,21 +543,26 @@ class TrainingSession return save_model(path, model_, spec, tokenizer_data, ModelRole::Checkpoint, ckpt); } - // 一站式续训：载入权重（同一 .nnckpt 文件）→ 读训练状态段 → 可选交叉 - // 校验（`current` = 本次训练的状态快照；配方/数据集身份不一致硬报错， - // Relaxed 只查 kind）→ 恢复优化器状态（t 先于/后于张量无序依赖，本入口 - // 固定先 t 后张量）→ 恢复会话计数器。返回加载出的状态（域读取数据集身份 - // 等字段用）。 + // 一站式续训：读训练状态段 → 可选交叉校验（`current` = 本次训练的状态 + // 快照；配方/数据集身份不一致硬报错，Relaxed 只查 kind）→ **校验通过后** + // 载入权重（同一 .nnckpt 文件）→ 恢复优化器状态（t 先于/后于张量无序依赖， + // 本入口固定先 t 后张量）→ 会话计数器镜像为优化器权威 t。返回加载出的 + // 状态（域读取数据集身份等字段用）。 + // + // 顺序即契约（issue #43）：权重载入放在校验之后，Strict 校验失败时模型 + // 参数保持调用前原样，不存在"半提交"（参数已改写却返回错误）。 [[nodiscard]] Result> resume(const std::string& path, const TrainState* current = nullptr, ResumePolicy policy = ResumePolicy::Strict) { - NN_TRY(lm, load_model(path, model_, nullptr, nullptr, WeightLoadIntent::Trainable)); NN_TRY(loaded, load_train_state(engine(), path)); if (current != nullptr) { NN_TRY(vr, verify_resume(*loaded->state, *current, policy)); } + NN_TRY(lm, load_model(path, model_, nullptr, nullptr, WeightLoadIntent::Trainable)); opt_->set_state_step(loaded->state->optimizer_step); NN_TRY(ls, opt_->load_state(loaded->optimizer_blobs)); - state_.optimizer_step = loaded->state->counters.global_step; + // 会话镜像 = 优化器权威 t（不是 TrainState.counters.global_step：后者是 + // "下一步位置"的绝对步号，与 Adam 的偏置修正计数不是同一量——issue #42）。 + state_.optimizer_step = opt_->state_step(); if (loaded->state->counters.lr > 0) set_lr(static_cast(loaded->state->counters.lr)); clear_accumulation_state_(); return std::move(loaded->state); diff --git a/src/session_checkpoint_test.cpp b/src/session_checkpoint_test.cpp index e2e88955..ed248b07 100644 --- a/src/session_checkpoint_test.cpp +++ b/src/session_checkpoint_test.cpp @@ -8,6 +8,8 @@ // [3] 交叉校验：配方不一致 Strict fail-fast、Relaxed 放行 // [4] 续训连通：resume 后继续 train_step，参数确实移动（优化器状态接上） // [5] 域字段不被覆盖：save_checkpoint 不清掉域预填的 dataset 身份 +// [6] 续训 → 再存档 → 再续训：优化器 t 连续累计不退化（issue #42） +// [7] 失败路径无半提交：Strict 校验失败后模型参数逐位不变（issue #43） // ───────────────────────────────────────────────────────────────────────── #include @@ -179,6 +181,8 @@ void test_roundtrip() auto resumed = snapshot(engine, session2.model()); CHECK(resumed == mid, "[1] resume 后权重与断点时刻逐位一致"); CHECK(session2.step_count() == 2, "[1] 会话计数器恢复 = 2"); + CHECK(session2.step_count() == session2.optimizer().state_step(), + "[1] 会话镜像 = 优化器权威 t（issue #42）"); // [3] 交叉校验：数据集身份变了 → Strict fail-fast；Relaxed 放行 auto bad = sample_domain_state(); @@ -186,6 +190,10 @@ void test_roundtrip() bad.dataset.sha256 = "0000000000000000000000000000"; auto strict_r = session2.resume(path, &bad, nn::ResumePolicy::Strict); CHECK(!strict_r.has_value(), "[3] Strict：数据集身份不一致 fail-fast"); + // issue #43：校验在载权重之前 ⇒ 失败路径不得改写模型参数 + auto after_strict = snapshot(engine, session2.model()); + CHECK(after_strict == mid, "[3] Strict 失败后模型参数与调用前逐位一致（无半提交）"); + CHECK(session2.step_count() == 2, "[3] Strict 失败后会话计数器不变"); auto relaxed_r = session2.resume(path, &bad, nn::ResumePolicy::Relaxed); CHECK(relaxed_r.has_value(), "[3] Relaxed：放行"); if (!relaxed_r) std::printf(" %s\n", relaxed_r.error().message.c_str()); @@ -201,6 +209,37 @@ void test_roundtrip() CHECK(before != after, "[4] resume 后参数确实移动"); CHECK(session2.step_count() == 3, "[4] 计数器继续累计 = 3"); } + + // ── [6] 续训 → 再存档 → 再续训：优化器 t 不退化（issue #42）── + const auto path2 = (scratch_dir() / "roundtrip2.nnckpt").string(); + auto domain2 = sample_domain_state(); + CHECK(session2.save_checkpoint(path2, tiny_spec(), domain2).has_value(), "[6] 再存档"); + auto peek2 = nn::peek_train_state(path2); + CHECK(peek2.has_value(), "[6] peek 二次断点"); + if (peek2) + { + const nn::TrainState& ps = **peek2; + CHECK(ps.optimizer_step == session2.optimizer().state_step(), + "[6] 二次断点里的 optimizer_step = 优化器真实 t"); + CHECK(ps.optimizer_step == 3, "[6] t 连续累计（2 → 3），不退化成本段步数"); + } + + auto model3_r = nn::build_gpt_model_from_spec(engine, tiny_spec(), nn::PrecisionProfile{}); + CHECK(model3_r.has_value(), "建模型 C"); + if (!model3_r) return; + auto ses3_r = nn::TrainingSession::create(std::move(*model3_r), cfg); + CHECK(ses3_r.has_value(), "create session C"); + if (!ses3_r) return; + auto session3 = std::move(*ses3_r); + auto cur2 = sample_domain_state(); + session3.fill_recipe(cur2.recipe); + auto st3_r = session3.resume(path2, &cur2, nn::ResumePolicy::Strict); + CHECK(st3_r.has_value(), "[6] 二次续训成功"); + if (st3_r) + { + CHECK(session3.step_count() == 3, "[6] 二次续训后会话镜像 = 3"); + CHECK(session3.optimizer().state_step() == 3, "[6] 二次续训后 Adam 的 t = 3"); + } } // ── [3] 配方不一致（非域字段）也应 fail-fast ────────────────────────────── @@ -243,9 +282,12 @@ void test_recipe_mismatch() auto session2 = std::move(*ses2_r); auto cur = sample_domain_state(); session2.fill_recipe(cur.recipe); // cfg2.optimizer = "sgd" ≠ 断点 "adam" + auto before = snapshot(engine, session2.model()); auto r = session2.resume(path, &cur, nn::ResumePolicy::Strict); CHECK(!r.has_value(), "[3] Strict：optimizer 配方不一致 fail-fast"); - // 注意：resume 内 load_model 已载入权重（断点参数），此处会话仍可用 + // issue #43：校验先于载权重 ⇒ 失败后模型仍是这次的初始化参数 + auto after = snapshot(engine, session2.model()); + CHECK(after == before, "[3] Strict 配方不一致失败后模型参数未被改写"); } } // namespace diff --git a/src/text_trainer_test.cpp b/src/text_trainer_test.cpp index e2e1662b..0bc2f8a9 100644 --- a/src/text_trainer_test.cpp +++ b/src/text_trainer_test.cpp @@ -424,8 +424,8 @@ void test_checkpoint_hook(const std::string& dir) CHECK(st.recipe.batch_size == cfg.batch_size && st.recipe.epochs == 2, "状态含训练配方（batch/epochs）"); CHECK(st.recipe.optimizer == "adam", "状态含优化器名"); - CHECK(st.optimizer_step == trainer_r->session().step_count(), - "状态含 Adam 的 t（= 优化器更新次数）"); + CHECK(st.optimizer_step == trainer_r->session().optimizer().state_step(), + "状态含 Adam 的 t（优化器权威 t，issue #42）"); CHECK(st.stride == cfg.stride, "状态含 stride（文本域扩展）"); // 落盘状态可被 peek/load 读回（与 CLI 断点同一文件格式）