自定义内核越界读不报错，也不截断

主线位置：算子主线的最底层，cubecl 把 #[cube] 函数编译成 MSL / WGSL 的那一步。主要影响自己写 #[cube(launch)] 内核的人；burn 的融合内核用 launch_unchecked（见生成融合内核），但 burn-cubecl 里也有一些内核用 launch，最后一节说。内核的写法和展开见 #[cube] 内核。

flowchart TB
  launch["launch 即 Checked"] --> ir["x 下标带 checked 标记"]
  ir --> pass["CheckedIo pass"]
  pass --> m{"base 是 RuntimeArray 吗"}
  m -->|"实际是 Pointer"| skip["不改写"]
  skip --> same["和 launch_unchecked 同一份代码"]
  m -.->|设计中| clamp["min 下标 len - 1"]

cubecl 0.11.3（经 burn 0.22.1 的 cubecl feature 引入）。实验 crate pc-gpu（开了 metal 在内的整组 feature，cubecl::Device::default() 走 wgpu<msl>）和 pc-oob-wgsl（只开 wgpu，走 wgpu<wgsl>），release 构建，M2 Max、macOS 27、rustc 1.99.0。CUDA、ROCm、Vulkan（SPIR-V）只读了源码。

现象

一个不带守卫的逐元素加法，launch 和 launch_unchecked 都生成；另有一个带守卫的版本做对照：

#[cube(launch, launch_unchecked)]
fn add<F: Float>(x: &[F], y: &[F], out: &mut [F]) {
  out[ABSOLUTE_POS] = x[ABSOLUTE_POS] + y[ABSOLUTE_POS];
}

#[cube(launch)]
fn add_guarded<F: Float>(x: &[F], y: &[F], out: &mut [F]) {
  if ABSOLUTE_POS < out.len() {
    out[ABSOLUTE_POS] = x[ABSOLUTE_POS] + y[ABSOLUTE_POS];
  }
}

要看出「夹取」和「越界」的区别，越界的下标必须落在内容可辨认的内存上。做法是上传 8 个 f32，再用 Handle::offset_end(8) 让句柄只露出前 6 个：

const X: [f32; 8] = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 100.0, 100.0];
const Y: [f32; 8] = [10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 1000.0, 1000.0];

/// A buffer of 8 f32 whose handle only shows the first 6 (`offset_end` of 8 bytes).
fn six_of_eight(handle: cubecl::server::Handle) -> BufferArg {
  // SAFETY: the handle owns 8 f32; the view covers the first 6 of them.
  unsafe { BufferArg::from_raw_parts(handle.offset_end(8), 6) }
}

内核里的 len() 不看 from_raw_parts 的第二个参数，而是启动时按句柄的可见字节数算（process_buffer 里 array.handle.size_in_used() / elem_size），所以先用一个小内核把 x.len() 写出来确认。然后启动 8 个单元：

读：x、y 只露 6 个，out 8 个。如果 launch 把下标夹到 len - 1，单元 6、7 读到的是 x[5] + y[5] = 66；直接越界读到的是后面的 100 + 1000 = 1100。
写：x、y 8 个，out 只露 6 个、初值全是 -1。夹取的话单元 6、7 写到 out[5]，out[6]、out[7] 保持 -1。

MSL（pc-gpu oob）：

device = Wgpu(WgpuDevice { kind: DefaultDevice, backend: Auto }), runtime = "wgpu<msl>"
x.len() seen by the kernel (8 f32, offset_end 8 bytes): 6
launch           read  (x, y show 6, out 8): [11.0, 22.0, 33.0, 44.0, 55.0, 66.0, 1100.0, 1100.0]
launch_unchecked read  (x, y show 6, out 8): [11.0, 22.0, 33.0, 44.0, 55.0, 66.0, 1100.0, 1100.0]
add_guarded      read  (x, y show 6, out 8): [11.0, 22.0, 33.0, 44.0, 55.0, 66.0, 1100.0, 1100.0]
launch           write (x, y 8, out shows 6): [11.0, 22.0, 33.0, 44.0, 55.0, 66.0, 1100.0, 1100.0]
launch_unchecked write (x, y 8, out shows 6): [11.0, 22.0, 33.0, 44.0, 55.0, 66.0, 1100.0, 1100.0]
add_guarded      write (x, y 8, out shows 6): [11.0, 22.0, 33.0, 44.0, 55.0, 66.0, -1.0, -1.0]

WGSL（pc-oob-wgsl，同一份代码）：

