一次 matmul

主线位置：算子主线从头到尾。Tensor::matmul → Storage::matmul → CpuStorage::matmul → gemm / sgemm；或者 → MetalStorage::matmul → call_mlx_gemm 编码进命令缓冲区 → 读数据时提交并等待。每一步的源码和细讲的页见怎么读这一卷的算子主线表，本页用实验把它们一层层露出来。

sequenceDiagram
  participant M as main 线程
  participant T as Tensor
  participant S as Storage
  participant K as MetalStorage
  participant Q as Commands
  participant G as GPU
  M->>T: a.matmul(b)
  T->>S: 形状检查后 matmul
  S->>K: Metal 对 Metal
  K->>Q: 取编码器 编码 gemm
  K-->>M: 返回新 Tensor
  M->>T: c.to_vec2()
  T->>K: to_cpu 编码 blit
  K->>Q: flush_and_wait_current
  Q->>G: commit
  G-->>Q: 完成
  Q-->>M: 读回的数据

实验程序

两个实验 crate，代码相同，只差 feature：intro-mm 依赖 candle-core 0.11.0 并开 metal、accelerate；intro-mm-noacc 只开 metal，CPU 上的 matmul 就落到 gemm crate。walk 模式是 base 里那个 2×2 的 matmul，外加一条时间线：

fn walk(dev: &str) -> anyhow::Result<()> {
  let device = device(dev)?;
  let a = Tensor::arange(0f32, 4.0, &device)?.reshape((2, 2))?;
  let t = Instant::now();
  let c = a.matmul(&a)?;
  let t_mm = t.elapsed();
  let v = c.to_vec2::<f32>()?;
  let t_read = t.elapsed();
  println!("{dev}: {v:?}; matmul returned after {}, to_vec2 returned after {}", us(t_mm), us(t_read));
  Ok(())
}

两个 crate 都用 [patch.crates-io] 换上一份 candle-metal-kernels 0.11.0 的拷贝，只在 load_library 和 load_pipeline_with_constants 里加了几行计时，设了环境变量 INTRO_TRACE 才打印 LIB / PIPE 两种行，其余代码和 tag 上一字不差。另一个只在本页用到的开关 INTRO_NONCE 见下文「第一次和之后」。

前半段：Tensor::matmul 和 Storage::matmul

Tensor::matmul 不碰数据，只做形状上的事：

    pub fn matmul(&self, rhs: &Self) -> Result<Self> {
        let a_dims = self.shape().dims();
        let b_dims = rhs.shape().dims();

        let dim = a_dims.len();

        if dim < 2 || b_dims.len() != dim {
            // ... Err(ShapeMismatchBinaryOp)
        }

        let m = a_dims[dim - 2];
        let k = a_dims[dim - 1];
        let k2 = b_dims[dim - 2];
        let n = b_dims[dim - 1];

        let c_shape = Shape::from(&a_dims[..dim - 2]).extend(&[m, n]);
        if c_shape.elem_count() == 0 || k == 0 {
            return Tensor::zeros(c_shape, self.dtype(), self.device());
        }
        let batching: usize = a_dims[..dim - 2].iter().product();
        let batching_b: usize = b_dims[..dim - 2].iter().product();
        if k != k2 || batching != batching_b {
            // ... Err(ShapeMismatchBinaryOp)
        }

        let storage = self.storage().matmul(
            &rhs.storage(),
            (batching, m, n, k),
            self.layout(),
            rhs.layout(),
        )?;
        let op = BackpropOp::new2(self, rhs, Op::Matmul);
        Ok(from_storage(storage, c_shape, op, false))
    }

两边的秩必须相同，批维不广播；要广播用 broadcast_matmul，它先 broadcast_as 再 contiguous，把批维真的复制出来。
批维被压成一个数 batching，后端只看到 (b, m, n, k) 四个数和两个 Layout；步长合不合要求由后端判断（见布局与视图）。
self.storage() 和 rhs.storage() 各取一把 RwLock 读锁，a.matmul(&a) 对同一把锁取两次读锁也没问题（见 Tensor 的结构）。
结果元素数为 0 时走 Tensor::zeros 提前返回，这条路上没有 BackpropOp。
算完之后 BackpropOp::new2 决定记不记 Op::Matmul；这里两个输入都是普通张量，记下的是 None（见计算图怎么记下来）。

