Halide中如何将rvar内联至var阶段并关联计算调度？

这种直接调用sum_elements.update().compute_with(output, i)的方式不可行，但可以通过调整计算定义和调度逻辑，实现你想要的单循环同时计算output和累加sum_elements的效果。

原方案不可行的原因 #

原代码中sum_elements的更新依赖独立的RDom r，而output的计算依赖Var i，两者的迭代空间没有直接映射关系。Halide无法将两个无关联的迭代逻辑强制合并到同一个循环中，直接执行原调度代码会触发编译错误。

实现目标效果的正确方式 #

要达到和目标C代码一致的单循环逻辑，需要让sum_elements的累加操作与output的迭代循环绑定，具体有两种实现思路：

思路1：复用Var i替代独立RDom #

直接用output的迭代变量i来完成sum_elements的累加，避免单独的RDom循环：

Input<Buffer<int>> vector {"vector", 1};
Output<Buffer<int>> output {"output", 1};
Output<Func> sum_elements {"sum_elements", Int(32), 0};

Var i;
// 定义output的计算逻辑
output(i) = 2 * vector(i);

// 初始化sum_elements为0
sum_elements() = 0;
// 复用output的迭代变量i完成累加，仅在i处于有效范围时执行
sum_elements() += output.dim(0).in_bounds() ? vector(i) : 0;

// 调度：将sum_elements的更新阶段嵌入output的循环
output.compute_root();
sum_elements.update().compute_with(output, i);

思路2：关联RDom与Var i的迭代空间 #

如果必须保留RDom，可以显式将其与output的迭代空间绑定，让Halide识别两者的对应关系：

Input<Buffer<int>> vector {"vector", 1};
Output<Buffer<int>> output {"output", 1};
Output<Func> sum_elements {"sum_elements", Int(32), 0};

Var i;
RDom r(0, vector.length());

output(i) = 2 * vector(i);
sum_elements() = 0;
sum_elements() += vector(r.x);

// 调度：先约束sum_elements的RDom范围与output一致，再绑定到output的循环
output.compute_root();
sum_elements.update().bound(r.x, 0, vector.length()).compute_with(output, i);

效果说明 #

两种方案最终都会生成与目标C代码逻辑一致的机器码：在同一个循环的每次迭代中，先计算output[i] = 2 * vector[i]，再将vector[i]累加到sum_elements中，避免了对输入向量的两次遍历，提升了执行效率。

内容的提问来源于stack exchange，提问作者Yotam Inbar