Rust in Everything
基于 candle-core 0.11.0

布局与视图

主线位置：算子主线的第三、四步之间。Storage::matmul(&rhs, (b, m, n, k), lhs_layout, rhs_layout) 除了两块数据，还收到两个 Layout；后端要靠它们知道数据从哪开始、每一维跨几个元素。本页讲 Layout 的三个字段、哪些操作只改它、后端拿到一个不连续的布局时怎么遍历；后端本身见 CPU 后端 和 Metal 后端。

flowchart TB
  view["视图操作"] --> lay["新 Layout"]
  lay --> q{"is_contiguous"}
  q -->|是| fast["按切片整段处理"]
  q -->|否| sb["strided_blocks"]
  sb --> single["SingleBlock"]
  sb --> uni["UniformBlocks"]
  sb --> multi["MultipleBlocks"]
  multi --> idx["StridedIndex 逐块"]
  bin["二元运算"] --> nd["NdIter 合并维度"]

实验 crate：exp-candle/tensor-exp（features metal accelerate）和 exp-candle/tensor-noaccel（同样的代码，去掉 accelerate）；M2 Max、macOS 27、rustc 1.99.0，计时都是 release。

三个字段

Layout：

#[derive(Debug, PartialEq, Eq, Clone)]
pub struct Layout {
    shape: Shape,
    // The strides are given in number of elements and not in bytes.
    stride: Vec<usize>,
    start_offset: usize,
}

第 (i0, i1, ..) 个元素在 storage 里的位置是 start_offset + Σ ik * stride[k]。步长以元素计，不以字节计；步长是 usize，不能为负，所以没有「翻转」视图，flip 用 index_select 复制；广播出来的维度步长为 0。Shape 是 Vec<usize> 的包装。新张量都从 Layout::contiguous 开始，行主序（row major），最后一维步长为 1。

只改布局的操作

视图操作在 Tensor 一层的写法都一样（见 Tensor 的结构），区别全在 Layout 上对应的方法，每个都只做几行算术：

  • narrow：那一维的长度换成 len，start_offset 加上 stride[dim] * start，步长不变。
  • transpose：交换两维的长度和步长。permute 是一般情况。
  • broadcast_as：左边补出来的维度、以及原来长度为 1 被拉长的维度，步长设为 0；其余维度必须长度相等，否则报 cannot broadcast。

layout 实验对一个 (2, 3, 4) 的张量打印几个视图的布局：

                     x (2,3,4): dims=[2, 3, 4] stride=[12, 4, 1] off=0 contiguous=true -> ...
               x.narrow(2,1,2): dims=[2, 3, 2] stride=[12, 4, 1] off=1 contiguous=false -> ...
               x.narrow(1,1,2): dims=[2, 2, 4] stride=[12, 4, 1] off=4 contiguous=false -> ...
              x.transpose(1,2): dims=[2, 4, 3] stride=[12, 1, 4] off=0 contiguous=false -> ...
                    x.i(0).t(): dims=[4, 3] stride=[1, 4] off=0 contiguous=false -> ...
 x.narrow(0,1,1).narrow(1,1,1): dims=[1, 1, 4] stride=[12, 4, 1] off=16 contiguous=true -> ...
            bcast (1,4)->(3,4): dims=[3, 4] stride=[0, 1] off=0 contiguous=false -> ...

（每行 -> 之后的 StridedBlocks 下面再看。）reshape 不在这个清单里：只有输入连续时它才是视图，不连续时会复制，见 Tensor 的结构。

is_contiguous 与 contiguous()

Layout::is_contiguous 转给 Shape::is_contiguous：

    pub fn is_contiguous(&self, stride: &[usize]) -> bool {
        if self.0.len() != stride.len() {
            return false;
        }
        let mut acc = 1;
        for (&stride, &dim) in stride.iter().zip(self.0.iter()).rev() {
            if dim > 1 && stride != acc {
                return false;
            }
            acc *= dim;
        }
        true
    }

从最后一维往前，要求每一维的步长等于它右边所有维度长度的乘积；长度为 1 的维度跳过，它的步长是多少都不影响寻址。它不看 start_offset，所以上面 x.narrow(0,1,1).narrow(1,1,1) 偏移 16 仍然算连续。contiguous_offsets 在连续时返回 (start_offset, start_offset + elem_count)，后端拿它直接切出一段切片。

