Rust in Everything
基于 candle-core 0.11.0

inplace_op2 卡死

主线位置：算子主线的「取 storage 锁」那一步。普通算子对输入只拿读锁；原地算子对第一个参数拿写锁，对其余参数拿读锁。锁在 Tensor_ 的 storage: Arc<RwLock<Storage>> 上，clone、narrow 之类的视图共享同一把锁，见 Tensor 的结构；原地算子的接口见 自定义算子。

sequenceDiagram
  participant U as 调用方
  participant L as 同一把 RwLock
  U->>L: storage_mut 拿写锁
  L-->>U: 写锁到手
  U->>L: rhs.storage 拿读锁
  Note over U,L: 写锁还在自己手里
  Note over U,L: 读锁永远等不到

实验 crate exp-candle/pb-pit，bin deadlock，debug，M2 Max、macOS 27、rustc 1.99.0。

现象

一个最普通的原地加法 lhs += rhs：

struct AddInto;

impl InplaceOp2 for AddInto {
  fn name(&self) -> &'static str {
    "add-into"
  }
  fn cpu_fwd(&self, s1: &mut CpuStorage, l1: &Layout, s2: &CpuStorage, l2: &Layout) -> Result<()> {
    let (CpuStorage::F32(a), CpuStorage::F32(b)) = (s1, s2) else {
      candle_core::bail!("add-into: f32 only")
    };
    let (Some((a0, a1)), Some((b0, b1))) = (l1.contiguous_offsets(), l2.contiguous_offsets()) else {
      candle_core::bail!("add-into: contiguous only")
    };
    a[a0..a1].iter_mut().zip(&b[b0..b1]).for_each(|(x, y)| *x += *y);
    Ok(())
  }
}

两个不同的张量没问题。接着两种写法：把张量 clone() 一份加到自己身上（self 模式），以及一个矩阵的第 1 行加到第 0 行（view 模式）：

// self
let x2 = x.clone(); // cheap clone: same Arc<RwLock<Storage>>
x.inplace_op2(&x2, &AddInto)?;

// view
let m = Tensor::new(&[[1f32, 2.], [3., 4.]], &Device::Cpu)?;
let row1 = m.narrow(0, 1, 1)?.squeeze(0)?;
let row0 = m.narrow(0, 0, 1)?.squeeze(0)?;
row0.inplace_op2(&row1, &AddInto)?;

后台启动，5 秒后看进程状态，再杀掉（只杀这个 pid）：

--- self
x += y -> [11.0, 22.0, 33.0, 44.0]
pid 18933: x += x.clone() ...
  PID STAT  %CPU ELAPSED COMMAND
18933 SN     0.0   00:05 ../target-pb/debug/deadlock self
exit=143 (killed after 5 s)
--- view
x += y -> [11.0, 22.0, 33.0, 44.0]
pid 18993: row0 += row1 (views of one matrix) ...
  PID STAT  %CPU ELAPSED COMMAND
18993 SN     0.0   00:05 ../target-pb/debug/deadlock view
exit=143 (killed after 5 s)

没有错误，没有 panic，CPU 占用 0.0%：线程睡在锁上。row0 和 row1 在数据上毫不重叠，照样卡住，因为锁是按 storage 加的，不是按元素范围。

出处

Tensor::inplace_op2：

    pub fn inplace_op2<C: InplaceOp2>(&self, rhs: &Self, c: &C) -> Result<()> {
        self.storage_mut()
            .inplace_op2(self.layout(), &rhs.storage(), rhs.layout(), c)
    }

self.storage_mut() 拿到写锁守卫，这个守卫在整条表达式结束前一直活着；接着 rhs.storage() 要同一把锁的读锁。两个张量共享 storage 时，读锁要等写锁释放，写锁要等这次调用返回，谁也等不到谁。inplace_op3 一样（写锁之后两个读锁），inplace_op1 只有一个参数，不会这样。锁在分派到设备之前就拿了，所以 Metal 上也一样：AddInto 没有写 metal_fwd，Metal 张量对自己的 clone 调用时照样卡住（metal-self 模式，5 秒后被杀），换成 copy() 才走到分派，报出 Err("Metal error no metal implementation for add-into")。

卡住的进程用 macOS 的 sample 看主线程（截取）：

    792 Thread_3995163   DispatchQueue_1: com.apple.main-thread  (serial)
                      792 _RINvMs_NtCs44JGuoGRrvl_11candle_core9custom_opNtNtB7_6tensor6Tensor11inplace_op2NtCsgwg5CnLzTaS_8deadlock7AddIntoEB1i_  (in deadlock) + 132  [0x1046767b4]  custom_op.rs:363
                          792 _RNvMs9_NtNtNtCs1SSuMAdGM58_3std4sync6poison6rwlockINtB5_6RwLockNtNtCs44JGuoGRrvl_11candle_core7storage7StorageE4readB13_  (in deadlock) + 72  [0x1046b9dcc]
                            792 _RNvMs1_NtNtNtNtCs1SSuMAdGM58_3std3sys4sync6rwlock5queueNtB5_6RwLock14lock_contended  (in deadlock) + 280  [0x10474e508]