Storage::matmul 先 same_device、same_dtype，再对 (self, rhs) 的变体 match。这段代码和它的报错在 Storage 与后端 trait 里有完整的引用和实测，这里只看它把调用交给了谁。

CPU：调用当场算完

在 BLAS 的入口 sgemm_ 下断点，跑 intro-mm walk cpu（debug 构建，参数和过长的类型名截断，标 ...）：

* thread #1, name = 'main', queue = 'com.apple.main-thread', stop reason = breakpoint 1.1
   * frame #0: 0x0000000187142db4 libBLAS.dylib`SGEMM
     frame #1: 0x00000001000c6acc intro-mm`candle_core::accelerate::sgemm(transa='N', transb='N', m=2, n=2, k=2, alpha=1, ...)
     frame #2: 0x000000010001c9f0 intro-mm`<candle_core::cpu_backend::MatMul as candle_core::cpu_backend::utils::Map2>::f::<f32>(...)
     frame #3: 0x0000000100026088 intro-mm`<candle_core::cpu_backend::MatMul as candle_core::cpu_backend::utils::Map2>::map(...)
     frame #4: 0x0000000100021030 intro-mm`<candle_core::cpu_backend::CpuStorage as candle_core::backend::BackendStorage>::matmul(...)
     frame #5: 0x000000010009f690 intro-mm`<candle_core::storage::Storage>::matmul(self=..., bmnk=(__0 = 1, __1 = 2, __2 = 2, __3 = 2), ...)
     frame #6: 0x0000000100086e9c intro-mm`<candle_core::tensor::Tensor>::matmul(...) at tensor.rs:1530:38
     frame #7: 0x0000000100005fe4 intro-mm`intro_mm::walk(dev=(data_ptr = "cpu", length = 3)) at main.rs:26:13
     frame #8: 0x00000001000052dc intro-mm`intro_mm::main at main.rs:76:15

intro-mm-noacc 在 gemm::gemm::gemm::<f32> 下断点，栈是同样的形状，只有最上面换了：

   * frame #0: 0x00000001000a72e4 intro-mm-noacc`gemm::gemm::gemm::<f32>(m=2, n=2, k=2, dst=0x00000001011bac60, dst_cs=1, dst_rs=2, read_dst=false, ...)
     frame #1: 0x000000010001ceb0 intro-mm-noacc`<candle_core::cpu_backend::MatMul as candle_core::cpu_backend::utils::Map2>::f::<f32>(...)
     frame #2: 0x0000000100027154 intro-mm-noacc`<candle_core::cpu_backend::MatMul as candle_core::cpu_backend::utils::Map2>::map(...)
     frame #3: 0x00000001000220fc intro-mm-noacc`<candle_core::cpu_backend::CpuStorage as candle_core::backend::BackendStorage>::matmul(...)
     frame #4: 0x0000000100089ef0 intro-mm-noacc`<candle_core::storage::Storage>::matmul(...)
     frame #5: 0x000000010008f5ec intro-mm-noacc`<candle_core::tensor::Tensor>::matmul(...) at tensor.rs:1530:38

从下往上：Tensor::matmul → Storage::matmul → CpuStorage::matmul（只有一行 MatMul(bmnk).map(self, lhs_l, rhs, rhs_l)）→ Map2::map 按 dtype 取出两边的切片 → MatMul::f。f 有三份，#[cfg] 按 feature 只留一份：不开 accelerate / mkl 时调 gemm crate，开 accelerate 时经 accelerate::sgemm 调系统 libBLAS.dylib 里的 SGEMM。candle 自己的代码只有四层，全程在 main 线程上，matmul 返回时结果已经在 Vec<f32> 里。Map2 怎么把「按 dtype 分派」写一次、各份 f 对步长有什么要求，见 CPU 后端。

Metal：matmul 只编码，读数据时才提交