Tensor::contiguous 连续时原样返回 self.clone()，否则分配一块未初始化的新 storage，让源 storage 按布局拷进去：

    pub fn contiguous(&self) -> Result<Tensor> {
        if self.is_contiguous() {
            Ok(self.clone())
        } else {
            let shape = self.shape();
            let mut storage = unsafe { self.device().alloc_uninit(shape, self.dtype())? };
            self.storage()
                .copy_strided_src(&mut storage, 0, self.layout())?;
            let op = BackpropOp::new1(self, Op::Copy);
            Ok(from_storage(storage, shape.clone(), op, false))
        }
    }

force_contiguous 是不做判断、总是复制的版本。

❗IMPORTANT
动脑筋：转置一个 (4, 1)
let odd = Tensor::zeros((4, 1), DType::F32, &Device::Cpu)?.t()?;
println!("{:?} {:?} {}", odd.dims(), odd.stride(), odd.is_contiguous());

转置后形状是 (1, 4)。步长是多少？is_contiguous() 返回 true 还是 false？

后端怎么遍历

StridedIndex：逐个元素

StridedIndex 是一个迭代器，按行主序的逻辑顺序吐出每个元素在 storage 里的位置。它记着一个多维下标 multi_index，next 像里程表一样从最后一维加一、满了进位：

        for ((multi_i, max_i), stride_i) in self
            .multi_index
            .iter_mut()
            .zip(self.dims.iter())
            .zip(self.stride.iter())
            .rev()
        {
            let next_i = *multi_i + 1;
            if next_i < *max_i {
                *multi_i = next_i;
                updated = true;
                next_storage_index += stride_i;
                break;
            } else {
                next_storage_index -= *multi_i * stride_i;
                *multi_i = 0
            }
        }

位置是增量维护的：进一位加一个步长，归零时减回去，不做乘法。x.transpose(1,2) 的完整序列：

strided_index of x.transpose(1,2): [0, 4, 8, 1, 5, 9, 2, 6, 10, 3, 7, 11, 12, 16, 20, 13, 17, 21, 14, 18, 22, 15, 19, 23]

每吐一个位置都要走一遍这个循环，元素多时这是主要开销。

StridedBlocks：按连续块

多数非连续布局并不是每个元素都跳着走，最右边几维往往还是连续的。Layout::strided_blocks 从右往左数出连续的维度，把它们合成一个块，只对左边剩下的维度做下标迭代：

        for (&stride, &dim) in self.stride().iter().zip(self.dims().iter()).rev() {
            // Size-1 dimensions are trivially contiguous regardless of their stride.
            if dim == 1 {
                contiguous_dims += 1;
                continue;
            }
            if stride != block_len {
                break;
            }
            block_len *= dim;
            contiguous_dims += 1;
        }
        let index_dims = self.dims().len() - contiguous_dims;
        match index_dims {
            0 => crate::StridedBlocks::SingleBlock {
                start_offset: self.start_offset,
                len: block_len,
            },
            1 => crate::StridedBlocks::UniformBlocks {
                start_offset: self.start_offset,
                block_len,
                count: self.dims()[0],
                src_stride: self.stride[0],
            },
            _ => {
                let block_start_index = crate::StridedIndex::new(
                    &self.dims()[..index_dims],
                    &self.stride[..index_dims],
                    self.start_offset,
                );
                crate::StridedBlocks::MultipleBlocks {
                    block_start_index,
                    block_len,
                }
            }
        }