custom_op.rs:363 就是 rhs.storage() 那一行，下面是 RwLock<Storage>::read 和 lock_contended。

candle 自己的原地方法知道这个问题。slice_set、scatter_set、Var::set 在拿锁之前先用 same_storage 比较两把 RwLock 的地址，相同就返回错误；inplace_op1/2/3 没有这道检查，而 same_storage 是 pub(crate)，用户代码调不到。

📝NOTE
没有蠢问题

问： 标准库不是说同一线程重复加锁「might panic」吗，为什么这里不 panic？

答： std::sync::RwLock 的文档对 read 和 write 都写的是「This function might panic when called if the lock is already held by the current thread」：可能 panic，也可能不，取决于平台的锁实现。macOS 上它就是一直阻塞。所以不能指望靠 panic 发现这个问题，只能在调用前避开。

修法

第二个参数和第一个共享存储时，先 copy() 出一份独立的数据。判断是否共享，可以用公开的 storage_and_layout：它拿的是读锁，两个读锁可以同时持有，比较两个守卫指向的 Storage 地址即可：

/// True when both tensors read from the same storage (two read locks at once are fine).
fn shares_storage(a: &Tensor, b: &Tensor) -> bool {
  let (sa, _) = a.storage_and_layout();
  let (sb, _) = b.storage_and_layout();
  std::ptr::eq(&*sa, &*sb)
}

let rhs = if shares_storage(&x, &x2) { x2.copy()? } else { x2 };
x.inplace_op2(&rhs, &AddInto)?;

let row0 = m.narrow(0, 0, 1)?.squeeze(0)?;
let row1 = m.narrow(0, 1, 1)?.squeeze(0)?.copy()?;
row0.inplace_op2(&row1, &AddInto)?;

x += y -> [11.0, 22.0, 33.0, 44.0]
shares_storage(x, x.clone()) = true
shares_storage(x, x.copy()) = false
shares_storage(x, y) = false
x += copy of x -> [22.0, 44.0, 66.0, 88.0]
row0 += copy of row1 -> m = [[4.0, 6.0], [3.0, 4.0]]
x + x (new tensor) -> [44.0, 88.0, 132.0, 176.0]
  • shares_storage 必须在调用 inplace_op2 之前、在它自己的作用域里比完：两个读守卫在函数返回时释放，之后再拿写锁才不冲突。
  • 原地写的是 row0 这个视图，m 跟着变了（第 0 行成了 [4, 6]），这正是原地算子的用处。
  • 多数时候根本不需要原地算子：最后一行的 &x + &x 生成新张量，两边都只拿读锁，读读可以并存。原地算子只在明确要省一块内存、且确认参数不共享 storage 时才用。
  • Metal 上 copy() 本身有一个上游问题：它不真正复制缓冲区，见 改了 copy 出来的张量，原张量也变了。在 Metal 上 copy() 能解开死锁（返回的张量有自己的 RwLock，上面 shares_storage 为 false），但两边仍指向同一块显存，原地写会同时改到两边。

怎么确认没中招

  • 进程卡住、CPU 占用为 0、没有任何输出时，用 sample <pid>（macOS）或 gdb -p 看主线程栈，像上面那样停在 RwLock<Storage>::read 里、上一帧是 inplace_op2 / inplace_op3，就是这个坑。
  • 代码里每个 inplace_op2 / inplace_op3 的调用处，问一句：后面的参数会不会是第一个参数的 clone、narrow、i(..)、reshape（连续时是视图）或者 detach？会就加上 shares_storage 判断或无条件 copy()。
  • 测试里用同一个张量的两个视图调一次，套 recv_timeout 或进程级超时，比等它在线上卡住强。

小结

  • inplace_op2 / inplace_op3 先拿第一个参数的写锁、再拿其余参数的读锁；共享 storage 的两个张量（clone、同一矩阵的两个视图）会让它永远阻塞，不报错、不 panic。
  • candle 的 slice_set 等方法用 same_storage 挡住了这种情况，原地算子没有，same_storage 也不公开。
  • 用 storage_and_layout 比较两个 Storage 的地址，共享就先 copy()；能用非原地运算就别用原地算子。

💬 关联讨论

发起讨论
加载中...