同一个程序跑 walk metal，在 call_mlx_gemm 下断点（debug）：

   * frame #0: 0x00000001000bfb94 intro-mm`candle_metal_kernels::kernels::mlx_gemm::call_mlx_gemm::<&candle_metal_kernels::metal::commands::CommandsGuard>
     frame #1: 0x00000001000add34 intro-mm`<candle_core::metal_backend::MetalStorage as candle_core::backend::BackendStorage>::matmul
     frame #2: 0x000000010009f890 intro-mm`<candle_core::storage::Storage>::matmul
     frame #3: 0x0000000100086e9c intro-mm`<candle_core::tensor::Tensor>::matmul
     frame #4: 0x0000000100005fe4 intro-mm`intro_mm::walk
     frame #5: 0x00000001000052dc intro-mm`intro_mm::main at main.rs:76:15

MetalStorage::matmul 做三件事：从缓冲池里要一块 b * m * n 个元素的输出缓冲区；self.device.command_encoder() 取当前命令缓冲区上的计算编码器（这会锁住 Commands 的状态）；把 dtype 换成 GemmDType 交给 call_mlx_gemm。dtype 只认 F32、F16、BF16，别的返回「mlx matmul doesn't support ...」。

call_mlx_gemm 按步长判断两边要不要转置（都不合要求就返回 MatMulNonContiguous），m 或 n 为 1 时转去 call_mlx_gemv；否则按形状、dtype 和设备型号选一组 tile 参数，拼出内核名，比如 2×2 时是 gemm_nn_f32_f32_32_32_16_2_2，1024² 时是 gemm_nn_f32_f32_64_64_16_2_2；对齐与否、有没有批维作为函数常量（function constants）传进去，和名字一起当管线缓存的键。然后设参数、算网格、dispatch，返回。栈里看到的 set_bytes::<GemmParams> 就是在往编码器里写这些参数。整个过程没有 commit，GPU 还没收到任何东西。

提交在哪里发生？在 CommandBuffer::commit 和 CommandBuffer::wait_until_completed 下断点（debug），整个 walk metal 各只停一次，而且都在 to_vec2 里面：

* thread #1, name = 'main', queue = 'com.apple.main-thread', stop reason = breakpoint 2.1
   * frame #0: 0x00000001000e836c intro-mm`<candle_metal_kernels::metal::command_buffer::CommandBuffer>::commit
     frame #1: 0x00000001000eb81c intro-mm`<candle_metal_kernels::metal::commands::Commands>::commit_swap_locked
     frame #2: 0x00000001000ebfd0 intro-mm`<candle_metal_kernels::metal::commands::Commands>::flush_and_wait_current
     frame #3: 0x00000001000c5ebc intro-mm`<candle_core::metal_backend::device::MetalDevice>::flush_and_wait_current
     frame #4: 0x00000001000a5034 intro-mm`<candle_core::metal_backend::MetalStorage>::to_cpu::<f32>
     frame #5: 0x00000001000abf2c intro-mm`<candle_core::metal_backend::MetalStorage as candle_core::backend::BackendStorage>::to_cpu_storage
     frame #6: 0x0000000100082538 intro-mm`<candle_core::tensor::Tensor>::to_vec2::<f32>
     frame #7: 0x00000001000060ac intro-mm`intro_mm::walk(dev=(data_ptr = "metal", length = 5)) at main.rs:28:13
* thread #1, name = 'main', queue = 'com.apple.main-thread', stop reason = breakpoint 3.1
   * frame #0: 0x00000001000e8278 intro-mm`<candle_metal_kernels::metal::command_buffer::CommandBuffer>::wait_until_completed
     frame #1: 0x00000001000eb6c0 intro-mm`<candle_metal_kernels::metal::commands::Commands>::ensure_completed(cb=0x000000016fdfabc0) at commands.rs:324:20
     frame #2: 0x00000001000ec0c0 intro-mm`<candle_metal_kernels::metal::commands::Commands>::flush_and_wait_current
     ...
     frame #6: 0x0000000100082538 intro-mm`<candle_core::tensor::Tensor>::to_vec2::<f32>

