Rust 加 WASM 的 FFI 性能测试：不同数据传递方式的吞吐量对比数据

Rust 加 WASM 的 FFI 性能测试：不同数据传递方式的吞吐量对比数据

一、FFI 是 WASM 的阿喀琉斯之踵

在上一篇评估报告中我说 WASM 运行时推理延迟比原生慢 30%-120%。当时没展开说这 30% 差在哪里。经过两周的精细 benchmark，答案很清楚：超过 70% 的开销来自 FFI 边界的数据序列化和内存拷贝。

WASM 的线性内存模型决定了：宿主和 WASM 模块之间没法直接共享对象，每次跨越边界都要做数据转换。这篇文章是我对三种主流数据传递方式的吞吐量对比，全部基于 Rust → WASM（wasmtime 运行时）实测。

二、三种数据传递方式

三、Benchmark 设计与实现

3.1 测试数据结构

use serde::{Serialize, Deserialize};

/// 模拟 AI 推理场景中的 token 批次数据
/// 包含 token ID 数组和对应的注意力掩码
#[derive(Clone, Serialize, Deserialize)]
pub struct TokenBatch {
    /// token ID 列表（可变长度，1~2048）
    pub tokens: Vec<u32>,
    /// 注意力掩码矩阵（二维数组，展平为一维）
    pub attention_mask: Vec<f32>,
    /// 批次元数据
    pub batch_size: u32,
    pub max_seq_len: u32,
}

3.2 方式一：JSON 序列化

/// 方式一：JSON 序列化传递
/// 将 TokenBatch 序列化为 JSON 字符串，写入 WASM 线性内存
fn benchmark_json_serialization(batch: &TokenBatch, store: &mut wasmtime::Store<()>, memory: &wasmtime::Memory) {
    // 序列化为 JSON 字符串
    let json = serde_json::to_string(batch).unwrap();
    let bytes = json.as_bytes();
    
    // 在 WASM 线性内存中分配空间
    let alloc = instance.get_typed_func::<i32, i32>(store, "alloc").unwrap();
    let ptr = alloc.call(store, bytes.len() as i32).unwrap();
    
    // 将数据写入 WASM 内存
    let mem_data = memory.data_mut(store);
    mem_data[ptr as usize..ptr as usize + bytes.len()].copy_from_slice(bytes);
    
    // 调用 WASM 处理函数
    let process = instance.get_typed_func::<(i32, i32), i32>(store, "process_json").unwrap();
    process.call(store, (ptr, bytes.len() as i32)).unwrap();
}

3.3 方式二：指针 + 布局描述

/// 方式二：通过结构体布局描述直接传递原始数据
/// 避免序列化开销，但依然需要拷贝内存
#[repr(C)] // 确保 Rust 和 WASM 端看到相同的字段排列
struct TokenBatchLayout {
    tokens_ptr: u32,     // 指向 token 数组的指针（在 WASM 内存中）
    tokens_len: u32,     // token 数组长度
    mask_ptr: u32,       // 指向掩码数组的指针
    mask_rows: u32,      // 掩码行数
    mask_cols: u32,      // 掩码列数
}

fn benchmark_raw_pointer(batch: &TokenBatch, store: &mut wasmtime::Store<()>, memory: &wasmtime::Memory) {
    // 步骤1：分配 WASM 内存并拷贝 token 数据
    let token_bytes = bytemuck::cast_slice(&batch.tokens);
    let token_ptr = wasm_alloc(store, token_bytes.len());
    memory.data_mut(store)[token_ptr..token_ptr + token_bytes.len()]
        .copy_from_slice(token_bytes);
    
    // 步骤2：分配并拷贝掩码数据
    let mask_bytes = bytemuck::cast_slice(&batch.attention_mask);
    let mask_ptr = wasm_alloc(store, mask_bytes.len());
    memory.data_mut(store)[mask_ptr..mask_ptr + mask_bytes.len()]
        .copy_from_slice(mask_bytes);
    
    // 步骤3：组装布局描述结构体
    let layout = TokenBatchLayout {
        tokens_ptr: token_ptr as u32,
        tokens_len: batch.tokens.len() as u32,
        mask_ptr: mask_ptr as u32,
        mask_rows: (batch.attention_mask.len() / batch.max_seq_len as usize) as u32,
        mask_cols: batch.max_seq_len,
    };
    
    // 步骤4：传递布局描述给 WASM
    let layout_ptr = wasm_alloc(store, std::mem::size_of::<TokenBatchLayout>());
    let layout_bytes = bytemuck::bytes_of(&layout);
    memory.data_mut(store)[layout_ptr..layout_ptr + layout_bytes.len()]
        .copy_from_slice(layout_bytes);
    
    let process = instance.get_typed_func::<i32, i32>(store, "process_raw").unwrap();
    process.call(store, layout_ptr as i32).unwrap();
}

/// WASM 端的辅助函数：分配线性内存
fn wasm_alloc(store: &mut wasmtime::Store<()>, size: usize) -> usize {
    let alloc = instance.get_typed_func::<i32, i32>(store, "alloc").unwrap();
    alloc.call(store, size as i32).unwrap() as usize
}

3.4 方式三：直接内存共享（wasm-bindgen 模式）

