OneDNN中跨原语共享Scratchpad的技术实现解析

背景与核心问题

在深度学习框架的底层优化中，OneDNN作为Intel推出的高性能计算库，其内存管理机制对性能有着决定性影响。开发者在实现自定义原语时，常会遇到需要在前向传播（forward）和反向传播（backward）之间共享临时内存的场景。本文针对Scratchpad（临时工作缓冲区）的跨原语复用问题，深入剖析技术原理和解决方案。

Scratchpad与Workspace的本质区别

OneDNN中存在两种易混淆的内存缓冲区：

1. Workspace

   • 专为训练场景设计，在forward和backward之间持久化存在
   • 由库自动管理生命周期，必须保持有效直到反向传播完成
   • 典型应用：卷积运算中的中间梯度数据缓存

2. Scratchpad

   • 临时性工作内存，仅存在于单次原语执行期间
   • 可配置为库管理或用户管理模式
   • 适用于任何传播类型（forward/inference/backward）

跨原语共享的技术挑战

当开发者尝试通过hint_fwd_pd传递原始指针时出现指针失效现象，这源于OneDNN的内存管理机制：

• Scratchpad默认在primitive执行完成后立即释放
• 直接传递裸指针会违反OneDNN的内存所有权模型
• 原语间的执行存在隐式内存屏障

解决方案与最佳实践

方案一：使用用户管理的Scratchpad

// 创建用户管理的scratchpad
memory::desc scratchpad_desc(...);
memory scratchpad_mem(scratchpad_desc, engine);

// 通过属性绑定到多个原语
primitive_attr attr;
attr.set_scratchpad_mode(scratchpad_mode::user);

// 前向和反向原语共享同一内存
forward_primitive(..., attr);
backward_primitive(..., attr, scratchpad_mem);

方案二：Workspace替代方案

对于必须持久化的数据：

1. 在primitive_desc创建时声明需要workspace
2. 通过query_workspace获取内存需求
3. 显式维护workspace内存的生命周期

关键注意事项

• 避免直接传递裸指针，使用OneDNN的内存抽象层
• 对于短期重用的临时缓冲区，优先考虑scratchpad
• 需要长期保持的数据应使用workspace
• 注意不同原语执行间的内存依赖关系

性能优化建议

1. 内存复用：通过memory对象池减少重复分配
2. 对齐配置：确保scratchpad满足硬件对齐要求
3. 大小预估：使用query_s64提前获取内存需求
4. 异步执行：配合stream控制内存可见性

总结

OneDNN通过精细的内存管理机制平衡了性能与安全性。理解scratchpad和workspace的设计哲学，掌握其正确的共享方式，能够帮助开发者在保持内存安全的前提下实现极致的计算性能。对于需要跨原语共享临时数据的场景，建议优先采用用户管理的scratchpad模式，并通过OneDNN提供的标准接口进行内存传递，这既符合库的设计规范，也能获得最佳的性能表现。