组合多个卤化物功能但保持尺寸相同

Question

我有三个 Halide 函数，它们具有以下输出维度：
40 x 40 x 64
40 x 40 x 128
40 x 40 x 64

我想将它们组合成一个函数，以便我得到一个函数句柄供以后使用。所以在这里，结果函数的维度应该是

40 x 40 x 256

我正在使用 Halide::select 但它会产生 4 个维度

concat(x,y,z,c)=Halide::select(c == 0, func_1(x, y, z), c == 1, func_2(x, y, z), func_3(x, y, z));

有什么方法可以生成合并的 3D 函数吗？

Answer 1

您可以使用 Tuple。由于 Tuple 的成员必须具有相同的大小，因此第 3 维的大小不同会有点复杂。您的 4d 解决方案也存在这种并发症。

result(x, y, z) = Tuple
    ( func_1(x, y, z)
    , func_2(x, y, z * 2 + 0) // Even z
    , func_2(x, y, z * 2 + 1) // Odd z
    , func_3(x, y, z)
    );

如果保留 4d 解决方案，请将 unroll(c) 添加到计划中，然后将在最内层循环内依次计算 3 个（现在是 4 个）函数。

与其在三维中连接三个函数，不如添加另一个维度或使用 Tuple 是更好的方法，恕我直言。

编辑：除了unroll(c)，您还需要reorder(c,x,y,z)来更改循环顺序。

Answer 2

您可能想要 return 管道对象而不是 Func。管道编译为具有多个输出参数的函数，这些参数可以是不同形状的缓冲区。

如果你确实想要一个 Func，你想要这样的东西：

concat(x, y, z) = 
  Halide::select(z < 64, func_1(x, y, clamp(z, 0, 63)), 
                 z < 192, func_2(x, y, clamp(z - 64, 0, 127)), 
                 func_3(x, y, clamp(z - 192, 0, 63)));

组合多个卤化物功能但保持尺寸相同

Combining multiple halide functions but keeping the dimensions same

halide