diff --git a/src/FindIntrinsics.cpp b/src/FindIntrinsics.cpp index 87cb325c3fc6..612f3063d7fb 100644 --- a/src/FindIntrinsics.cpp +++ b/src/FindIntrinsics.cpp @@ -122,7 +122,13 @@ class FindIntrinsics : public IRMutator { } }; - if (inner_t.is_int_or_uint() && inner_t != x.type()) { + // Only strip a genuine widening cast, i.e. one whose inner type is no + // wider than x. If the outer cast is a narrowing cast, inner_t is wider + // than x, and considering it would return an expression wider than x, + // which can send the surrounding narrowing rewrites into infinite + // recursion. + if (inner_t.is_int_or_uint() && inner_t != x.type() && + inner_t.bits() <= x.type().bits()) { consider(inner_t); } consider(x.type().narrow()); 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/correctness/lossless_cast.cpp b/test/correctness/lossless_cast.cpp index 8959b72ccd53..e9d3d99055d0 100644 --- a/test/correctness/lossless_cast.cpp +++ b/test/correctness/lossless_cast.cpp @@ -274,6 +274,45 @@ int main(int argc, char **argv) { } } + // Runtime test: a widening_mul whose operand is a narrowing cast + // (int32 -> int16) of a widening_add used to send FindIntrinsics into + // infinite recursion. strip_widening_cast could return an expression wider + // than its input, which the widening_mul narrowing rewrite would re-widen, + // ping-ponging the multiply between its int32 and int64 forms forever. This + // must simply compile, and the vectorized result must match a scalar + // reference. + { + Var x("x"); + Buffer buf(64, "buf"); + for (int i = 0; i < 64; i++) { + buf(i) = (int8_t)(i - 32); + } + + Expr t = buf(x); + Expr inner = cast(Int(16), cast(UInt(8), t) + cast(UInt(8), 68)); + Expr a = widen_right_add(inner, t); + Expr b = widen_right_mul(a, cast(Int(8), a) * cast(Int(8), 39)); + Expr wide_add = widening_add(a, b); + Expr e = cast(Int(64), widening_mul(cast(Int(16), 73), cast(Int(16), wide_add))); + + Func f_vec, f_ref; + f_vec(x) = e; + f_ref(x) = e; + f_vec.vectorize(x, 4, TailStrategy::RoundUp); + + Buffer out_vec = f_vec.realize({64}); + Buffer out_ref = f_ref.realize({64}); + + for (int i = 0; i < 64; i++) { + if (out_vec(i) != out_ref(i)) { + std::cerr << "Recursion regression test: mismatch at x=" << i + << ": ref=" << out_ref(i) << " vectorized=" << out_vec(i) << "\n"; + found_error = true; + break; + } + } + } + if (found_error) { return 1; } 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);