Fix infinite recursion in find intrinsics

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());

src/runtime/HalideBuffer.h

@@ -2463,7 +2463,10 @@ class Buffer {
     template<typename Fn>
     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<Fn>(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<Fn>(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<float, 3> 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<typename Fn,
              typename = std::enable_if_t<!std::is_arithmetic_v<std::decay_t<Fn>>>>
     Buffer<T, Dims, InClassDimStorage> &fill(Fn &&f) {
-        // We'll go via for_each_element. We need a variadic wrapper lambda.
-        FillHelper<Fn> wrapper(std::forward<Fn>(f), this);
-        return for_each_element(wrapper);
+        if constexpr (std::is_invocable_v<Fn &>) {
+            // 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<Fn> wrapper(std::forward<Fn>(f), this);
+            for_each_element(wrapper);
+        }
+        return *this;
     }
 
     /** Check if an input buffer passed extern stage is a querying

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",

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<int8_t> 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<int64_t> out_vec = f_vec.realize({64});
+        Buffer<int64_t> 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;
     }

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<uint8_t> 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<uint8_t> dst_image = Buffer<uint8_t>::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);