MetalStorage::to_cpu 先分配一块 CPU 可见的缓冲区、编码一次 blit（从结果缓冲区拷过去），然后 Commands::flush_and_wait_current：commit_swap_locked 结束当前编码器、commit 当前命令缓冲区、换上一个新的，再对刚提交的那个 waitUntilCompleted。命令队列按提交顺序执行，等到它完成，之前的 gemm 也一定完成了。

一个 Metal 程序里，命令缓冲区被提交的时机只有这几种：

时机	入口	等不等
读数据：`to_vec*`、`to_scalar`、`to_device(&Cpu)`	`to_cpu` → `flush_and_wait_current`	等这一个命令缓冲区
`device.synchronize()`	`Device::synchronize` → `MetalDevice::wait_until_completed` → `flush_and_wait`	等最后一个在途的命令缓冲区
攒够了：`Commands::command_encoder` 里计数到 `compute_per_buffer`（默认 50，环境变量 `CANDLE_METAL_COMPUTE_PER_BUFFER`）	`commit_swap_locked`	不等
`Commands` 被析构	`Drop` 里 `flush`	不等

这个计数加在每一次 command_encoder() / blit_command_encoder() 调用上，也就是大致每个 op 加一；同一个命令缓冲区里的 op 共用一个计算编码器，遇到 blit 才换。细节、栅栏（fence）和 #3595 修过的读回竞态见 Metal 后端。

📝NOTE

冷知识：Metal 的 matmul 换过三次内核

现在 MetalStorage::matmul 只有一条路：call_mlx_gemm。之前有过两条，第二条是一个预编译的库。2023-12 的提交 0404a3eb「Removed MPSMatrix entirely (buggy).」把 Metal Performance Shaders 的矩阵乘拿掉，换成一个 116184 字节的 libMetalFlashAttention.metallib，include_bytes! 进二进制，用 new_library_with_data 直接加载，不需要编译（这个提交早于本地仓库最早的 tag 0.5.1）。

#2473「Hook the MLX matmul kernels in candle-core.」（首次发布于 0.7.0）在旁边加了 MLX 的 gemm，要设环境变量 CANDLE_USE_MLX_MM 才走它；#2547「Switch to using the MLX matmul by default.」（首次发布于 0.8.0）把默认值翻过来，改成设 CANDLE_USE_MFA_MM 才走旧路；#2742「Remove the old MFA gemm kernels.」删掉了旧的那条分支和 call_gemm，两天后 #2755「Remove the MFA gemm library.」删掉了 .metallib 文件本身，两个都首次发布于 0.8.3。四个提交的说明都没有写原因。和本页有关的后果是：matmul 的内核从「随二进制带来的已编译库」变成了 mlx_gemm.metal 这 1483 行源码，进程里第一次 matmul 要在运行期编译它，下一节的 3.8 s 就是这么来的。

第一次和之后

计时：1024×1024

time 模式：Tensor::randn 建两个 1024×1024 的 f32 矩阵、synchronize，然后同一个 matmul 跑 6 次，每次记 matmul 返回的时间、synchronize 之后的时间、再读出 c[0][0] 之后的时间。release 构建：

