Rust in Everything
基于 candle 0.11.0

整数运算

主线位置：张量主线的「计算」一步，只看整数张量：除法、和浮点标量的运算、溢出、除以 0，以及什么 dtype 能当下标。
版本：PyTorch 2.15.0（CPU、MPS）、candle 0.11.0（CPU，accelerate；Metal）
flowchart TB
  in["I64 张量乘 0.5"] --> pt["PyTorch"]
  in --> cpu["candle CPU"]
  in --> mtl["candle Metal"]
  pt --> p1["提升成 float32"]
  p1 --> p2["得 -3.5"]
  cpu --> c1["0.5 先转 i64 成 0"]
  c1 --> c2["得 0"]
  mtl --> m1["转 f32 再 fma"]
  m1 --> m2["截断得 -3"]

实验代码是一个 candle 程序 intops（参数选 CPU 或 Metal）和一个 PyTorch 脚本 intops.py（参数选 cpu 或 mps）。candle 侧 debug 和 release 各跑一遍，下文只在两者不同时分开写。平台 M2 Max、macOS 27、rustc 1.99.0。

除法：只有截断

四种符号组合各取一个：

a = torch.tensor([-7, -7, 7, 7])
b = torch.tensor([2, -2, 2, -2])
print(a / b, a // b, torch.div(a, b, rounding_mode="trunc"), a % b)

let a = Tensor::new(&[-7i64, -7, 7, 7], &dev)?;
let b = Tensor::new(&[2i64, -2, 2, -2], &dev)?;
let q = (&a / &b)?;

# PyTorch（CPU 和 MPS 相同）
a / b                  float32  [-3.5, 3.5, 3.5, -3.5]
a // b                 int64    [-4, 3, 3, -4]
div trunc              int64    [-3, 3, 3, -3]
a % b                  int64    [1, -1, 1, -1]
# candle（CPU 和 Metal 相同）
a / b                  I64      [-3, 3, 3, -3]

两点不同：

  • PyTorch 对整数张量写 / 是真除法，结果是 float32；candle 的 / 还是 I64，是整数除法。
  • candle 的整数除法向零截断，对应 torch.div(rounding_mode="trunc")，不是 //。

两个后端都直接用了语言自带的整数除法。CPU 上，Tensor::div 由 binary_op! 宏生成，后端按 dtype 调用 BinaryOpT 里对应的函数；这些函数由 op.rs 的 bin_op! 宏生成，Div 那一行是 bin_op!(Div, div, |v1, v2| v1 / v2)，同一个闭包展开成 u8、u32、i64 等十个版本。Rust 的整数 / 向零截断。Metal 上，MetalStorage::binary_impl 把 B::KERNEL（"bdiv"）拼上 dtype 得到内核名 bdiv_i64，内核在 binary.metal 里是 define_binary_op(bdiv, x / y)。Metal 着色语言的整数除法同样向零截断，所以两边结果一致。减法同理：CPU 是 |v1, v2| v1 - v2，Metal 是 define_binary_op(bsub, x - y)。

candle 没有 //，也没有取余：candle-core/src 里没有给 Tensor 实现 Rem，也没有 remainder / fmod 方法。要 PyTorch 的 //，用截断的商修正一步，只用整数运算：

fn floor_div(a: &Tensor, b: &Tensor) -> Result<Tensor> {
  // truncating quotient, then subtract 1 where the remainder and divisor have opposite signs
  let q = (a / b)?;
  let r = (a - (&q * b)?)?;
  let fix = (&r * b)?.lt(0i64)?.to_dtype(DType::I64)?;
  &q - &fix
}

floor_div(a, b)        I64      [-4, 3, 3, -4]

r = a - q·b 是截断商对应的余数，它和除数异号（r·b < 0）时，向下取整的商比截断的商小 1。lt 返回 U8 掩码，转成 I64 再减。CPU 和 Metal 输出相同。两边各输出 floor_div 和 trunc_div 两个字段，用 compare.py 对比，容差 0：

floor_div                n=  4  max|diff|=0.00e+00
trunc_div                n=  4  max|diff|=0.00e+00
worst=0.00e+00 tol=0e+00 OK

整数运算没有舍入，误差正好是 0。反向验证：把 candle 侧的 floor_div 换成截断的 a / b，compare.py 报 floor_div ... max|diff|=1.00e+00、MISMATCH，退出码 1。

burn 的 Int 除法同样向零截断，但它有一个跟除数同号的取余，两者不配套，见 burn 卷的整数运算；candle 干脆没有取余。

整数张量和浮点标量

PyTorch 里 a * 0.5 会把整数张量提升成 float32。candle 不做类型提升（见创建与 dtype），结果还是 I64，但 CPU 和 Metal 算出来的数不一样：

print(a * 0.5, a * 2.5, a + 0.5)

show("a * 0.5", &a * 0.5);
show("a * 2.5", &a * 2.5);
show("a + 0.5", &a + 0.5);

# PyTorch（CPU 和 MPS 相同）
a * 0.5                float32  [-3.5, -3.5, 3.5, 3.5]
a * 2.5                float32  [-17.5, -17.5, 17.5, 17.5]
a + 0.5                float32  [-6.5, -6.5, 7.5, 7.5]
# candle CPU
a * 0.5                I64      [0, 0, 0, 0]
a * 2.5                I64      [-14, -14, 14, 14]
a + 0.5                I64      [-7, -7, 7, 7]
# candle Metal
a * 0.5                I64      [-3, -3, 3, 3]
a * 2.5                I64      [-17, -17, 17, 17]
a + 0.5                I64      [-6, -6, 7, 7]

a / 2.0 和 a * 0.5 的结果相同，两个后端都一样。CPU 上乘 0.5 得 0、乘 2.5 等于乘 2、加 0.5 什么也没做；Metal 上是先算出浮点结果再向零截断。

原因在 Tensor 和 f64 的运算符。bin_trait! 宏给 Tensor op f64 生成的实现全部转成 affine(mul, add)：

        impl std::ops::$trait<f64> for Tensor {
            type Output = Result<Tensor>;

            fn $fn1(self, rhs: f64) -> Self::Output {
                self.affine($mul(rhs), $add(rhs))
            }
        }
// ...
bin_trait!(Add, add, |_| 1., |v| v);
bin_trait!(Sub, sub, |_| 1., |v: f64| -v);
bin_trait!(Mul, mul, |v| v, |_| 0.);
bin_trait!(Div, div, |v| 1. / v, |_| 0.);

a * 0.5 是 affine(0.5, 0.)，a / 2.0 是 affine(1. / 2., 0.)，a - 1.0 是 affine(1., -1.)。两个后端对 affine 的实现不同：

  • CPU 的 Affine 先把两个系数转成张量的元素类型，再在元素类型里算：

    impl Map1 for Affine {
    fn f<T: WithDType>(&self, vs: &[T], layout: &Layout) -> Result<Vec<T>> {
        let mul = T::from_f64(self.0);
        let add = T::from_f64(self.1);
        Ok(unary_map(vs, layout, |v| v * mul + add))
    }
}

    i64 的 from_f64 是 with_dtype! 生成的 |v: f64| v as i64。0.5 as i64 是 0，2.5 as i64 是 2，于是乘 0.5 变成乘 0，加 0.5 变成加 0。

  • Metal 的 MetalStorage::affine 把系数以 mul as f32、add as f32 传给内核，affine_kernel 在 float 里算完再转回去：

        for (uint i = tid; i < dim; i += step) {
        output[i] = static_cast<T>(fma(float(input[i]), mul, add));
    }

    -7 × 0.5 = -3.5，转成整数向零截断得 -3。

Metal 这条路还有一个代价：经过 float，超过 2²⁴ 的整数会丢精度。

let big = Tensor::new(&[(1i64 << 24) + 1], &dev)?;
show("i64 2^24+1 + 1.0", &big + 1.0);
show("i64 2^24+1 + ones", &big + &big.ones_like()?);

# PyTorch：+ 1.0 提升成 float32 丢了精度，+ 1 是整数加法
i64 2^24+1 + 1.0       float32  [16777216.0]
i64 2^24+1 + 1         int64    [16777218]
# candle CPU（release）
i64 2^24+1 + 1.0       I64      [16777218]
i64 2^24+1 + ones      I64      [16777218]
# candle Metal
i64 2^24+1 + 1.0       I64      [16777216]
i64 2^24+1 + ones      I64      [16777218]

两个整数张量相加走的是 badd_i64，不经过浮点，结果正确。

main 上 CPU 的 Affine、with_dtype! 和 affine.metal 都没有改。

📝NOTE
冷知识：Metal 上 I64 的 affine 原来直接报错

PR #3116「Fix beit on metal by adding additional affine implementations」（首次发布于 0.9.2-alpha.1，第一个正式版是 0.10.0）之前，MetalStorage::affine 的连续分支只认 F32、F16、BF16、U8、U32，I64 落进 dtype => crate::bail!("Metal contiguous affine {dtype:?} not implemented")；跨步分支连 U8、U32 都没有。这个 PR 在两个分支里补上 affine_i64（跨步分支同时补上 U8、U32），在 affine.metal 里加了一行 AFFINE(affine_i64, int64_t)。当时的 AFFINE 宏已经是 output[id] = T(fma(float(input[id]), mul, add));，所以 I64 一接进来就是「转成 float、算完截断」的语义，和 CPU 上「系数先转成整数」不同。提交说明只有标题，没有讨论两个后端的差别。U8、U32 的连续版本更早，来自 PR #2376「Enable the affine kernel for u8/u32.」（首次发布于 0.7.0）。

要可移植的结果，就别让整数张量和浮点标量直接运算：想要浮点结果，先 a.to_dtype(DType::F32)? 再乘；想要整数结果，用整数张量做运算数（&a + &a.ones_like()?），或者整数除法 &a / &two。

❗IMPORTANT
动脑筋：U32 减 1.0
let u = Tensor::new(&[3u32, 5], &Device::Cpu)?;
let r = (&u - 1.0)?;

r 是多少？在 Metal 上呢？

溢出：debug 下 panic，release 下回绕

u8 = torch.tensor([200, 100], dtype=torch.uint8)
print(u8 + torch.tensor([100, 0], dtype=torch.uint8), u8.sum())
print(torch.tensor([3, 5], dtype=torch.uint32) - torch.tensor([5, 3], dtype=torch.uint32))

("u32 [3,5] - [5,3]", Box::new(|| &Tensor::new(&[3u32, 5], &dev)? - &Tensor::new(&[5u32, 3], &dev)?)),
("u8 [200,100] + [100,0]", Box::new(|| &Tensor::new(&[200u8, 100], &dev)? + &Tensor::new(&[100u8, 0], &dev)?)),
("u8 sum_all", Box::new(|| Tensor::new(&[200u8, 100], &dev)?.sum_all())),
("u32 sum 2^31+2^31", Box::new(|| Tensor::new(&[1u32 << 31, 1 << 31], &dev)?.sum_all())),
("i64 MAX + 1 (tensor)", Box::new(|| &Tensor::new(&[i64::MAX], &dev)? + &Tensor::new(&[1i64], &dev)?)),
// 每一项包在 catch_unwind 里，panic 由 hook 打印

# PyTorch CPU
u32 [3,5] - [5,3]      ERR NotImplementedError: "add_stub" not implemented for 'UInt32'
u8 [200,100] + [100,0] uint8    [44, 100]
u8 sum                 int64    [300]
u32 sum 2^31+2^31      int64    [4294967296]
# candle CPU，debug
                       PANIC attempt to subtract with overflow (op.rs:457)
u32 [3,5] - [5,3]      (panicked, see line above)
                       PANIC attempt to add with overflow (op.rs:456)
u8 [200,100] + [100,0] (panicked, see line above)
                       PANIC attempt to add with overflow (kernels.rs:44)
u8 sum_all             (panicked, see line above)
                       PANIC attempt to add with overflow (kernels.rs:44)
u32 sum 2^31+2^31      (panicked, see line above)
                       PANIC attempt to add with overflow (op.rs:456)
i64 MAX + 1 (tensor)   (panicked, see line above)
# candle CPU release；Metal 的 debug 和 release 也是这几行
u32 [3,5] - [5,3]      U32      [4294967294, 2]
u8 [200,100] + [100,0] U8       [44, 100]
u8 sum_all             U8       [44]
u32 sum 2^31+2^31      U32      [0]
i64 MAX + 1 (tensor)   I64      [-9223372036854775808]

PyTorch 的 MPS 上，uint32 的减法和求和都报 Failed to create function state object for: ...，uint8 那两行和 CPU 相同；int64 的 MAX + 1 在 CPU 和 MPS 上都回绕成 MIN。

几点：

  • PyTorch 的 uint32 不能做减法，candle 的 U32 可以，结果回绕。从 PyTorch 过来，没有可以对照的习惯，只能自己记住 3 - 5 是 4294967294。

  • 整数求和不提升。PyTorch 对 uint8、uint32 求和返回 int64，结果正确；candle 的 sum_all 保持元素类型，200 + 100 回绕成 44。U8 正是比较运算（ge、eq）返回的掩码类型，对一个长度超过 255 的掩码直接 sum_all 来数个数，release 下会静默得到错的数。先转成更宽的类型再求和，CPU 和 Metal 上都得到 300：

    u8 to I64 then sum     I64      [300]

  • debug 和 release 不一样。CPU 的逐元素运算就是 bin_op! 里的 v1 + v2、v1 - v2，求和是 vec_reduce_sum 里的 *res += *xs.add(i)，都是普通的 Rust 整数运算。Cargo 的 dev 配置对依赖也打开溢出检查，所以 debug 构建 panic，release 构建回绕。Metal 内核没有溢出检查，两种构建都回绕。同一份代码，开发时 panic，上线后安静地出错。

除以 0 和 MIN / -1

torch.tensor([7, -7]) // torch.tensor([0, 0])
torch.div(torch.tensor([7, -7]), torch.tensor([0, 0]), rounding_mode="trunc")

let z = Tensor::new(&[0i64, 0], &dev)?;
show("i64 [7,-7] / [0,0]", &Tensor::new(&[7i64, -7], &dev)? / &z);
let zu = Tensor::new(&[0u32, 0], &dev)?;
show("u32 [7,9] / [0,0]", &Tensor::new(&[7u32, 9], &dev)? / &zu);

# PyTorch CPU
i64 [7,-7] / 0 (int)   ERR RuntimeError: ZeroDivisionError
i64 [7,-7] div trunc 0 ERR RuntimeError: ZeroDivisionError
# PyTorch MPS
i64 [7,-7] / 0 (int)   int64    [0, -1]
i64 [7,-7] div trunc 0 int64    [0, 0]
# candle CPU（debug 和 release 相同）
thread 'main' (8946082) panicked at .../candle-core-0.11.0/src/op.rs:459:28:
attempt to divide by zero
# candle Metal（debug 和 release 相同）
i64 [7,-7] / [0,0]     I64      [0, 0]
u32 [7,9] / [0,0]      U32      [4294967295, 4294967295]

（panic 行里的本地路径前缀换成了 ...。）CPU 上的 panic 来自 Rust 整数除法自己的检查，release 下也在，消息里没有张量、没有算子名，也不是 Err，? 接不住。Metal 不报错，给出 0 或全 1 的位模式。PyTorch 自己在 CPU 上报错、在 MPS 上同样静默，搬到 candle 后两种设备都要自己先处理掉 0 除数。

i64::MIN / -1：PyTorch 的 // 和 trunc 在 CPU 和 MPS 上都得到 MIN；candle Metal 也是 MIN；candle CPU 在 debug 和 release 下都 panic「attempt to divide with overflow」（同样在 op.rs:459）。

下标的 dtype

index_select、gather、embedding 要整数下标。两边接受的类型不同：

row = torch.tensor([10.0, 20.0, 30.0])
for dt in [torch.uint8, torch.int32, torch.int64, torch.int16]:
  row.index_select(0, torch.tensor([2, 0], dtype=dt))
row[torch.tensor([1, 0, 1], dtype=torch.uint8)]

let ids = Tensor::new(&[2i64, 0], &dev)?;
for dt in [DType::U8, DType::U32, DType::I64, DType::I32, DType::I16, DType::F32] {
  let r = ids.to_dtype(dt).and_then(|i| row.index_select(&i, 0));
  show(&format!("index_select {dt:?}"), r);
}
show("index_select u8 [1,0,1]", row.index_select(&Tensor::new(&[1u8, 0, 1], &dev)?, 0));

# PyTorch（CPU 和 MPS 相同）
index_select uint8     ERR RuntimeError: index_select(): Expected dtype int32 or int64 for index
index_select int32     float32  [30.0, 10.0]
index_select int64     float32  [30.0, 10.0]
index_select int16     ERR RuntimeError: index_select(): Expected dtype int32 or int64 for index
row[uint8 [1,0,1]]     float32  [10.0, 30.0]
# candle CPU
index_select U8        F32      [30, 10]
index_select U32       F32      [30, 10]
index_select I64       F32      [30, 10]
index_select I32       ERR unsupported dtype F32 for op index-select
index_select I16       ERR unsupported dtype F32 for op index-select
index_select F32       ERR unsupported dtype F32 for op index-select
index_select u8 [1,0,1] F32      [20, 10, 20]
# candle Metal（前三行和最后一行同 CPU）
index_select I32       ERR Metal contiguous to_dtype I64 I32 not implemented
index_select I16       ERR Metal contiguous to_dtype I64 I16 not implemented
index_select F32       ERR Metal contiguous index_select F32 F32 not implemented

（PyTorch 对 uint8 下标还打印一条 UserWarning: indexing with dtype torch.uint8 is now deprecated。）

  • PyTorch 的下标是 int32 / int64；candle 是 U8、U32、I64，I32 不行。

  • PyTorch 里 row[uint8 张量] 是布尔掩码（选出第 0、2 个）；candle 的 U8 下标就是下标，[1, 0, 1] 取出第 1、0、1 个元素。从 PyTorch 搬掩码选择时别直接换成 index_select。

  • CPU 的报错说「unsupported dtype F32」，但 F32 是被索引的 row 的类型，不是下标的类型。CpuStorage::index_select 按下标的变体分派，落进 _ 分支时报的是 self.dtype()：

            match ids {
            Self::U8(ids) => IndexSelect { ids, ids_l, dim }.map(self, l),
            Self::U32(ids) => IndexSelect { ids, ids_l, dim }.map(self, l),
            Self::I64(ids) => IndexSelect { ids, ids_l, dim }.map(self, l),
            _ => Err(Error::UnsupportedDTypeForOp(self.dtype(), "index-select").bt()),
        }
    读到这条报错，要去查的是下标的 dtype。

  • Metal 上连 I32 都造不出来：cast.metal 只给 f32、f16、bf16、i64、u32、u8 生成转换内核。

argmax 返回 U32（见归约与索引），可以直接当下标用，但不能和 I64 张量做运算：

argmax_keepdim         U32      [2]
argmax + i64           ERR dtype mismatch in add, lhs: U32, rhs: I64

浮点转整数两边都向零截断：[-1.5, 2.7] 转 int64 / I64 都是 [-1, 2]；candle 转 U32 时负数变成 0（[0, 2]），CPU 和 Metal 相同。

动脑筋答案

# candle CPU
u32 [3,5] - 1.0        U32      [3, 5]
# candle Metal
u32 [3,5] - 1.0        U32      [2, 4]

CPU 上 r 还是 [3, 5]，减法什么也没做。&u - 1.0 是 affine(1., -1.)，CPU 的 Affine 先算 u32::from_f64(-1.0)，即 -1.0 as u32；Rust 的浮点转整数是饱和转换，负数变成 0，于是变成「乘 1 加 0」。&u + (-1.0) 同样得到 [3, 5]，&u - 4.0 也是。Metal 在 float 里算 3 - 1、5 - 1，得到 [2, 4]；&u - 4.0 在 Metal 上得 [0, 1]（3 - 4 = -1.0 转 U32 成了 0）。PyTorch 的 u - 1.0 提升成 float32，得 [2.0, 4.0]。

小结

  • candle 的整数 / 向零截断，结果还是整数，CPU 和 Metal 一致；没有 // 和 %，向下取整的除法用 q - (r·b < 0) 自己凑，和 PyTorch 误差为 0。
  • 整数张量和浮点标量运算走 affine：CPU 先把系数转成整数（a * 0.5 得 0，u32 - 1.0 不变），Metal 经 f32 计算再截断（a * 0.5 得 -3，超过 2²⁴ 丢精度）。要么先转浮点，要么用整数张量做运算数。
  • 整数溢出在 CPU debug 下 panic，在 CPU release 和 Metal 上回绕；U8 / U32 求和不提升，数掩码之前先转 U32 或 I64。
  • 整数除以 0：CPU panic（不是 Err），Metal 静默给出 0 或 0xFFFFFFFF。
  • 下标只收 U8、U32、I64；U8 是下标不是掩码；CPU 下标类型不对时报错里写的是被索引张量的 dtype。

💬 关联讨论

发起讨论
加载中...