c++ - Halide:将C++函数传递给Halide Func

我有一个二进制图像，并希望使用Halide从图像的顶部开始为每一列查找第一个非零像素。

在c++中，给定名为mask的图像，它看起来像这样:

vector<int> top_y;
top_y.reserve(mask.n_cols);
for (size_t x = 0; x < mask.n_cols; ++x) {
    for (size_t y = 0; y < mask.n_rows; ++y) {
        if (mask(y,x) != 0) {
            top_y[x] = y;
            break;
        } else if (y == mask.n_rows-1) {
            top_y[x] = mask.n_rows);
        }
    }
}

我已经看到了这种for, for, if结构化循环的示例(例如，使用RDom::where指令-参见tutorial lesson 17)，但是这种情况因break;的使用而不同。

考虑到外循环的并行性，也许可以将C++函数(由内循环函数组成，包括break)传递给Halide Func，然后在图像的所有列上实现该Func。

如果是这样，您能指导我举例说明如何实现吗？

最佳答案

您想要的可以在(image(x，y)!= 0)上的纯Halide argmax中获得-它会返回第一个真值的索引。但这不会产生中断行为。这是我们一直想实现的优化，但尚未实现。

您可以使用Func::define_extern陷入任意C++阶段。您可以在它们上使用compute_at来对某个使用方Func的每一列进行外部调用，然后对使用方Func进行常规的Halide调度以并行遍历各列。

有关define_extern用法的示例，请参见:https://github.com/halide/Halide/blob/master/test/correctness/extern_stage.cpp