Flex：CPU 后端
主线位置：算子主线在 CPU 上的终点。
Dispatch按DispatchTensorKind::Flex选中这一支（见 Dispatch），调Flex::float_matmul，它直接在当前线程上算完返回，没有命令队列，也没有融合。本页讲 Flex 的张量怎么存、几类 op 在哪里接进 SIMD、gemm和 rayon。
flowchart TB
op["Flex float_* 方法"] --> dt{"dtype"}
dt -->|F32| simd["macerator SIMD 快路径"]
dt -->|"F16 BF16"| conv["逐元素转 f32 再转回"]
dt -->|F64| gen["泛型循环"]
simd --> par{"元素多于阈值"}
par -->|是| ray["rayon 分块"]
par -->|否| one["单线程"]
op --> mm["matmul"]
mm --> gemm["gemm crate"]
op --> view["转置 切片 翻转 广播"]
view --> lay["只改 Layout"]
burn-flex 在 0.21 开发中进仓库（PR #4761，首次进入 v0.21.0-pre.4），之前是一个独立的仓库；0.22 开发中 burn-ndarray 被标为弃用（PR #5344，首次进入 v0.22.0-pre.3），提交说明让用户迁到 burn-flex 或某个 CubeCL 后端。全 crate 约 3.7 万行，其中卷积、注意力、FFT、池化、插值这些 ModuleOps 占了大头；本页只读存储、逐元素、归约、矩阵乘四块。依赖里没有 C 库：gemm（0.19）做矩阵乘，macerator（0.5）做 SIMD，rayon（1.12）做并行，后两者是默认开启的 feature simd、rayon。
一个张量类型装所有 dtype
Flex 是空结构体，FlexDevice 是零大小的单元结构体，Debug 打印成 Cpu。浮点、整数、布尔三种原语是同一个 FlexTensor：
pub struct FlexTensor {
/// Shared byte storage. Clone increments refcount.
data: Arc<Bytes>,
/// Layout describing shape, strides, and offset.
layout: Layout,
/// Runtime data type.
dtype: DType,
}
Bytes 是 burn-std 从 cubecl 重新导出的字节缓冲，和 bytes 卷的 bytes::Bytes 不是一个东西。元素类型在运行期由 dtype 决定，取数据时用 storage::<E>() 按 E 重新解释字节（bytemuck::cast_slice），先断言 E::dtype() 和 dtype 一致；布尔按一字节存，Bool(Native) 和 Bool(U8) 共用 u8 视图，Bool(U32) 不支持。所以每个 op 的实现都以一个 match tensor.dtype() 开头，按类型分到泛型函数上。
clone 只加一次 Arc 计数。要写的时候，storage_mut 调 Arc::make_mut：独占时原地给出可变切片，共享时先整块复制一份（写时复制）。is_unique 就是 Arc::strong_count(&self.data) == 1，它也是 TensorMetadata::can_mut 的答案。
Layout：同一块缓冲的不同读法
Layout 是形状、步长、起始偏移：
pub struct Layout {
shape: Shape,
/// Strides in elements. Negative strides enable zero-copy flip.
strides: Vec<isize>,
start_offset: usize,
}
步长是 isize，带符号。几种「视图」op 都只算一个新的 Layout，用同一个 Arc<Bytes>：
float_swap_dims/float_permute：交换步长（Layout::transpose）。float_slice：起点移到切片开头，步长乘以切片的 step（Layout::slice），负 step 也行。float_flip：起点移到该维最后一个元素，步长取负（Layout::flip）。float_expand：被广播的维步长为 0。float_reshape：只有连续时才是视图（Layout::reshape对不连续的返回None），否则FlexTensor::reshape先to_contiguous()复制。
backend-explore（burn =0.22.1，debug）在一个 3×4 的 f32 张量上直接调 Flex 的 op，打印结果的布局，same_buffer 比较 bytes().as_ptr() 和原张量是否相同：
let t = Flex::float_from_data(TensorData::new((0..12).map(|v| v as f32).collect::<Vec<_>>(), [3, 4]), &FlexDevice);
let base = t.bytes().as_ptr();
show("original", &t, base);
show("swap_dims", &Flex::float_swap_dims(t.clone(), 0, 1), base);
show("slice", &Flex::float_slice(t.clone(), &[Slice::new(1, Some(3), 1), Slice::new(0, None, 2)]), base);
show("flip", &Flex::float_flip(t.clone(), &[1]), base);
let row = Flex::float_slice(t.clone(), &[Slice::new(0, Some(1), 1)]);
show("expand", &Flex::float_expand(row, Shape::new([3, 4])), base);
show("reshape", &Flex::float_reshape(t.clone(), Shape::new([4, 3])), base);
let tt = Flex::float_swap_dims(t.clone(), 0, 1);
show("T.reshape", &Flex::float_reshape(tt, Shape::new([12])), base);
original shape=[3, 4] strides=[4, 1] offset=0 contiguous=true same_buffer=true
swap_dims shape=[4, 3] strides=[1, 4] offset=0 contiguous=false same_buffer=true
slice shape=[2, 2] strides=[4, 2] offset=4 contiguous=false same_buffer=true
flip shape=[3, 4] strides=[4, -1] offset=3 contiguous=false same_buffer=true
expand shape=[3, 4] strides=[0, 1] offset=0 contiguous=false same_buffer=true
reshape shape=[4, 3] strides=[3, 1] offset=0 contiguous=true same_buffer=true
T.reshape shape=[12] strides=[1] offset=0 contiguous=true same_buffer=false
只有最后一行复制了：转置后的布局不连续，压成一维前必须先排成连续。
缓冲（Arc<Bytes>）是一本印好的书，Layout 是一张读书说明：从第几页开始（start_offset），读完一行跳几页、一个字跳几页（strides），一共读多少（shape）。转置是把「跳行」和「跳字」两个数对调；切片是换个起始页、跳得更远；翻转是从末尾起往回跳（负步长）；广播是读完一行不翻页，接着读同一行（步长 0）。几个人拿着各自的说明读同一本书，书只有一本（Arc 计数）。谁要在书上写字，先看是不是只有自己在读，不是就复印一本再写（Arc::make_mut）。
比方到这里为止：书只能按说明读，但 Flex 的很多快路径只接受「从第 0 页起、一页不跳」的连续布局，碰到别的读法会先把内容抄成一本新书（into_contiguous），这一步并不免费。
into_contiguous 本身也分好几条路：先合并能合并的维、去掉长度为 1 的维，然后按「一维」「二维正步长（转置一类，按块复制以照顾缓存）」「二维负步长」「批量二维转置」「最内维连续」「通用」选一个。注释举的例子是 ConvNeXt 的 layer-norm 输入：一个置换过的 [N, H, W, C] 合并后变成二维 [H*W, C] 的转置，能用分块复制。
逐元素：先想办法不分配
二元浮点 op（float_add 等）都进 binary_op：先把两个操作数广播成同一形状（只改步长），再按 dtype 分：
match dtype {
DType::F32 => binary_op_f32(lhs, rhs, f32_op, simd_hint),
DType::F64 => binary_op_typed(lhs, rhs, f64_op),
DType::F16 => binary_op_typed(lhs, rhs, |a: f16, b: f16| {
f16::from_f32(f32_op(a.to_f32(), b.to_f32()))
}),
DType::BF16 => binary_op_typed(lhs, rhs, |a: bf16, b: bf16| {
bf16::from_f32(f32_op(a.to_f32(), b.to_f32()))
}),
_ => panic!("binary_op: unsupported dtype {:?}", dtype),
}
f32 的 binary_op_f32（simd feature 开着时）和通用的 binary_op_typed 按这个顺序试：
| 顺序 | 条件 | 做法 |
|---|---|---|
| 1 | lhs 独占、连续、偏移为 0，rhs 连续 | SIMD 原地写进 lhs（f32）/ 普通循环原地写（其它类型） |
| 2 | 第 1 条不成立，rhs 独占、连续、偏移为 0，lhs 连续（只有 f32） | SIMD 原地写进 rhs；减法、除法用反向内核 rsub / rdiv |
| 3 | 某一边独占，合并后的循环嵌套能从它的第 0 个元素起顺序写 | zip_apply_inplace，带广播也原地写 |
| 4 | 都不行 | 分配新缓冲：两边都连续时一个 zip 收集，否则走合并循环 / 二维带步长 / 逐元素迭代器 |
按值传递张量的意义就在这里：调用方不再持有的操作数，它的缓冲可以直接当输出（为什么 Tensor 的方法按值消耗 self，见 Tensor 与 Kind）。第 3 条前的注释还记了一笔：以前这里有手写的 SharedRow / PerRowScalar 快路径，量出来比合并后的循环嵌套更慢，因为它们先复制稠密操作数、再多一遍读改写，于是删掉了。
a 还被别处 clone 持有、b 独占，Flex::float_add(a, b) 的结果写进谁的缓冲？两边都被持有呢？let data = || TensorData::new(vec![1.0f32; 1024], [1024]);
let a = Flex::float_from_data(data(), &FlexDevice);
let b = Flex::float_from_data(data(), &FlexDevice);
let (pa, pb) = (a.bytes().as_ptr(), b.bytes().as_ptr());
let c = Flex::float_add(a, b.clone());
println!("unique a: out==a {} out==b {}", c.bytes().as_ptr() == pa, c.bytes().as_ptr() == pb);
let a = Flex::float_from_data(data(), &FlexDevice);
let pa = a.bytes().as_ptr();
let keep = a.clone();
let b = Flex::float_from_data(data(), &FlexDevice);
let pb = b.bytes().as_ptr();
let c = Flex::float_add(a, b);
println!("a shared: out==a {} out==b {}", c.bytes().as_ptr() == pa, c.bytes().as_ptr() == pb);
let b = Flex::float_from_data(data(), &FlexDevice);
let c = Flex::float_add(keep.clone(), b.clone());
println!("both shared: out==a {} out==b {}", c.bytes().as_ptr() == keep.bytes().as_ptr(), c.bytes().as_ptr() == b.bytes().as_ptr());
f16 和 bf16 没有 SIMD 快路径：闭包里每个元素先 to_f32、算、再 from_f32，走第 1、3、4 条里的普通循环。f16 归约（sum、prod）同样转成 f32 累加再转回。半精度在 Flex 上省的是内存，不是逐元素的时间；矩阵乘是例外，见下文。
SIMD：macerator
simd/portable.rs 用一个宏 define_inplace_f32_op! 生成 add_inplace_f32 这一族，每个都是三个函数：公开入口、顺序版、并行版。顺序版挂着 #[macerator::with_simd]：
#[macerator::with_simd]
#[allow(clippy::assign_op_pattern)]
fn $seq_fn<S: Simd>(a: &mut [f32], b: &[f32]) {
let lanes = S::lanes32();
let len = a.len();
let simd_len = len / lanes * lanes;
let mut i = 0;
while i < simd_len {
unsafe {
let $a = vload_unaligned(a.as_ptr().add(i));
let $b = vload_unaligned(b.as_ptr().add(i));
vstore_unaligned::<S, _>(a.as_mut_ptr().add(i), $body);
}
i += lanes;
}
// ...（尾部逐元素）
}
函数对一个抽象的 S: Simd 写一遍，with_simd 为每种指令集生成一份，运行期用 is_x86_feature_detected! / is_aarch64_feature_detected! 选：x86_64 上是 AVX-512、AVX2、SSE 三档，aarch64 上是 NEON，wasm32 上是 SIMD128，其它平台退到标量。lanes32() 是一个向量里能放几个 f32。归约用同样的写法，macerator_sum 开 8 个独立的累加器，让加法流水线不互相等待。
关掉 simd feature 时，这些函数换成 simd/scalar.rs 里的标量版本（只剩布尔运算和比较），binary_op_f32 直接落到泛型循环。
rayon：只在大张量上开
并行与否看元素个数，各处阈值不同：
| 位置 | 阈值 | 粒度 |
|---|---|---|
f32 原地二元运算（portable.rs） | 4 × 1024 × 1024 个元素 | 每块 4096 个 |
f32 求和（sum_f32_contiguous） | 4 × 1024 × 1024 个元素 | 每块 64 × 1024 个，块内 SIMD |
按维求极值（reduce.rs） | 32 × 1024 个元素 | 按行 |
矩阵乘（get_parallelism） | M·N·K ≥ 192³ | 交给 gemm 的 Parallelism::Rayon(0)（用全部线程） |
求和那里的注释给了阈值的来由：求和受内存带宽限制，小于约 16 MiB 的数据在 Apple M 系列上还在 L2 里，rayon 的任务开销不划算。极值的阈值低得多，因为每个元素的工作（一次比较加一次条件写）便宜，但按维分组后每组够大。
求和还有一个不复制的技巧：张量不连续时，如果布局恰好把缓冲里每个元素访问一次（置换、翻转过的整块张量），元素顺序不影响和，就直接对整块缓冲求和，不用按步长逐个取。
矩阵乘：gemm
matmul 检查秩和内维，按 dtype 分：f32、f64、f16 进 matmul_gemm::<T>，bf16 进 matmul_bf16。二维的 matmul_2d_strided 把两个输入的行步长、列步长原样传给 gemm::gemm：
let (lhs_row_stride, lhs_col_stride) = get_2d_strides(lhs.layout());
let (rhs_row_stride, rhs_col_stride) = get_2d_strides(rhs.layout());
let lhs_data: &[T] = lhs.storage();
let rhs_data: &[T] = rhs.storage();
let lhs_ptr = unsafe { lhs_data.as_ptr().add(lhs.layout().start_offset()) };
let rhs_ptr = unsafe { rhs_data.as_ptr().add(rhs.layout().start_offset()) };
gemm 接受任意行列步长，所以 a.matmul(b.transpose()) 里转置过的 b 不需要先复制成连续。三维以上按批维广播，再三选一：批数至少 4 且总运算量 ≥ 128³ 时在批维上用 rayon、每个矩阵单线程；否则单个矩阵够大（≥ 192³）就让 gemm 在矩阵内部并行；都不够就单线程。
bf16 是另一回事：gemm 不支持它，matmul_bf16 先 to_contiguous，把两边逐元素转成 f32，算 f32 矩阵乘，再把结果转回 bf16。整数矩阵乘不走 gemm，是朴素循环（i32 有可选的 SIMD）。
问： f16 的逐元素运算没有 SIMD、要逐个转 f32，那 f16 矩阵乘是不是也比 f32 慢？
答： 不是。矩阵乘不经过 binary_op，matmul_gemm::<f16> 直接把 f16 指针交给 gemm（Flex 开了 gemm 的 f16 feature），由 gemm 自己的半精度内核算。下一节的实测里，1024×1024 的 f16 矩阵乘中位数约 2.0 ms，f32 约 4.0 ms；bf16 要来回转换，约 4.4 ms，比 f32 还慢一点。
和 NdArray 比
backend-matmul（release）只依赖 burn =0.22.1，features std flex ndarray rayon simd；rayon 同时打开 Flex 的 rayon 和 NdArray 的 multi-threads（ndarray/rayon 加 matrixmultiply/threading），NdArray 没开任何 BLAS。每个尺寸先热身 3 次，然后计时若干次，每次计时前后 device.sync()，取中位数和最小值：
fn time_matmul(device: &Device, n: usize, dtype: DType, reps: usize) -> (f64, f64) {
let a = Tensor::<2>::random([n, n], Distribution::Default, device).cast(dtype);
let b = Tensor::<2>::random([n, n], Distribution::Default, device).cast(dtype);
for _ in 0..3 {
let _ = a.clone().matmul(b.clone());
}
device.sync().unwrap();
let mut v = Vec::with_capacity(reps);
for _ in 0..reps {
let start = Instant::now();
let c = a.clone().matmul(b.clone());
device.sync().unwrap();
v.push(start.elapsed().as_secs_f64() * 1e3);
drop(c);
}
v.sort_by(|x, y| x.partial_cmp(y).unwrap());
(v[0], v[reps / 2])
}
M2 Max、macOS 27、rustc 1.99.0，连跑两遍，第一遍：
square matmul, f32, warmup 3, median/min of reps, release
n reps flex ms (med/min) ndarray ms (med/min) ratio
64 200 0.008/0.008 0.019/0.012 2.38
128 100 0.053/0.052 0.040/0.036 0.76
256 50 0.236/0.156 0.151/0.140 0.64
512 20 0.698/0.636 0.884/0.858 1.27
1024 10 3.786/3.481 6.871/6.258 1.81
2048 5 28.763/27.584 56.595/55.854 1.97
flex only, n=1024, reps 10
F32: median 3.939 ms, min 3.769 ms
F16: median 2.071 ms, min 2.023 ms
BF16: median 4.446 ms, min 3.890 ms
F64: median 7.666 ms, min 7.046 ms
第二遍的比值是 2.64、0.91、0.86、1.48、1.63、1.90，趋势相同。ratio 是 NdArray 中位数除以 Flex 中位数，大于 1 表示 Flex 快。512 以上 Flex 快 1.3 到 2 倍；128、256 这两档 NdArray 的 matrixmultiply 持平或略快，两遍之间也晃得最厉害，不宜读出确定的结论。64 这一档 Flex 快一倍多，固定开销小。f64 约是 f32 的两倍时间，和数据量翻倍一致。
PR #4761（「Add burn-flex CPU backend」，首次进入 v0.21.0-pre.4）把独立仓库里的 burn-flex 搬进来，同时接进 burn-dispatch，用 burn-backend-tests 的 1598 个一致性测试通过 Dispatch 跑了一遍。提交说明记下了跑出来的问题：float_is_nan / float_is_inf 收下了 out_dtype: BoolDType 参数却没用，Dispatch 要 Bool(U8) 时返回的是 Bool(Native)；复查又找到约 28 处同样的问题，根子在 ops/comparison.rs 把 BoolStore::Native 写死了。今天 Flex 的比较函数都带一个 out_dtype 参数，backend.rs 里还留着 comparison_preserves_out_dtype_u8 这样的测试。同一个提交说明还写着，为了套进 Dispatch 当时的宏，Flex 先被定义成带两个幻影类型参数的 Flex<E = f32, I = i32>，「Making it truly generic over element types is a follow-up」；后来走的是相反的方向，类型参数被删掉，dtype 改由设备设置决定（见 Backend trait）。
动脑筋答案
unique a: out==a true out==b false
a shared: out==a false out==b true
both shared: out==a false out==b false
- 第一行：
a独占、连续、偏移 0，b连续，走表里第 1 条，结果就是a的缓冲。b.clone()让b共享，不影响。 - 第二行：
a被keep持有，第 1 条不成立；b独占，第 2 条把结果写进b。加法可交换，减法、除法会换成rsub/rdiv内核。 - 第三行：两边都被持有，只能分配新缓冲。
小结
- Flex 的三种原语是同一个
FlexTensor：Arc<Bytes>+Layout+ 运行期dtype；clone加计数，写时独占就原地写，否则复制。 Layout的步长带符号，转置、切片、翻转、广播都只改布局；不连续的张量 reshape 时才复制。- 逐元素运算按「写进 lhs → 写进 rhs → 带广播原地写 → 分配」的顺序找出路；f32 有 macerator SIMD，运行期选指令集，f16 / bf16 逐元素经 f32 转换。
- rayon 只在大张量上用，阈值按操作是否受内存带宽限制而不同；矩阵乘交给 gemm，直接吃步长，bf16 先转 f32。
- 实测 512 以上的 f32 矩阵乘 Flex 比 NdArray 快 1.3 到 2 倍，f16 比 f32 快一倍。