device = Wgpu(WgpuDevice { kind: DefaultDevice, backend: Auto }), runtime = "wgpu<wgsl>"
x.len() seen by the kernel (8 f32, offset_end 8 bytes): 6
launch           read  (x, y show 6, out 8): [11.0, 22.0, 33.0, 44.0, 55.0, 66.0, 1100.0, 1100.0]
launch_unchecked read  (x, y show 6, out 8): [11.0, 22.0, 33.0, 44.0, 55.0, 66.0, 1100.0, 1100.0]
add_guarded      read  (x, y show 6, out 8): [11.0, 22.0, 33.0, 44.0, 55.0, 66.0, 1100.0, 1100.0]
launch           write (x, y 8, out shows 6): [11.0, 22.0, 33.0, 44.0, 55.0, 66.0, 1100.0, 1100.0]
launch_unchecked write (x, y 8, out shows 6): [11.0, 22.0, 33.0, 44.0, 55.0, 66.0, 1100.0, 1100.0]
add_guarded      write (x, y 8, out shows 6): [11.0, 22.0, 33.0, 44.0, 55.0, 66.0, -1.0, -1.0]

内核看到的 x.len() 是 6，长度信息是对的。
launch 和 launch_unchecked 在两条路径上结果完全一样：读到了 1100（不是 66），写到了 out[6]、out[7]（不是保持 -1）。
add_guarded 的守卫只判断 out，所以「读」那一行照样越界读了 x、y；「写」那一行守卫挡住了 out[6]、out[7]。

开 CUBECL_DEBUG_LOG=stdout 看生成的代码，add 在每条路径上各编译了 2 次（KernelId 里 mode: Checked 和 mode: Unchecked 各一个）。MSL 两份逐行比较（diff），除了入口名的哈希（add_f_f32_82718798 / add_f_f32_5b505577）和四个内建变量的编号顺序，完全相同；访存部分是：

device float const* const v17 = &v1[v60];
float const v18 = *v17;
device float const* const v23 = &v5[v60];
float const v24 = *v23;
float const v25 = v18 + v24;
device float * const v30 = &v9[v60];
*v30 = v25;

没有 min。info_st 里的 static_meta[4] 放着各缓冲区的长度（add_guarded 的 MSL 里 out.len() 就是读 info_v35.static_meta[2]），add 的两份都没有读它。WGSL 两份的差别只有入口名（add_f_f32_9f1b4126 / add_f_f32_3251ecf7），访存同样是裸下标：

let v17 = &(*v136)[v57];
let v18 = *v17;
let v23 = &(*v135)[v57];
let v24 = *v23;
let v25 = v18 + v24;
let v30 = &(*v134)[v57];
*v30 = v25;

出处

按设计，launch 和 launch_unchecked 的区别只在执行模式：宏生成 launch 时调 launch_body(ExecutionMode::Checked)，ExecutionMode 的注释写着「Checked kernels are safe.」。MSL（cubecl-cpp）、WGSL（cubecl-wgpu）和 SPIR-V（cubecl-spirv）三个编译器都挂了同一个 CheckedIoPass，它应当把带 checked 标记的下标改写成 checked_index 的 index.min(buffer_len - 1)。

问题在 CheckedIo 的匹配条件：

impl MatchRewrite for CheckedIo {
    fn r#match(&mut self, ctx: &Context, op: Ptr<Operation>) -> bool {
        Operation::get_op::<IndexOp>(op, ctx)
            .is_some_and(|it| it.checked(ctx) && is_runtime_array(ctx, it.base(ctx)))
    }
    // ...
}

fn is_runtime_array(ctx: &Context, value: Value) -> bool {
    let ty = value.get_type(ctx).deref(ctx);
    ty.downcast_ref::<RuntimeArrayType>().is_some()
}

它要求 IndexOp 的 base 的类型本身是 RuntimeArrayType。可是 IndexOp 声明了 OperandNOfType<0, PointerType>，base 必须是指针；全局缓冲在 Scope::global 里建的正是 PointerType::get(ctx, ty_arr.into(), AddressSpace::Global(buffer_pos))，指向运行期数组的指针。外层永远是 PointerType，downcast_ref::<RuntimeArrayType>() 永远落空，这个改写从不发生。.deref(ctx) 是 pliron 取出类型对象，不是解引用指针。