StridedBlocks 的三个变体：整个张量就是一段（SingleBlock）；只剩一维要迭代，块的起点是等差数列（UniformBlocks，起点 start_offset + i * src_stride）；要迭代两维以上，块起点交给一个 StridedIndex（MultipleBlocks）。上面那组布局对应的结果：

                     x (2,3,4): ... -> SingleBlock { start_offset: 0, len: 24 }
               x.narrow(2,1,2): ... -> MultipleBlocks { block_len: 2, first starts: [1, 5, 9, 13, 17, 21] }
               x.narrow(1,1,2): ... -> UniformBlocks { start_offset: 4, block_len: 8, count: 2, src_stride: 12 }
              x.transpose(1,2): ... -> MultipleBlocks { block_len: 1, first starts: [0, 4, 8, 1, 5, 9] }
                    x.i(0).t(): ... -> MultipleBlocks { block_len: 1, first starts: [0, 4, 8, 1, 5, 9] }
 x.narrow(0,1,1).narrow(1,1,1): ... -> SingleBlock { start_offset: 16, len: 4 }
            bcast (1,4)->(3,4): ... -> UniformBlocks { start_offset: 0, block_len: 4, count: 3, src_stride: 0 }

在中间一维上 narrow，后两维合成长 8 的块，只剩第 0 维要迭代，得到 UniformBlocks；在最后一维上 narrow，块长 2，要迭代前两维，是 MultipleBlocks。转置把步长 1 换到了中间，最后一维步长 4，一个连续维度都没有，块长为 1，退化成逐元素的 StridedIndex。

CPU 后端的一元运算（整数类型走 unary_map，浮点类型走带向量化版本的 unary_map_vec）按这三种情况各写一个循环：SingleBlock 整段处理，另两种逐块处理，块长为 1 时逐元素调用标量函数。复制（contiguous 背后的 copy_strided_src_）在 SingleBlock 上是一次 copy_from_slice，UniformBlocks 交给二维复制 copy2d_，MultipleBlocks 每块一次 copy_from_slice，块长为 1 时逐元素赋值。Metal 后端的 copy_strided_src 在 dtype 相同、最后一维步长为 1、布局是 UniformBlocks 时改用 copy2d 内核（步长恰好等于块长时是一次 blit），其余情况走通用的 strided copy 内核。

📝NOTE
冷知识：UniformBlocks 是 0.11.0 才有的

PR #3495（「Introduce new StridedBlocks type UniformBlocks」，首次进入 0.11.0）之前，strided_blocks 只返回 SingleBlock 和 MultipleBlocks 两种：只要有一维要迭代，就建一个 StridedIndex，每个块起点都走一遍上面那个里程表循环。这个 PR 加了 UniformBlocks，把「只剩一维」的情况（例如 narrow 一个中间维、把一行广播成多行）变成等差数列，CPU 的复制和一元运算、Tensor::cat、Metal 的 strided copy 都加了这一支；同时在扫描连续维度时跳过长度为 1 的维度，以前长度为 1、步长不规整的维度会让扫描提前停下。PR 一起加了 cat 和 contiguous 的基准（基准注释写的场景是「representative of transformer KV-cache growth」）。

二元运算：NdIter

二元运算要同时按两个布局走，而且两边可能一个转置、一个广播。binary_map 现在用 NdIter（PR #3579，0.11.0）：把相邻、步长允许合并的维度合成一维，每次 next 给出两边各自的起点，内层按 inner_strides 走 inner_size 个元素；在 binary_map_vec 里，两边内层步长都是 1 时交给向量化的 f_vec，左边 1、右边 0（右边是广播的标量）时交给 f_scalar_vec，其余组合逐元素算。

NdIter 的维度数组是定长的：

/// Maximum tensor rank supported by [`NdIter`].
pub const MAX_DIMS: usize = 8;

秩超过 8 的张量，合并之后仍超过 8 维就会越界。check 把一个全 2 的 8 维、9 维张量按维度反序 permute（一维也合不了）再做加法；check2 对一个连续的 9 维张量做加法。debug 下：

rank 8 permuted add: Ok(Ok([2, 2, 2, 2, 2, 2, 2, 2]))

thread 'main' (3689907) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/candle-core-0.11.0/src/nditer.rs:47:9:
rank 9 exceeds MAX_DIMS=8

check2 的连续 9 维张量在 debug 下是同一个位置、同一句 panic。release 下：

rank 8 permuted add: Ok(Ok([2, 2, 2, 2, 2, 2, 2, 2]))

thread 'main' (3690356) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/candle-core-0.11.0/src/nditer.rs:94:21:
index out of bounds: the len is 8 but the index is 8