intro-mm cpu 1024x1024 run 0: matmul returned 1.485 ms, after synchronize 1.485 ms (1446 GFLOP/s), after reading c[0][0] 1.500 ms (=32.256)
intro-mm cpu 1024x1024 run 1: matmul returned 890.8 µs, after synchronize 890.8 µs (2411 GFLOP/s), after reading c[0][0] 892.0 µs (=32.256)
intro-mm cpu 1024x1024 run 2: matmul returned 846.2 µs, after synchronize 846.3 µs (2538 GFLOP/s), after reading c[0][0] 850.8 µs (=32.256)
...
intro-mm-noacc cpu 1024x1024 run 0: matmul returned 3.787 ms, after synchronize 3.787 ms (567 GFLOP/s), after reading c[0][0] 3.790 ms (=11.596)
intro-mm-noacc cpu 1024x1024 run 1: matmul returned 4.268 ms, after synchronize 4.268 ms (503 GFLOP/s), after reading c[0][0] 4.270 ms (=11.596)
intro-mm-noacc cpu 1024x1024 run 2: matmul returned 3.228 ms, after synchronize 3.228 ms (665 GFLOP/s), after reading c[0][0] 3.229 ms (=11.596)
...
## metal, OS shader cache warm
intro-mm metal 1024x1024 run 0: matmul returned 2.321 ms, after synchronize 3.908 ms (549 GFLOP/s), after reading c[0][0] 5.138 ms (=-29.871)
intro-mm metal 1024x1024 run 1: matmul returned 18.5 µs, after synchronize 1.458 ms (1473 GFLOP/s), after reading c[0][0] 2.353 ms (=-29.871)
intro-mm metal 1024x1024 run 2: matmul returned 14.3 µs, after synchronize 1.496 ms (1436 GFLOP/s), after reading c[0][0] 2.121 ms (=-29.871)
intro-mm metal 1024x1024 run 3: matmul returned 10.0 µs, after synchronize 1.155 ms (1860 GFLOP/s), after reading c[0][0] 1.762 ms (=-29.871)
...
## metal, INTRO_NONCE (OS cache miss forced)
LIB  Random 61.1 ms
PIPE Random rand_normal_f32 0.8 ms
LIB  Gemm 3820.0 ms
PIPE Gemm gemm_nn_f32_f32_64_64_16_2_2 0.1 ms
intro-mm metal 1024x1024 run 0: matmul returned 3820.986 ms, after synchronize 3825.006 ms (1 GFLOP/s), after reading c[0][0] 3826.370 ms (=-29.871)
intro-mm metal 1024x1024 run 1: matmul returned 28.1 µs, after synchronize 1.470 ms (1461 GFLOP/s), after reading c[0][0] 2.399 ms (=-29.871)
...

CPU 上 randn 没有固定种子，两个 crate 的随机矩阵不同，c[0][0] 也就不同。

	第一次（到 synchronize）	之后每次	第一次里有什么
CPU + Accelerate	1.5 ms	0.85–0.95 ms	线程和缓存预热
CPU，`gemm` crate	3.8 ms	3.2–4.3 ms	同上
Metal，macOS 的 Metal 库缓存命中（热）	3.9 ms	1.15–1.5 ms	从系统缓存取回 gemm 库（约 2–3 ms）、建管线
Metal，缓存不命中（冷）	3.83 s	1.15–1.5 ms	`newLibraryWithSource` 把 `mlx_gemm.metal` 编译一遍，3.82 s

数字只对这台机器（M2 Max、macOS 27、rustc 1.99.0）有效。几点：

在这个尺寸上，开 accelerate 的 CPU 最快，约 2.4 TFLOP/s；gemm crate 慢约 4 倍；Metal 稳定后约 1.5 TFLOP/s，不如 Accelerate。M 系列芯片上的 Accelerate 是很强的对手，「上 GPU 一定更快」在小矩阵上不成立。
Metal 上 matmul 那一行稳定后 10–28 µs 就返回，时间花在 synchronize 里；CPU 上 matmul 返回时已经算完，synchronize 是空操作（Device::synchronize 对 Cpu 直接 Ok(())）。所以 Metal 计时要在结束处 synchronize，否则测到的只是编码的时间。
冷的时候，3.82 s 全在 matmul 这一行里，而且在调用线程上同步完成：load_library 拿着写锁调 newLibraryWithSource，编译完才往下编码。它是按源文件编译的：第一次 matmul 编的是整个 mlx_gemm.metal，之后换任何形状、任何 tile 参数都只是从这个库里取函数、建管线（PIPE 那几行都是零点几毫秒）。
冷和热由 macOS 决定，不由 candle 决定。candle 只在进程里缓存，进程一退出就没了；macOS 在用户缓存目录 $(getconf DARWIN_USER_CACHE_DIR)com.apple.metal/ 里另有一份跨进程的编译缓存，命中时同一个库只要 2–3 ms。是否命中不受 candle 控制，同一个二进制第二次运行通常是热的，但不保证。本页为了稳定地复现冷的情形，在 patch 过的拷贝里设 INTRO_NONCE，往每段 .metal 源码末尾追加一行带时间戳的注释，源码文本变了，系统缓存就对不上；设 CANDLE_METAL_ENABLE_FAST_MATH=false 改变编译选项，也能起到同样作用。这层缓存的机制和更多实测见 Metal 后端。

