diff --git a/src/runtime/HalideBuffer.h b/src/runtime/HalideBuffer.h index b49272d0f962..f95541221c2b 100644 --- a/src/runtime/HalideBuffer.h +++ b/src/runtime/HalideBuffer.h @@ -2463,7 +2463,10 @@ class Buffer { template HALIDE_ALWAYS_INLINE static void for_each_element(double, int dims, const for_each_element_task_dim *t, Fn &&f) { int args = num_args(0, std::forward(f)); - assert(dims >= args); + assert(args == dims && + "Callable passed to for_each_element (or fill) must accept either a " + "const int *, or a number of int arguments equal to the dimensionality " + "of the buffer."); for_each_element_variadic(0, args - 1, t, std::forward(f)); } @@ -2481,19 +2484,14 @@ class Buffer { public: /** Call a function at each site in a buffer. This is likely to be * much slower than using Halide code to populate a buffer, but is - * convenient for tests. If the function has more arguments than the - * buffer has dimensions, the remaining arguments will be zero. If it - * has fewer arguments than the buffer has dimensions then the last - * few dimensions of the buffer are not iterated over. For example, - * the following code exploits this to set a floating point RGB image - * to red: + * convenient for tests. The function must take a number of int + * arguments equal to the dimensionality of the buffer. For example, + * the following code sets a floating point RGB image to red: \code Buffer im(100, 100, 3); - im.for_each_element([&](int x, int y) { - im(x, y, 0) = 1.0f; - im(x, y, 1) = 0.0f; - im(x, y, 2) = 0.0f: + im.for_each_element([&](int x, int y, int c) { + im(x, y, c) = (c == 0) ? 1.0f : 0.0f; }); \endcode @@ -2572,13 +2570,25 @@ class Buffer { /** Fill a buffer by evaluating a callable at every site. The * callable should look much like a callable passed to * for_each_element, but it should return the value that should be - * stored to the coordinate corresponding to the arguments. */ + * stored to the coordinate corresponding to the arguments. A callable + * that takes no arguments (for example a random number generator) is + * evaluated once per element, ignoring the coordinate. */ template>>> Buffer &fill(Fn &&f) { - // We'll go via for_each_element. We need a variadic wrapper lambda. - FillHelper wrapper(std::forward(f), this); - return for_each_element(wrapper); + if constexpr (std::is_invocable_v) { + // A callable that takes no coordinates has nothing to + // distinguish sites, so fill every value with its result. + // for_each_value writes through raw references, so it does not + // mark the host buffer dirty the way operator() does. + set_host_dirty(); + for_each_value([&](T &v) { v = f(); }); + } else { + // We'll go via for_each_element. We need a variadic wrapper lambda. + FillHelper wrapper(std::forward(f), this); + for_each_element(wrapper); + } + return *this; } /** Check if an input buffer passed extern stage is a querying diff --git a/test/correctness/for_each_element.cpp b/test/correctness/for_each_element.cpp index 85fcfc2380fd..487d1e0e16b8 100644 --- a/test/correctness/for_each_element.cpp +++ b/test/correctness/for_each_element.cpp @@ -27,8 +27,10 @@ int main(int argc, char **argv) { im(pos) *= 3; }); - im.for_each_element([&](int x, int y) { - for (int c = 0; c < 3; c++) { + // Slice off a channel to get a two-dimensional domain to iterate + // over, then handle all channels inside the callable. + im.sliced(2, im.dim(2).min()).for_each_element([&](int x, int y) { + for (int c : im.dim(2)) { int correct = (10 * x + 5 * y + c) * 3; if (im(x, y, c) != correct) { printf("im(%d, %d, %d) = %d instead of %d\n", diff --git a/test/performance/rgb_interleaved.cpp b/test/performance/rgb_interleaved.cpp index 3adfbeda4f52..68b53fc81d86 100644 --- a/test/performance/rgb_interleaved.cpp +++ b/test/performance/rgb_interleaved.cpp @@ -33,7 +33,7 @@ void test_deinterleave() { // Setup dst to be planar, with no extra padding between channels or rows. Buffer dst_image(1 << 12, 1 << 12, 3); - src_image.for_each_element([&](int x, int y) { + src_image.sliced(2, 0).for_each_element([&](int x, int y) { src_image(x, y, 0) = 0; src_image(x, y, 1) = 128; src_image(x, y, 2) = 255; @@ -54,7 +54,7 @@ void test_deinterleave() { printf("Interleaved to planar bandwidth %.3e byte/s.\n", dst_image.number_of_elements() / t1); - dst_image.for_each_element([&](int x, int y) { + dst_image.sliced(2, 0).for_each_element([&](int x, int y) { assert(dst_image(x, y, 0) == 0); assert(dst_image(x, y, 1) == 128); assert(dst_image(x, y, 2) == 255); @@ -69,7 +69,7 @@ void test_deinterleave() { dst.realize(dst_image); }); - dst_image.for_each_element([&](int x, int y) { + dst_image.sliced(2, 0).for_each_element([&](int x, int y) { assert(dst_image(x, y, 0) == 0); assert(dst_image(x, y, 1) == 128); assert(dst_image(x, y, 2) == 255); @@ -111,7 +111,7 @@ void test_interleave(bool fast) { // Setup dst to be interleaved Buffer dst_image = Buffer::make_interleaved(1 << 12, 1 << 12, 3); - src_image.for_each_element([&](int x, int y) { + src_image.sliced(2, 0).for_each_element([&](int x, int y) { src_image(x, y, 0) = 0; src_image(x, y, 1) = 128; src_image(x, y, 2) = 255; @@ -136,7 +136,7 @@ void test_interleave(bool fast) { printf("Planar to interleaved bandwidth %.3e byte/s.\n", dst_image.number_of_elements() / t); - dst_image.for_each_element([&](int x, int y) { + dst_image.sliced(2, 0).for_each_element([&](int x, int y) { assert(dst_image(x, y, 0) == 0); assert(dst_image(x, y, 1) == 128); assert(dst_image(x, y, 2) == 255);