WGSL 这一侧还有一层：cubecl 建 shader 模块时关掉了 wgpu 自己的越界检查（bounds_checks: false），理由写在旁边，就是依赖 cube 的 checked 模式：

                let checks = wgpu::ShaderRuntimeChecks {
                    // Cube does not need wgpu bounds checks - OOB behaviour is instead
                    // checked by cube (if enabled).
                    // This is because the WebGPU specification only makes loose guarantees that Cube can't rely on.
                    bounds_checks: false,

下面的 create_shader_module_trusted 前面的 SAFETY 注释是「Cube guarantees OOB safety when launching in checked mode」。MSL 走的是 create_shader_module_passthrough，源码直接交给 Metal。两条路上，checked 模式失效之后都没有别的保护。

现在的匹配条件来自 cubecl #1402（「refactor: Make everything Pliron」，首次进入 v0.11.0-pre.3）。改之前 CheckedIoVisitor 按 Memory::Index(op) if op.checked 加 op.list.has_buffer_length() 判断，不看类型；那次提交的说明只有标题，看不出条件改成按类型判断是不是有意的。过程见 #[cube] 内核的冷知识。

📝NOTE

没有蠢问题

问：越界读到的是别的缓冲区的数据，这不是 GPU 驱动该管的吗？

答：本页的越界都发生在一块 8 个元素的分配之内、句柄露出的 6 个之外，从 Metal 和 wgpu 看都是同一块内存里的合法地址；cubecl 的缓冲区又是从内存池的大页里切出来的，相邻的切片在同一页上。在 WGSL 路径上 cubecl 还明确关掉了 wgpu 的越界检查（见上）。越出整块分配、整页时驱动会不会拦，本页没有测。

修法

内核自己写守卫，launch 和 launch_unchecked 都一样。上面的 add_guarded 就是一种写法，实测它挡住了对 out 的越界写。它也说明了守卫的范围：只判断了 out.len()，「读」那一行里 x、y 比 out 短，照样越界读了。几个缓冲区长度可能不同时，每个都要判断，或者在启动前确认它们一样长。启动的单元数通常向上取整到 CubeDim 的倍数，多出来的单元靠守卫挡住，#[cube] 内核里的 saxpy 就是这样写的。

这是 cubecl 0.11.3 的上游问题（实测 MSL 和 WGSL；源码上 SPIR-V 和 CUDA / HIP 挂的是同一个 pass、同一个条件），修复要改 CheckedIo::match，让它看指针指向的类型。在那之前，不要把 launch 当成越界保护。

burn 自己的内核

burn-cubecl 里有 9 个内核用的是 #[cube(launch, ..)]，比如 cast_element、mask_where_kernel、select_assign_kernel。select_assign_kernel 的注释写明了依赖 checked 模式：

/// Uses checked launch mode because user-provided `indices` may contain out-of-bounds values
/// that would cause invalid writes into `tensor`. Checked mode clamps these accesses rather
/// than producing undefined behavior.
#[cube(launch, address_type = "dynamic")]
fn select_assign_kernel<F: Numeric, I: Numeric, Op: BinaryOpFamily>(

它的写入另外用了显式判断的 write_checked（index < list.len()），读 tensor[index_tensor] 没有。拿一个越界的下标试（[3, 2] 的零张量，沿维 1 用下标 2 做 select_assign 加 1；pc-gpu select）：

metal: [0.0, 0.0, 1.0, 0.0, 1.0, 0.0]

thread 'main' (3597248) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-flex-0.22.1/src/ops/gather_scatter.rs:144:5:
index 2 out of bounds for dimension of size 2
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
flex: panicked

Flex 当场 panic；Metal 不报错，第 0、1 行的「第 2 列」落到了下一行的第 0 列，第 2 行的那一次超出整个张量，被 write_checked 跳过。这个错位不是本页的问题造成的：checked 模式和 write_checked 夹的都是整个缓冲区的长度，不管某一维越界。Tensor::select_assign 的文档也写了「Not all backends have runtime bound checks for the indices」。本页的问题在这里只多了一次越界读（读的值随后被丢掉）。下标来自用户数据时，在 GPU 上要自己保证它在范围内。

怎么确认没中招

照本页的办法：缓冲区多放几个可辨认的值，用 Handle::offset_end 只露出前一部分，看越界的单元读写到了什么。
开 CUBECL_DEBUG_LOG=stdout，在生成的 MSL / WGSL 里找下标处有没有 min 或比较；这个开关会让启动变成阻塞的，只在调试时开。
每个直接用 x[ABSOLUTE_POS] 的内核，检查前面有没有 if ABSOLUTE_POS < ...len()。

小结

已核实（M2 Max 实测 + 源码）：cubecl 0.11.3 的 launch 在 wgpu<msl> 和 wgpu<wgsl> 上生成的代码和 launch_unchecked 相同，越界读写直接发生，不夹取、不报错。
原因：CheckedIo::match 要求下标的 base 类型是 RuntimeArrayType，而 IndexOp 的 base 必须是 PointerType，全局缓冲正是指针，改写从不触发。
没有核实：SPIR-V、CUDA、HIP 路径只读了源码（同一个 pass）；越出整块分配时驱动会不会拦。
修法：内核里自己写守卫，每个可能越界的缓冲区都要判断；GPU 上的下标数据（如 select_assign 的 indices）要自己保证在范围内。

自定义内核越界读不报错，也不截断

现象

出处

修法

burn 自己的内核

怎么确认没中招

小结

💬 关联讨论

登录