跳转到内容
查看源码

执行模式

Kernel 以六种模式之一运行,由它的主体关键字决定。归约、scan、scatter、prefix 都“遍历”一个输入,每个元素运行一次主体,并用 over 指明遍历哪个输入;输入的声明顺序没有意义。

模式 主体 每次运行对应 输出 可融合
Map compute { … } 输出定义域的每个元素 同定义域,可多个输出 是
归约 compute over x { … } + merge(a, b) { … } x 的每个元素 恰好一个元素的 Rows 否
Scan scan over x with y = init { … } x 的每个元素,按顺序 与 x 等长 否
Scatter scatter over x into t with sum|min|max|last { … } x 的每个元素 目标 t 的形状 否
Prefix prefix over x with sum|min|max { … } x 的每个元素 与 x 等长的 Rows —
Group group N { … } 每 N 个元素一组协同运行 同 map —

各后端都实现全部模式;crates/backends/conformance 用“36 个语言特性 × 7 种执行形式”的矩阵检查 JIT 与 GPU 对照解释器的一致性。

对输出定义域的每个元素运行一次。Map 可以在同一定义域上写多个输出,用记录一次给出所有输出:

kernel shade.surface@1 {
requires { field.write; }
output albedo: Image;
output normal: Grid<Vec3>;
output depth: Grid<Float>;
compute {
let hit = cast(eye, ray());
if hit.t < 0.0 {
return { albedo: sky, normal: vec3(0.0, 0.0, 0.0), depth: 100.0 };
}
{ albedo: vec4(hit.albedo, 1.0), normal: hit.normal, depth: hit.t }
}
}

图中用 node.albedo 读取。当一个消费者是 map 所有输出的唯一读者,并且都通过 sample 读取时,两者融合,不写出任何中间输出。

kernel stats.luma@1 {
requires { field.read; field.write; }
input source: Image;
output result: Rows<Light>; // 恰好输出一个元素
struct Light { sum: Float; max: Float; count: Int; }
compute over source { // 对 source 的每个元素
let l = luma(sample(source));
return Light { sum: l, max: l, count: 1 };
}
merge(a, b) { // 两两合并
return Light { sum: a.sum + b.sum, max: max(a.max, b.max), count: a.count + b.count };
}
}

merge 必须满足结合律和交换律:合并树的形状由后端决定,整数与布尔精确,浮点在舍入误差内一致。归约节点不接受 domain,也从不融合。

kernel audio.lowpass@1 {
requires { field.read; field.write; }
input samples: Rows<Float>;
param k: Float = 0.1;
output result: Rows<Float>; // 每个输入元素一个输出元素
scan over samples with y = 0.0 { // 遍历的输入、携带值与初值
return y + (sample(samples) - y) * k;
}
}

主体按顺序对每个元素运行:y 是上一个结果(第一次是 init),返回值是当前输出元素,也是下一个 y;index() 是元素下标。Scan 在单线程运行,从不融合。

kernel machine.step@1 {
requires { field.read; field.write; }
input ticks: Rows<Int>; // 每个元素一步
memory ram: Rows<Int>; // 像输入一样绑定,像输出一样读取
output result: Rows<Int>;
scan over ticks with pc = 0 {
let op = at(ram, pc); // 能读到之前的步骤写入的值
put(ram, pc + 100, op * 2); // 按步骤顺序写入
return pc + 1;
}
}

ram <~ machine.ram from zeros.result 让内存跨帧保留,ram <- fresh.result 让每次求值从头开始。put 越界的写入被丢弃(不夹取);绑定的原 Field 从不被修改,求值失败时它保持原样。仓库中的 CHIP-8 和 ARM Thumb 模拟器(examples/machines)就是这样写的。

kernel particles.splat@1 {
requires { field.read; field.write; }
input particles: Rows<Particle>;
input canvas: Grid<Vec3>;
output result: Grid<Vec3>; // 初始为 canvas
struct Splat { cell: Vec2; heat: Vec3; }
scatter over particles into canvas with sum {
let p = sample(particles);
return Splat { cell: floor(p.position * size(canvas)), heat: tint(p) };
}
}
  • 主体返回一个两字段的结构体:目标中的地址(rows 用 Int,grid 用 Vec2,volume 用 Vec3)和一个目标元素类型的值。
  • 合并规则:sum 逐字段相加,min/max 取逐字段极值,last 替换。结果等同于按元素顺序依次应用。
  • 目标外的地址(包括 NaN)被丢弃。输出与目标同形状、同布局、同范围。
  • emit(pair); 可以额外发出若干对(只能在 min、max 或非浮点 sum 下使用)。
kernel compact.offsets@1 {
requires { field.read; field.write; }
input flags: Rows<Int>;
output result: Rows<Int>; // 每个输入元素一个输出元素
prefix over flags with sum { // 每个保留元素的名次
return if sample(flags) != 0 { 1 } else { 0 };
}
}

输出在 i 处是 0..=i 的值用 sum、min 或 max 合并的结果。Prefix 可能以任意分组合并,所以对含浮点的元素做 sum 会被拒绝(请用 scan),last 也不允许。

kernel fft.row@1 {
requires { field.read; field.write; }
input line: Rows<Vec2>;
output result: Rows<Vec2>;
group 64 { // 每组 64 个
let x = sample(line);
let pair = exchange(x, lane() ^ 1);
return if lane() % 2 == 0 { x + pair } else { pair - x };
}
}

元素 e 是第 e / N 组的第 e % N 个 lane(lane())。exchange(v, k) 取得 lane k 在同一次 exchange 时持有的 v。组内所有 lane 必须一致地到达每次 exchange,因此 exchange 不能位于依赖 lane 私有数据的分支、match、循环次数或 until 之下(依赖参数、字面量、定义域的可以)。在 GPU 上一组就是一个 workgroup,exchange 是共享内存加屏障;交换的数据超过 workgroup 内存时,该 group 改在 JIT 上运行。packages/signal(FFT)就建立在 group 之上。

这些写在场景的节点上,而不是 Kernel 里:

写法 作用
async = true; 节点离帧在自己的线程上运行;完成前读者拿到上一次的输出
times = N; + carry port <- output; 节点放置 N 次,后一次读前一次的输出(如 Jacobi 迭代)
awake = value(source, row); 该行为 0 时节点休眠,不运行

详见图与状态。