PyTorch 写时复制存储Copy-on-Write Storage解析从线程安全模型到源码实现【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch导读c10/core/impl/README-cow.md描述了 PyTorch 中为张量Tensor引入写时复制存储Copy-on-Write Storage即惰性拷贝 lazy copy的核心设计。它的关键约束是继续维护张量互为别名当且仅当它们共享同一 Storage这一 PyTorch 不变式——即惰性拷贝得到的张量拥有各自独立distinct的 Storage而这些 Storage 之间再共享同一份底层数据分配。读完本文你将理解为什么惰性拷贝天然与多线程安全模型冲突、运行时如何通过共享锁复制 / 独占锁窃取两类原语同时保证正确性与性能以及这套机制在c10/core/impl/下各源码文件中的具体落地形态。一、设计动机在保持 Storage 别名不变式的前提下实现惰性拷贝PyTorch 的语义约定是两个张量若共享同一块底层存储则互为别名任何一方就地写入另一方可见反之若张量拥有各自独立的 Storage则它们之间不存在别名关系。传统的深拷贝clone会立刻复制一份完整数据代价高昂而写时复制希望把复制推迟到真正发生写入的那一刻。为此设计引入了一个新的中间概念——写时复制上下文copy-on-write context每个参与惰性拷贝的张量拥有独立的StorageImpl但这些 Storage 的DataPtr共同引用同一个共享的 COW 上下文而真正的数据分配由该上下文持有。用 COW.h 中的注释来说这是对给定存储创建一个 COW 克隆若该存储本身还不是 COW则会先把它转换成 COW 存储。这种设计保持了对用户可见的语义稳定性只要你不写入多个惰性拷贝可以像共享底层存储一样高效零数据复制一旦发生写入运行时会自动把数据物化materialize把该张量与其兄弟副本真正拆分开。二、线程安全的前提约定与惰性拷贝引入的新难题原文档明确指出整套设计的正确性建立在一个 PyTorch 既有的用户约定之上——用户有责任保证写入不会与读取或其他写入并发发生。这是整个 C 生态默认的编程假设也是 PyTorch 长期以来对多线程使用张量的基本要求。惰性拷贝给这个模型带来了新的复杂度用户无需知道某个张量是否已存在惰性拷贝因此用户也没有义务在多个惰性拷贝之间做写入串行化。举例来说两个 Storage 不同、但共享同一份 COW 数据上下文的张量完全可能被分别交给两个线程随意读写此时运行时必须自行保证安全。原文档随后论证这其实并不难防护——因为写时复制的语义决定了处理一次写入只需在写入时物化该张量。如果对每一份拷贝都无条件做一次真实复制那么整个过程甚至可以完全不需要同步但工程上存在一个更聪明的通用优化——对最后剩余的那一份引用可以省去复制、直接窃取数据而这个优化正是需要等待所有进行中的拷贝完成的原因。三、核心模型lazy-clone 是读materialization 是写设计将影响张量 COW 状态的运算抽象为两类lazy-clone惰性克隆包括显式调用或由 reshape 之类的算子内部隐藏触发逻辑上等价于只读操作materialization物化即任何对张量的写入写入时触发数据实体化。3.1 共享同一 Storage 的张量集合内部设计的关键洞察是lazy-clone 在逻辑上是读操作materialization 在逻辑上是写操作。因此对于一组共享同一 Storage 的张量而言只要有物化正在进行就不会有包括 lazy-clone 在内的任何读操作与之并发——这一约束可以直接沿用上述用户负责读写同步的既有模型无需新增同步。3.2 跨 Storage 共享 COW 上下文的世界上面这条洞察只适用于共享同一 Storage的张量集合。设计还必须处理更危险的情形Storage 不同、但共享同一个写时复制上下文的张量。在这个世界里materialization 可能与 lazy-clone 竞争甚至可能与其他 materialization 竞争。此时安全性来自一个结构性事实既然存在竞争就必然至少有两个对上下文的引用意味着你正在执行 lazy-clone 时上下文不可能凭空消失因此 lazy-clone 只需要一次**原子引用计数自增atomic refcount bump**即可保证安全。3.3 最复杂的情形所有惰性拷贝并发物化原文档特别推演了最坏情形——所有惰性拷贝同时开始物化因为物化是写操作此时必然没有任何进行中的 lazy-clone需要保证的是所有物化操作能够并发地安全读取共享数据并各自复制出自己的副本若没有最后一份引用窃取数据的优化这里可以做到完全不加锁每个物化各复制一份并递减引用计数即可但由于该优化存在物化竞争中的失败方loser必须等待其他进行中的拷贝结束然后再直接窃取数据而非复制。四、锁策略共享锁复制、独占锁窃取原文档为上述推理给出了最终实现策略这也是整套机制最精炼的总结复制数据时持共享锁shared lock多个物化可以并发地读共享数据并各自复制互不阻塞窃取数据时持独占锁exclusive lock获取独占锁这一动作天然保证——在窃取发生前所有持共享锁的进行中拷贝都已完成从而可以安全地零复制收编最后一份数据。结合仓库中 COWDeleter.h 的代码注释可以印证COWDeleterContext::decrement_refcount()通过返回值区分两种情况分别对应上述两种加锁路径// 共享锁仍存在其他引用调用方需在共享锁保护下复制数据 using NotLastReference std::shared_lockstd::shared_mutex; // 独占这是上下文最后一份引用且在返回前已保证所有进行中的拷贝完成 using LastReference std::unique_ptrvoid, DeleterFnPtr; std::variantNotLastReference, LastReference decrement_refcount();五、源码级实现剖析原文档只给出设计层面的描述仓库在 c10/core/impl 下提供了完整实现可分为三个文件理解存储克隆与物化入口 COW.cpp、公开 API COW.h以及引用计数与锁的管理器 COWDeleter.cpp。5.1 上下文与引用计数管理器 COWDeleterContextCOWDeleter.h 定义了核心数据结构COWDeleterContext其成员即实现要点std::shared_mutex mutex_; // 支撑共享锁复制/独占锁窃取 std::unique_ptrvoid, DeleterFnPtr data_; // 唯一一份真正共享的数据分配 std::atomicstd::int64_t refcount_ 1; // 原子引用计数increment_refcount()仅做一次原子自增refcount_且断言自增后大于 1——这正是 lazy-clone 在 3.2 节中所说的只需要原子引用计数自增。decrement_refcount()是整套机制的关键见 COWDeleter.cpp先将计数原子自减当减到 0 时持unique_lock把data_移出、释放锁、随后delete this销毁上下文并返回LastReference否则返回一个std::shared_lock(mutex_)即NotLastReference让调用方在共享锁保护下去复制数据。构造函数断言传入数据的 deleter 绝不是cow_deleter自身防止把 COW 上下文再包一层 COW析构函数则断言引用计数已归零见 COWDeleter.cpp。同时COW.h 声明了cow_deleter(void* ctx)它被用作DataPtr的ctx_deleter实现见 COWDeleter.cpp——只是把指针还原为COWDeleterContext*并调用decrement_refcount()。5.2 如何判定一个 DataPtr 是否为 COWCOW.cpp 提供了两个判定函数is_cow_data_ptr直接比较该 DataPtr 的 deleter 函数指针是否为cow::cow_deleter用于快速判断是否已处于 COW 状态has_simple_data_ptr当 Storage 存在 Allocator 时委托给allocator-is_simple_data_ptr(data_ptr)否则要求上下文指针get_context()与数据指针get()相等——即无异常上下文的简单 DataPtr这是能否把存储转换为 COW 的判定依据。5.3 lazy_clone_storage 的三种情形lazy_clone_storageCOW.cpp依据数据指针的形态分三条路径处理这与原文档的线程安全讨论一一对应简单 DataPtr普通存储从原 Storage 取出原始 context 所有权用new COWDeleterContext(...)包装成 COW 上下文再把当前 Storage 本身改写成 COW 状态set_data_ptr_noswapset_materializer(materialize_cow)最后新建一个共享同一上下文的 StorageImpl 返回。此情形无需加锁。已是 COW DataPtr原 Storage 必然已挂有 materializer有内部断言只需对上下文做一次引用计数自增并返回一个新的 StorageImpl。此情形同样无需加锁——lazy-clone 本身是读而当前上下文持有活引用不会消失。存在非 COW 的异常 context无法安全转换直接返回nullptr表示不支持。新建的 StorageImpl 会保留原 Storage 的字节大小、Allocator、resizable标志与设备类型并挂上同样的 materializer见 COW.cpp。5.4 materialize_cow 的两条物化路径materialize_cowCOW.cpp即上文所说的写时物化逻辑与原文档第五节完全对应auto result ctx-decrement_refcount(); if (std::holds_alternativecow::COWDeleterContext::LastReference(result)) { // 唯一引用失败方等待所有 pending 拷贝结束后直接窃取数据、零复制 new_data_ptr DataPtr(data.release(), data_ptr.get(), data.get_deleter(), ...); } else { // 非最后引用持有共享锁结果对象调用 allocator 的 clone 复制数据 new_data_ptr storage-allocator()-clone(data_ptr.get(), storage-nbytes()); }LastReference 分支这是并发物化竞争中的失败方路径。此时其他并发拷贝都已通过共享锁完成或放弃上下文把唯一一份数据移交给它materialize_cow直接以原数据构造新的 DataPtr完全避免一次复制——这就是最后剩余引用窃取数据优化的实现点。NotLastReference 分支返回的shared_lock对象在函数存活期间保持持有代码注释明确说明不需要消费该结果它只是一个确保数据在复制期间存活下去的共享锁随后调用allocator()-clone做一次真实深拷贝。两条路径末尾都用set_data_ptr_no_materialize换上新的 DataPtr并用release_context()释放对旧上下文的那份引用避免引用计数被二次递减因为decrement_refcount已经做过一次。另外值得注意materialize_cow开头断言当前不处于at::parallel_for的并行循环体内!c10::ParallelGuard::is_enabled()即物化在并行 for 的循环函数中被禁止见 COW.cpp。六、从 ATen 算子到 Python 层的接入COW 的顶层入口是 ATen 算子_lazy_clone。在 AutogradComposite.cpp 的合成实现中它先取得输入张量的 Storage然后直接调用c10::impl::cow::lazy_clone_storage(*self_storage)返回的新 Storage 重新包装出张量从而实现显式调用触发惰性克隆。该算子在native_functions.yamlaten/src/ATen/native/native_functions.yaml中亦有登记。在 Python 侧Tensor._lazy_clone()正是把这一机制暴露给用户的入口。torch/_tensor_docs.py见 torch/_tensor_docs.py 一带对相关方法的文档明确指出一旦调用写入类方法就会**物化materialize**由_lazy_clone创建的底层惰性拷贝。这从 API 语义上验证了本文第三节的抽象——lazy-clone 只做引用层面的共享任何真实写入如copy_之类都会触发存储物化并打断共享。值得注意的是_lazy_clone这种带隐藏内部语义的算子还会出现在 torch/overrides.py 等对__torch_function__协议做处理的位置暗示编译器/分发系统需要把它当作特殊算子对待。七、设计要点小结与阅读指引将原文档的推演与源码相互印证可以提炼出几条必须牢记的结论设计要点说明别名不变式张量互为别名 ⟺ 共享同一 Storage惰性拷贝间共享的是COW 上下文/数据分配而非 Storage 本身线程安全前提用户负责保证对普通共享 Storage张量的读写不并发这一约定依旧有效读写性质抽象lazy-clone 读materialization 写仅在同 Storage 集合内成立跨 Storage 竞争只要存在竞争上下文引用数 ≥ 2lazy-clone 仅需原子自增即可安全并发物化复制数据走共享锁最后一份窃取数据走独占锁独占锁保证所有 pending 复制先结束代价边界除最后一份引用可零复制收编外其余每次物化都触发一次真实分配复制希望继续深入阅读的读者可按以下顺序逐文件精读设计文档c10/core/impl/README-cow.md公开 API 与语义注释c10/core/impl/COW.h存储克隆、判定与物化实现c10/core/impl/COW.cpp引用计数与锁管理器c10/core/impl/COWDeleter.h、c10/core/impl/COWDeleter.cppATen/Python 接入AutogradComposite.cpp、torch/_tensor_docs.py这套设计把读/写并发由用户负责这一传统模型与惰性拷贝存在但用户无需感知的新需求做了清晰的切分所有需要在运行时兜底的并发风险都被收敛到一次原子引用计数与一把读写锁之上从而在正确性、实现复杂度与零复制性能优化三者之间取得了平衡。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考