而 check2 在 release 下打印 rank 9 contiguous add: Ok(9)。

debug 下 NdIter::new 开头的 debug_assert! 按原始秩判断，连续的 9 维张量也 panic；release 下 debug_assert! 不存在，连续张量合并成一维后能算，不能合并的在数组下标处 panic。两种情况都是 panic，不是 Err。

📝NOTE
没有蠢问题

问： matmul 吃转置张量不用先 contiguous()，为什么 exp 吃转置张量会慢好几倍？两者不都是「按步长读」吗？

答： 读法不一样。matmul 把步长整个交给矩阵乘实现：gemm crate 接收任意的行步长和列步长，Accelerate 的 BLAS 和 Metal 的 MLX 内核接收「转置」标志，它们都按块读进缓存或线程组内存，访问顺序由矩阵乘自己安排。exp 走的是 unary_map_vec：转置后块长为 1，每个元素都要经过一次 StridedIndex::next，一次读一个元素、调一次标量 exp；块内的向量化版本（开 accelerate 时是 vForce 的 vs_exp）只在块长大于 1 时用得上。

matmul 接受哪些步长

三个 matmul 实现对布局的要求不同：

实现最后两维批维
CPU，gemm crate（不开 accelerate）任意行步长、列步长ab_skip 要求批维能用一个步长描述，否则报 MatMulUnexpectedStriding
CPU，accelerate只有「行主序」和「转置」两种（BLAS 的 N / T）同上
Metal，call_mlx_gemm只有 N / T 两种，否则 MatMulNonContiguous只取倒数第三维的步长，见下

layout 实验里三种输入：转置的 w.t()；从 (16, 32) 里 narrow 出前 16 列（行步长 32、列步长 1，既不是行主序也不是转置）；一个 (2, 3, 4, 5) 的张量 permute((1, 0, 2, 3)) 之后做批量 matmul。

tensor-exp（开 accelerate）：

== Cpu
a @ w.t(): ok [8, 32]
a @ w.narrow(1,0,16)... (16 cols of 32x32): Err: MatMulUnexpectedStriding { lhs_l: Layout { shape: [8, 16], stride: [16, 1], start_offset: 0 }, rhs_l: Layout { shape: [16, 16], stride: [32, 1], start_offset: 0 }, bmnk: (1, 8, 16, 16), msg: "non-contiguous rhs" }
permuted batch matmul: Err: MatMulUnexpectedStriding { lhs_l: Layout { shape: [3, 2, 4, 5], stride: [20, 60, 5, 1], start_offset: 0 }, rhs_l: Layout { shape: [3, 2, 5, 6], stride: [60, 30, 6, 1], start_offset: 0 }, bmnk: (6, 4, 6, 5), msg: "non-contiguous lhs" }
== Metal(MetalDevice(DeviceId(1)))
a @ w.t(): ok [8, 32]
a @ w.narrow(1,0,16)... (16 cols of 32x32): Err: Metal error Invalid matmul arguments [16, 1] [32, 1] (8, 16, 16)
permuted batch matmul: ok [3, 2, 4, 6]

tensor-noaccel（CPU 走 gemm）的 CPU 部分：

== Cpu
a @ w.t(): ok [8, 32]
a @ w.narrow(1,0,16)... (16 cols of 32x32): ok [8, 16]
permuted batch matmul: Err: MatMulUnexpectedStriding { ... msg: "non-contiguous lhs" }

同一段代码，换一个 feature 就从能算变成报错。Metal 上的批量 matmul 「ok」更值得怀疑：check 把同样的计算在 CPU 上先 contiguous() 再乘，和 Metal 的结果比（release）：

metal permuted batch matmul, max abs diff vs cpu = 6.7067194

结果是错的。原因在 call_mlx_gemm 算批步长的地方：

        let a_stride = if lhs_stride.len() > 2 {
            lhs_stride[lhs_stride.len() - 3] as isize
        } else {
            (m * k) as isize
        };

所有批维被压成一个批维 b，第 i 个矩阵的起点按 i * stride[-3] 算，前面几个批维自己的步长被丢掉了。CPU 的 ab_skip 检查的正是这一点，不满足就报错；Metal 没有检查。check2 换成更常见的写法，对 (2, 4, 3, 5) 在第 1 维 narrow(1, 0, 2)（步长 [60, 15, 5, 1]），逐个批比较（debug）：

