I32 求和静默回绕

主线位置：张量主线的「计算」一步，Int 张量的全局归约和前缀和。

版本：PyTorch 2.15.0（CPU）、burn 0.22.1（Flex 设备，debug 和 release 构建）

flowchart TB
  len["长度 I32 两个 15 亿"] --> pts["PyTorch sum"]
  len --> bs["burn sum"]
  len --> bc["burn cumsum"]
  pts --> ok["int64 30 亿"]
  bs --> wrap["I32 回绕成负数"]
  bc --> dbg["debug panic"]
  bc --> rel["release 回绕"]

实验 crate pita-t 的 p02（feature 同 floor-div），PyTorch 侧 p02_int_overflow.py。平台 M2 Max、macOS 27、rustc 1.99.0。panic 行里的本地路径前缀换成了 ...。

PyTorch 的写法

数据集里每篇文档的 token 数常常是 numpy 的 int32，拿来求总数和偏移量：

lengths = torch.from_numpy(np.array([1_500_000_000, 1_500_000_000], dtype=np.int32))
total = lengths.sum()
offsets = lengths.cumsum(0)

lengths dtype torch.int32
sum     3000000000 torch.int64
cumsum  [1500000000, 3000000000] torch.int64
lengths + lengths [-1294967296, -1294967296] torch.int32

逐元素加法在 int32 里回绕，但 sum 和 cumsum 的结果类型被提升成了 int64，30 亿是对的。习惯了这一点，就不会去想求和会不会溢出。

搬到 burn 的现象

Flex 设备上 Int 的默认位宽就是 I32（ndarray 上是 I64，见 burn 卷 dtype-per-device）。数据给的是 i64，建出来照样是 I32：

let data = TensorData::new(vec![1_500_000_000i64, 1_500_000_000], [2]);
let lengths = Tensor::<1, Int>::from_data(data.clone(), &device);
let total = lengths.clone().sum();
// 另一次运行：
let offsets = lengths.cumsum(0);

## debug sum
lengths dtype I32 values [1500000000, 1500000000]
sum     I32 [-1294967296]
## debug cumsum
lengths dtype I32 values [1500000000, 1500000000]

thread 'main' (8762847) panicked at .../burn-flex-0.22.1/src/ops/cumulative.rs:19:54:
attempt to add with overflow
## release sum
sum     I32 [-1294967296]
## release cumsum
cumsum  I32 [1500000000, -1294967296]

（release 两段的 lengths dtype 行和 debug 相同，略去。）

• sum 在 debug 和 release 下都回绕成 -1294967296，类型还是 I32，没有任何提示。
• cumsum 在 debug 下 panic，在 release 下回绕。只在 debug 下跑过测试的代码，到了 release 构建会悄悄给出负的偏移量。

原因

两件事叠在一起。

第一，建张量时 dtype 由设备决定，不看数据的元素类型：from_data 不写 dtype 时取设备默认值，Flex 上是 I32（dtype 与设备「numpy 来的 f64」一节对浮点讲过同一件事）。值本身放得下 i32，转换就不报错；放不下才报错，见下面的「怎么检查」。

第二，归约的输出类型等于输入类型。Tensor::sum 返回 Tensor<1, K>，dtype 跟着输入。Flex 的整数求和其实在 i64 里累加，widening_scalar_reduce 最后一步截回原类型：

            // Truncate back to target type (wrapping, matches PyTorch)
            let result: E = truncate_i64_to_pod(total);

注释说「matches PyTorch」，上面 PyTorch 的输出说明并不是：PyTorch 提升了输出类型，没有截回。cumsum 用的是普通的 acc + val，跟着 Rust 的溢出检查走，所以 debug panic、release 回绕。同一个后端里 sum 和 cumsum 对溢出的处理不一样。各后端、各构建的完整对照表在整数运算的「溢出」一节。这两处在 origin/main（2026-10-28）上没有改动（git diff 这两个文件为空）。

改法

可能超过 2³¹ 的量，建张量时就写成 I64，或者归约前 cast：

let wide = lengths.cast(IntDType::I64);
println!("cast I64: sum {:?}, cumsum {:?}", iv(wide.clone().sum()), iv(wide.cumsum(0)));
let born = Tensor::<1, Int>::from_data(data, (&device, DType::I64));

cast I64: sum [3000000000], cumsum [1500000000, 3000000000]
created as I64: I64 sum [3000000000]

debug 和 release 输出相同，和 PyTorch 的 3000000000、[1500000000, 3000000000] 一致。cast 之前的 I32 张量里每个元素本身没有溢出，所以先 cast 再求和是安全的；如果逐元素运算已经回绕（比如上面的 lengths + lengths），cast 也救不回来，要在那之前就换成 I64。

怎么检查

• 建张量时，单个值超出 i32 会当场 panic，这一步是有检查的：
  Copy

  thread 'main' (8778823) panicked at .../burn-std-0.22.1/src/element/cast.rs:218:1:
  Element cannot be represented in the target type: "i64"(3000000000) => "i32"
  

  （from_data(vec![3_000_000_000i64], &device)，debug 和 release 相同。）检查只在建张量这一刻有；之后的加法、求和、前缀和都不检查。所以不能因为「数据装进去没报错」就认为结果不会溢出。

• 对计数、长度、偏移量、像素和这一类量，代码里显式写 DType::I64，不要依赖设备默认值。

• 测试要跑 release 构建：cumsum 的溢出在 debug 下会 panic，在 release 下只是数错了。

• 怀疑哪里溢出时，用 cast(IntDType::I64) 再算一遍，和 I32 的结果比。

小结

• PyTorch 把 int32 的 sum / cumsum 提升成 int64；burn 的输出类型等于输入类型，Flex 上 Int 默认又是 I32。
• I32 的 sum 静默回绕；cumsum 在 debug 下 panic，release 下回绕。
• 建张量时超出范围会报错，之后的运算不查。
• 大的计数一开始就用 I64，或者归约前 cast(IntDType::I64)。