c++ - 有没有办法在 Halide Generator 中将两个输入图像堆叠到一个 4 维缓冲区中？

Question

我的生成器管道中有两个输入图像 (rgb)。每个都经过自定义裁剪并调整为 NxNx3 图像。发电机的输出Output<Buffer<float>> {"batch", 4}; // (N, N, 3 ,2)我设置

batch(x, y, c, batch_size) = cast<float>(0); 
batch(x, y, c, 0) = image_1_resized(x, y, c);
batch(x, y, c, 1) = image_2_resized(x, y, c);

但是当我去编译它时它会出现段错误(我无法检索有用的调试信息)。如果我尝试这样做:

clamped_image_1 = BoundaryConditions(image_1_resized, 0 ,{{0, N}, {0, N});

batch(x, y, c, batch_size) = cast<float>(0); 
batch(x, y, c, 0) = image_1_resized(x, y, c);
batch(x, y, c, 1) = clamped_image_1(x + 10, y, c);

它编译并运行良好。我很想知道，如果使用 Halide 可以实现我想要实现的目标？

Answer 1

与其将其作为多阶段函数来执行，不如尝试使用 select 在单个阶段中填充不同的 channel 。

batch(x, y, c, n) = select(n==0, image_1_resized(x, y, c), image_2_resized(x, y, c));

然后你可以通过绑定(bind)和展开 n 来安排它，halide 将删除分支:

batch.bound(n, 0, 2).unroll(n);