/// 方式三：宿主和 WASM 协商好内存布局，零拷贝
/// 要求：双方使用相同的 #[repr(C)] 结构体定义
fn benchmark_zero_copy(batch: &TokenBatch, store: &mut wasmtime::Store<()>, memory: &wasmtime::Memory) {
    // 直接获取 WASM 线性内存的可变引用
    let mem = memory.data_mut(store);
    
    // 在 WASM 内存栈上（预分配区域）直接写入数据
    // 注意：生产环境需要用更安全的分配策略
    const OFFSET: usize = 1024; // 预留给 WASM 内部使用的区域之后
    
    // 写入 token 计数 + 数据
    let token_count = batch.tokens.len() as u32;
    mem[OFFSET..OFFSET + 4].copy_from_slice(&token_count.to_le_bytes());
    
    let token_slice = &mem[OFFSET + 4..OFFSET + 4 + batch.tokens.len() * 4];
    // 注意：这里不能直接 cast，需要用安全的方式写入
    // 实际项目中建议用 bytemuck 或 zerocopy crate
    
    // WASM 函数只接收一个偏移量参数
    let process = instance.get_typed_func::<i32, i32>(store, "process_zerocopy").unwrap();
    process.call(store, OFFSET as i32).unwrap();
}

生产翻车现场：上面这个零拷贝方案在 x86 机器上跑得飞快，但部署到 Apple Silicon (M2) 上时，推理结果全是乱码。排查了一整天发现：#[repr(C)] 结构体在 x86 上是紧凑排列的，但在 ARM64 上结构体末尾会插入 padding 字节。WASM 模块用的是 x86 内存布局，宿主是 ARM64 机器，copy_from_slice 时多读了 4 字节的 padding 数据。修复方案：在 TokenBatchLayout 上显式标注 #[repr(C, packed)] 并用 bytemuck::AnyBitPattern 做对齐检查。

生产实战经验：MessagePack 的内存泄露陷阱

除了 ARM64 对齐问题，序列化格式的选择也有坑。起初用 JSON 传 2048 tokens 数据，FFI 耗时 3.8ms。换成 MessagePack 降到 2.1ms（省 45%），但连续推理 100 次后，Chrome DevTools Memory 面板显示 WASM 线性内存从 12MB 涨到 47MB——存在渐进式内存泄露。排查发现 MessagePack 反序列化时每个字段都触发 WASM 侧 alloc 分配小块内存，释放不及时导致碎片累积。

最终用 bytemuck::Pod 定义固定内存布局，双方直接按 struct 读取：

/// 用 bytemuck 定义跨 FFI 的原始字节布局
/// 无需序列化/反序列化，双方按字节拷贝即可
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
#[repr(C)]
struct TokenBatchFfi {
    tokens_ptr: u32,
    tokens_len: u32,
    mask_ptr: u32,
    mask_rows: u32,
    mask_cols: u32,
}

bytemuck::Pod 保证 struct 是"纯数据"——无指针、无 padding、无复杂字段，可安全跨边界按字节传输。改造后 FFI 耗时从 2.1ms 降到 0.8ms，连续推理 500 次内存稳定在 13MB。

四、Benchmark 结果

结论：

• 小数据量（<1KB）：三种方式差异不大，用 JSON 最简单；
• 中等数据（1KB~100KB）：原始指针是性价比最高的方案；
• 大数据量（>100KB）：零拷贝方案优势明显，尤其在 AI 推理场景中，每次传递 2048 tokens 的 embedding 矩阵（约 8MB）。

线上实测数据：在我们的 AI CLI 插件中，一次典型推理调用需要传递 token_ids（~2KB）+ attention_mask（~8KB）+ 模型元数据（~500B）。用 JSON 方案时 FFI 耗时 3.8ms，原始指针方案 1.1ms。但因为 FFI 只占整个推理延迟（45ms）的一小部分，我们最终选了 JSON 方案——开发效率提升远大于 2.7ms 的性能损失。不是所有场景都需要零拷贝，先 benchmark 再决定。

端到端延迟拆解：WASM vs 原生

对 2048 tokens 场景按环节拆解各阶段延迟：

WASM 推理计算本身只比原生慢约 5%（19.1ms vs 18.2ms），差距不大。主要额外开销来自 FFI 序列化——JSON 方案的 FFI 相关延迟达 5ms（占总延迟 20%），bytemuck 方案降至 1.1ms（占比 5%）。数据量 <1KB 时 JSON 的简便性远大于性能损失，超过 10KB 后直接内存布局的收益非常明显。

五、总结

Rust + WASM 的 FFI 性能瓶颈不在 WASM 指令执行本身，而在数据跨越边界的成本。三个实用建议：

1. 优先减少 FFI 调用次数——与其传 100 次小数据，不如合并成一次大数据调用；
2. 对性能敏感路径用原始指针——serde 反序列化在 WASM 侧的 alloc 是主要瓶颈；
3. 零拷贝适用于固定长度数据——如果你的数据结构大小在编译期已知，#[repr(C)] + 内存布局协商是最优解。

这次 benchmark 也让我重新审视了 AI CLI 工具的架构：如果未来要支持 WASM 插件做自定义推理后处理，FFI 路径的设计会直接决定插件的性能天花板。好消息是，Rust 的类型系统让我们在追求零拷贝性能的同时，不至于完全放弃内存安全。

下一篇预告：AI Agent 的用户体验设计——loading 状态、错误提示和置信度展示的最佳实践。