metal lhs stride [60, 15, 5, 1]
batch (0,0) max abs diff = 0.00000011920929
batch (0,1) max abs diff = 0.00000023841858
batch (1,0) max abs diff = 7.059169
batch (1,1) max abs diff = 5.521973

前两个矩阵对，后两个读到了错误的位置。在 Metal 上做四维以上的批量 matmul 前，如果输入可能来自 narrow 或 permute，要先 contiguous()。

实测：转置前后

1024×1024 f32，b.t() 不连续，btc = b.t().contiguous()；每项先跑一次记作「首次」，再预热 3 次，取 20 次平均；Metal 每次计时前后 synchronize()。

操作（ms）CPU accelerateCPU gemmMetal
a.matmul(&btc)0.8714.7570.689
a.matmul(&b.t())0.9364.7100.486
a.matmul(&b.t().contiguous())3.2087.0690.696
&a + &btc0.0730.1180.071
&a + &b.t()1.0780.9730.484
btc.exp()0.4581.5960.027
b.t().exp()2.3562.4050.352
b.t().contiguous()2.1762.1720.334
b.narrow(1, 0, 512).contiguous()0.0430.0590.037

（Metal 一列取自 tensor-exp 那次运行。）几点：

  • matmul 吃转置不吃亏：三个实现都直接按转置读，CPU 上和连续输入差别在 10% 以内，Metal 上转置输入反而更快。先 contiguous() 再乘是净损失，CPU 上多出的 2.3 ms 和单独一次转置复制（2.2 ms）相当。
  • 逐元素运算吃转置吃亏：加法慢 7 到 15 倍，exp 在 CPU 上慢 1.5 到 5 倍，在 Metal 上慢 13 倍。如果同一个转置张量要参加多次逐元素运算，先 contiguous() 一次更划算。
  • contiguous() 的代价取决于块长：转置是块长 1 的 MultipleBlocks，CPU 上 4 MB 要 2.2 ms；在最后一维 narrow 一半是 UniformBlocks，块长 512，同样复制 2 MB 只要 0.04 ms。
  • Metal 的「首次」差别很大：同一个 &a + &btc，tensor-exp 那次首次是 585.93 ms，tensor-noaccel 那次是 3.02 ms。前者是这个进程里第一次用到二元运算的内核，要现场编译；为什么第二个进程快得多，见 Metal 后端。

动脑筋答案

(4,1).t(): dims=[1, 4] stride=[1, 1] contiguous=true

(4, 1) 的步长是 [1, 1]（最后一维长度 1，步长按规则算出 1；第 0 维步长是右边长度之积 1）。转置交换两维，形状 (1, 4)、步长还是 [1, 1]。按行主序，(1, 4) 的连续步长应是 [4, 1]，第 0 维的步长对不上；但 is_contiguous 跳过长度为 1 的维度，只检查第 1 维的步长 1，所以是连续的。这也是对的：长度为 1 的维度只有下标 0，步长乘上 0，从不影响寻址。

小结

  • Layout 是形状、按元素计的 usize 步长、起始偏移；narrow 改长度和偏移，transpose / permute 换步长，broadcast_as 把步长设为 0，都不碰数据。
  • is_contiguous 只看步长、忽略长度为 1 的维度和偏移；contiguous() 不连续时分配新 storage 并按布局复制。
  • 后端按 strided_blocks 的三种结果遍历：一整段、等差的块（0.11.0 新增的 UniformBlocks）、StridedIndex 驱动的块；二元运算用 NdIter 合并维度，秩超过 8 会 panic。
  • matmul 的三个实现都直接接受转置，不用先 contiguous()；其他步长 CPU 的 gemm 能吃、accelerate 和 Metal 报错。Metal 的批量 matmul 只用倒数第三维的步长，批维来自 narrow / permute 时结果是错的。
  • 逐元素运算对转置输入慢 1.5 到 15 倍；contiguous() 的成本由块长决定。

💬 关联讨论

发起讨论
加载中...