Reject split factors not representable as int32 (#9143)

src/Func.cpp

@@ -1073,9 +1073,19 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
     return intm;
 }
 
-void Stage::split(const string &old, const string &outer, const string &inner, const Expr &factor, bool exact, TailStrategy tail) {
+void Stage::split(const string &old, const string &outer, const string &inner, const Expr &factor_arg, bool exact, TailStrategy tail) {
     debug(4) << "In schedule for " << name() << ", split " << old << " into "
-             << outer << " and " << inner << " with factor of " << factor << "\n";
+             << outer << " and " << inner << " with factor of " << factor_arg << "\n";
+
+    user_assert(factor_arg.defined())
+        << "In schedule for " << name() << ", split factor for splitting "
+        << old << " is undefined.\n";
+    user_assert(Int(32).can_represent(factor_arg.type()))
+        << "In schedule for " << name() << ", split factor for splitting "
+        << old << " has type " << factor_arg.type()
+        << ", which is not representable as int32.\n";
+    Expr factor = cast<int32_t>(factor_arg);
+
     vector<Dim> &dims = definition.schedule().dims();
 
     definition.schedule().touched() = true;

test/correctness/CMakeLists.txt

@@ -303,6 +303,7 @@ tests(GROUPS correctness
       specialize_to_gpu.cpp
       specialize_trim_condition.cpp
       split_by_non_factor.cpp
+      split_factor_type.cpp
       split_fuse_rvar.cpp
       split_reuse_inner_name_bug.cpp
       split_store_compute.cpp

test/correctness/split_factor_type.cpp

@@ -0,0 +1,55 @@
+#include "Halide.h"
+#include <cstdio>
+
+using namespace Halide;
+
+int main(int argc, char **argv) {
+    // Split factors with types not representable as int32 (e.g. uint32)
+    // should be rejected up front with a clean error, not produce an
+    // internal compiler error during lowering. See issue #9143.
+
+    // Types that fit in int32 should still work.
+    {
+        Var x, xo, xi;
+        Func f("f");
+        f(x) = x;
+        f.bound(x, 0, 64);
+        f.split(x, xo, xi, Expr((int16_t)8));
+        f.compile_jit();
+    }
+    {
+        Var x, xo, xi;
+        Func f("g");
+        f(x) = x;
+        f.bound(x, 0, 64);
+        f.split(x, xo, xi, Expr((uint16_t)8));
+        f.compile_jit();
+    }
+
+#if HALIDE_WITH_EXCEPTIONS
+    // A uint32 split factor is not representable as int32, so it should
+    // be rejected at split() time. Only check this if exceptions are
+    // available at runtime, so we can recover from the error.
+    if (Halide::exceptions_enabled()) {
+        Var x, xo, xi;
+        Func f("h");
+        f(x) = x;
+        f.bound(x, 0, 64);
+
+        bool error = false;
+        try {
+            f.split(x, xo, xi, Expr((uint32_t)8));
+        } catch (const Halide::CompileError &e) {
+            error = true;
+            printf("Expected compile error:\n%s\n", e.what());
+        }
+        if (!error) {
+            printf("There was supposed to be an error!\n");
+            return 1;
+        }
+    }
+#endif
+
+    printf("Success!\n");
+    return 0;
+}