❗IMPORTANT

动脑筋：Metal 上算完、synchronize 也做了，再读 1024×1024 结果里的一个元素 c.get(0)?.get(0)?.to_scalar::<f32>()?，要多久？

提示：GPU 已经空闲，get 只是改 Layout 的视图。猜一个数量级：微秒，还是毫秒？

动脑筋答案

零点几毫秒到将近 1 毫秒，比 CPU 上读同一个元素（1–5 µs）慢两个数量级。上面的输出里，Metal 每一轮「after reading c[0][0]」比「after synchronize」多 0.6–0.9 ms：

intro-mm metal 1024x1024 run 1: matmul returned 18.5 µs, after synchronize 1.458 ms (1473 GFLOP/s), after reading c[0][0] 2.353 ms (=-29.871)
intro-mm metal 1024x1024 run 3: matmul returned 10.0 µs, after synchronize 1.155 ms (1860 GFLOP/s), after reading c[0][0] 1.762 ms (=-29.871)
intro-mm cpu 1024x1024 run 1: matmul returned 890.8 µs, after synchronize 890.8 µs (2411 GFLOP/s), after reading c[0][0] 892.0 µs (=32.256)

to_scalar 在 Metal 上调 storage.to_cpu_storage()，再按 layout().start_offset() 取一个数。to_cpu_storage 不看 Layout：to_cpu 按 self.count 拷整个缓冲区，也就是 100 万个 f32、4 MB。
这次拷贝要走一遍完整的 GPU 往返：分配一块共享缓冲区、编码 blit、提交一个新的命令缓冲区、waitUntilCompleted，最后 read_to_vec 在 CPU 上再拷一次到 Vec。
CPU 上 storage 本来就是 Vec，直接按偏移取，不复制，也不等谁。

时间主要花在往返上，不在 4 MB 上。intro-mm read（release，热缓存）把同一个元素用两种方式读 4 轮：一种直接读视图，一种先在 GPU 上 narrow 再 contiguous() 成一个单元素张量再读，后者只拷 4 字节，却要多编码一个拷贝内核：

run 0: view to_scalar 973.6 µs (11.699), narrow+contiguous+to_vec1 622.3 µs (11.699)
run 1: view to_scalar 833.9 µs (11.699), narrow+contiguous+to_vec1 340.6 µs (11.699)
run 2: view to_scalar 678.6 µs (11.699), narrow+contiguous+to_vec1 625.8 µs (11.699)
run 3: view to_scalar 608.5 µs (11.699), narrow+contiguous+to_vec1 611.1 µs (11.699)

两者都在 0.3–1 ms 之间，差别不稳定。读回的固定成本是一次提交加一次等待，所以生成循环每个 token 读一次 argmax 的结果，这一次往返躲不掉（见生成一个 token）。

小结

Tensor::matmul 只做形状：秩相同、k 和批维相等，批维压成一个数，(b, m, n, k) 加两个 Layout 交给 Storage::matmul。
CPU 上 candle 自己只有四层：Tensor::matmul → Storage::matmul → CpuStorage::matmul → MatMul::f，然后是 gemm crate 或 Accelerate 的 SGEMM，返回时已经算完。
Metal 上 MetalStorage::matmul 把 MLX 的 gemm 编码进当前命令缓冲区就返回；提交发生在读数据、synchronize、每攒约 50 个 op、Commands 析构这几个时刻，2×2 的例子里只有 to_vec2 那一次。
1024² f32：Accelerate 约 0.9 ms，gemm crate 约 3.8 ms，Metal 稳定后约 1.4 ms。
Metal 的第一次 matmul 要编译整个 mlx_gemm.metal：macOS 的 Metal 库缓存命中时约 2–3 ms，不命中时 3.8 s，都发生在 matmul 那一行里。
Metal 上 to_scalar / to_vec* 会把整个缓冲区拷回 CPU，不管视图有多小；读回一次的固定成本是一次提交加等待，约 0.3–1 ms。