Add change_type() directive with reduction-aware overflow checks

Makefile

@@ -447,6 +447,7 @@ SOURCE_FILES = \
   AddImageChecks.cpp \
   AddParameterChecks.cpp \
   AddSplitFactorChecks.cpp \
+  AddTypeChangeChecks.cpp \
   AlignLoads.cpp \
   AllocationBoundsInference.cpp \
   ApplySplit.cpp \
@@ -649,6 +650,7 @@ HEADER_FILES = \
   AddImageChecks.h \
   AddParameterChecks.h \
   AddSplitFactorChecks.h \
+  AddTypeChangeChecks.h \
   AlignLoads.h \
   AllocationBoundsInference.h \
   ApplySplit.h \

python_bindings/src/halide/halide_/PyFunc.cpp

@@ -223,6 +223,7 @@ void define_func(py::module &m) {
             .def("eager_inline", [](Func &func, const py::args &args) -> Func & {
                 return func.eager_inline(args_to_vector<Func>(args));
             })
+            .def("change_type", &Func::change_type, py::arg("type"), py::arg("unsafe") = false)
             .def("compute_root", &Func::compute_root)
             .def("store_root", &Func::store_root)
 

src/AddTypeChangeChecks.cpp

@@ -0,0 +1,35 @@
+#include "AddTypeChangeChecks.h"
+#include "Function.h"
+#include "IR.h"
+#include "IROperator.h"
+#include "Schedule.h"
+#include "Simplify.h"
+
+namespace Halide {
+namespace Internal {
+
+Stmt add_type_change_checks(const Stmt &s, const std::map<std::string, Function> &env) {
+    std::vector<Stmt> stmts;
+
+    for (const auto &p : env) {
+        const Function &f = p.second;
+        for (const auto &[condition, message] : f.schedule().type_change_checks()) {
+            if (!condition.defined()) {
+                continue;
+            }
+            Expr proven = simplify(condition);
+            if (is_const_one(proven)) {
+                // Statically proven; no runtime check needed.
+                continue;
+            }
+            Expr error = requirement_failed_error(condition, {Expr(message)});
+            stmts.push_back(AssertStmt::make(condition, error));
+        }
+    }
+
+    stmts.push_back(s);
+    return Block::make(stmts);
+}
+
+}  // namespace Internal
+}  // namespace Halide

src/AddTypeChangeChecks.h

@@ -0,0 +1,29 @@
+#ifndef HALIDE_ADD_TYPE_CHANGE_CHECKS_H
+#define HALIDE_ADD_TYPE_CHANGE_CHECKS_H
+
+/** \file
+ * Defines the lowering pass that injects the overflow-safety preconditions
+ * recorded by Func::change_type() into the pipeline's assertion block.
+ */
+
+#include <map>
+#include <string>
+
+#include "Expr.h"
+
+namespace Halide {
+namespace Internal {
+
+class Function;
+
+/** Prepend assertions for any static preconditions that Func::change_type()
+ * recorded on the funcs in `env` (that it could not discharge at schedule time,
+ * e.g. because a reduction extent was symbolic). Statically-true conditions are
+ * dropped. Like the other check passes, the resulting asserts are removed later
+ * when the no_asserts target feature is set. */
+Stmt add_type_change_checks(const Stmt &s, const std::map<std::string, Function> &env);
+
+}  // namespace Internal
+}  // namespace Halide
+
+#endif

src/AssociativeOpsTable.cpp

@@ -101,6 +101,11 @@ struct TableKey {
 
 map<TableKey, vector<AssociativePattern>> pattern_tables;
 
+std::mutex &ops_table_lock() {
+    static std::mutex lock;
+    return lock;
+}
+
 #define declare_vars(t, index)                                        \
     Expr x##index = Variable::make((t), "x" + std::to_string(index)); \
     Expr y##index = Variable::make((t), "y" + std::to_string(index)); \
@@ -354,8 +359,7 @@ const vector<AssociativePattern> &get_ops_table(const vector<Expr> &exprs) {
     const vector<AssociativePattern> &table = [&]() -> decltype(auto) {
         // get_ops_table_helper() lazily initializes the table, so ensure
         // that multiple threads can't try to do so at the same time.
-        static std::mutex ops_table_lock;
-        std::scoped_lock lock_guard(ops_table_lock);
+        std::scoped_lock lock_guard(ops_table_lock());
 
         return get_ops_table_helper(types, exprs[0].node_type(), exprs.size());
     }();
@@ -368,5 +372,23 @@ const vector<AssociativePattern> &get_ops_table(const vector<Expr> &exprs) {
     return table;
 }
 
+Expr get_associative_identity(Type type, IRNodeType root) {
+    std::scoped_lock lock_guard(ops_table_lock());
+
+    const vector<AssociativePattern> &table = get_ops_table_helper({type}, root, 1);
+    if (table.empty()) {
+        return Expr();
+    }
+
+    const Expr &identity = table.front().identities.front();
+    for (const AssociativePattern &pattern : table) {
+        internal_assert(pattern.size() == 1);
+        if (!equal(pattern.identities.front(), identity)) {
+            return Expr();
+        }
+    }
+    return identity;
+}
+
 }  // namespace Internal
 }  // namespace Halide

src/AssociativeOpsTable.h

@@ -8,6 +8,7 @@
 #include "IREquality.h"
 #include "IROperator.h"
 
+#include <optional>
 #include <utility>
 #include <vector>
 
@@ -71,6 +72,10 @@ struct AssociativePattern {
 
 const std::vector<AssociativePattern> &get_ops_table(const std::vector<Expr> &exprs);
 
+/** Return the identity for a single-output associative op, if the table has one
+ * and all matching patterns agree on it. */
+Expr get_associative_identity(Type type, IRNodeType root);
+
 }  // namespace Internal
 }  // namespace Halide
 

src/CMakeLists.txt

@@ -57,6 +57,7 @@ target_sources(
     AddImageChecks.h
     AddParameterChecks.h
     AddSplitFactorChecks.h
+    AddTypeChangeChecks.h
     AlignLoads.h
     AllocationBoundsInference.h
     ApplySplit.h
@@ -239,6 +240,7 @@ target_sources(
     AddImageChecks.cpp
     AddParameterChecks.cpp
     AddSplitFactorChecks.cpp
+    AddTypeChangeChecks.cpp
     AlignLoads.cpp
     AllocationBoundsInference.cpp
     ApplySplit.cpp

src/ConstantBounds.cpp

@@ -9,11 +9,12 @@ namespace Internal {
 namespace {
 ConstantInterval bounds_helper(const Expr &e,
                                Scope<ConstantInterval> &scope,
-                               std::map<Expr, ConstantInterval, ExprCompare> *cache) {
+                               std::map<Expr, ConstantInterval, ExprCompare> *cache,
+                               const FuncValueBounds *func_bounds) {
     internal_assert(e.defined());
 
     auto recurse = [&](const Expr &e) {
-        return bounds_helper(e, scope, cache);
+        return bounds_helper(e, scope, cache, func_bounds);
     };
 
     auto get_infinite_bounds = [&]() {
@@ -129,6 +130,13 @@ ConstantInterval bounds_helper(const Expr &e,
                 // We can't do much with the other bitwise ops, but we can treat
                 // bitwise_not as an all-ones bit pattern minus the argument.
                 return recurse(make_const(e.type(), -1) - op->args[0]);
+            } else if (func_bounds && op->call_type == Call::Halide && op->type.is_int_or_uint()) {
+                auto it = func_bounds->find({op->name, op->value_index});
+                if (it != func_bounds->end()) {
+                    return ConstantInterval::make_intersection(
+                        ConstantInterval::bounds_of_type(op->type),
+                        covering_constant_interval(it->second));
+                }
             }
             // If you add a new intrinsic here, also add it to the expression
             // generator in test/correctness/lossless_cast.cpp
@@ -163,10 +171,11 @@ ConstantInterval bounds_helper(const Expr &e,
 
 ConstantInterval constant_integer_bounds(const Expr &e,
                                          const Scope<ConstantInterval> &scope,
-                                         std::map<Expr, ConstantInterval, ExprCompare> *cache) {
+                                         std::map<Expr, ConstantInterval, ExprCompare> *cache,
+                                         const FuncValueBounds *func_bounds) {
     Scope<ConstantInterval> sub_scope;
     sub_scope.set_containing_scope(&scope);
-    return bounds_helper(e, sub_scope, cache);
+    return bounds_helper(e, sub_scope, cache, func_bounds);
 }
 
 }  // namespace Internal

src/ConstantBounds.h

@@ -1,6 +1,7 @@
 #ifndef HALIDE_CONSTANT_BOUNDS_H
 #define HALIDE_CONSTANT_BOUNDS_H
 
+#include "Bounds.h"
 #include "ConstantInterval.h"
 #include "Expr.h"
 #include "Scope.h"
@@ -19,15 +20,19 @@ namespace Internal {
  * negated, be incremented, etc without risking overflow.
  *
  * Also optionally accepts a scope containing the integer bounds of any
- * variables that may be referenced, and a cache of constant integer bounds on
- * known Exprs, which this function will update. The cache is helpful to
- * short-circuit large numbers of redundant queries, but it should not be used
- * in contexts where the same Expr object may take on different values within a
- * single Expr (i.e. before uniquify_variable_names).
+ * variables that may be referenced, a cache of constant integer bounds on
+ * known Exprs, which this function will update, and previously-computed
+ * FuncValueBounds for any Halide Call nodes encountered, which lets a call to
+ * a producer Func (e.g. one known to be the result of a clamp) get a tighter
+ * bound than its type's full range. The cache is helpful to short-circuit
+ * large numbers of redundant queries, but it should not be used in contexts
+ * where the same Expr object may take on different values within a single
+ * Expr (i.e. before uniquify_variable_names).
  */
 ConstantInterval constant_integer_bounds(const Expr &e,
                                          const Scope<ConstantInterval> &scope = Scope<ConstantInterval>::empty_scope(),
-                                         std::map<Expr, ConstantInterval, ExprCompare> *cache = nullptr);
+                                         std::map<Expr, ConstantInterval, ExprCompare> *cache = nullptr,
+                                         const FuncValueBounds *func_bounds = nullptr);
 
 }  // namespace Internal
 }  // namespace Halide

src/ConstantInterval.cpp

@@ -1,8 +1,10 @@
 #include "ConstantInterval.h"
 
+#include "ConstantBounds.h"
 #include "Error.h"
 #include "IROperator.h"
 #include "IRPrinter.h"
+#include "Interval.h"
 
 namespace Halide {
 namespace Internal {
@@ -101,6 +103,14 @@ bool ConstantInterval::contains(uint64_t x) const {
     }
 }
 
+bool ConstantInterval::contains(const ConstantInterval &other) const {
+    // Every value in `other` must lie within this interval. Where `other` is
+    // unbounded, this must be unbounded on the same side to contain it.
+    const bool too_small = min_defined && (!other.min_defined || other.min < min);
+    const bool too_large = max_defined && (!other.max_defined || other.max > max);
+    return !(too_small || too_large);
+}
+
 ConstantInterval ConstantInterval::make_union(const ConstantInterval &a, const ConstantInterval &b) {
     ConstantInterval result = a;
     result.include(b);
@@ -140,6 +150,17 @@ ConstantInterval ConstantInterval::make_intersection(const ConstantInterval &a,
     return result;
 }
 
+ConstantInterval covering_constant_interval(const Interval &in) {
+    auto min_bounds = constant_integer_bounds(in.min);
+    auto max_bounds = constant_integer_bounds(in.max);
+    ConstantInterval result;
+    result.min = min_bounds.min;
+    result.min_defined = min_bounds.min_defined;
+    result.max = max_bounds.max;
+    result.max_defined = max_bounds.max_defined;
+    return result;
+}
+
 void ConstantInterval::operator+=(const ConstantInterval &other) {
     (*this) = (*this) + other;
 }

src/ConstantInterval.h

@@ -13,6 +13,8 @@ struct Type;
 
 namespace Internal {
 
+struct Interval;
+
 /** A class to represent ranges of integers. Can be unbounded above or below,
  * but they cannot be empty. */
 struct ConstantInterval {
@@ -64,6 +66,11 @@ struct ConstantInterval {
     /** Test if the interval contains a particular unsigned value */
     bool contains(uint64_t x) const;
 
+    /** Test if this interval contains every value of another interval. An
+     * unbounded side of the other interval is contained only if this interval is
+     * also unbounded on that side. */
+    bool contains(const ConstantInterval &other) const;
+
     /** Construct the smallest interval containing two intervals. */
     static ConstantInterval make_union(const ConstantInterval &a, const ConstantInterval &b);
 
@@ -101,6 +108,12 @@ struct ConstantInterval {
     static ConstantInterval bounds_of_type(Type);
 };
 
+/** Convert a symbolic Interval to a ConstantInterval, keeping only endpoints that
+ * are already constant integers. A symbolic or infinite bound becomes unbounded.
+ * This does no bounds analysis of its own; any tightening must have already been
+ * done to the Interval upstream. */
+ConstantInterval covering_constant_interval(const Interval &in);
+
 /** Arithmetic operators on ConstantIntervals. The resulting interval contains
  * all possible values of the operator applied to any two elements of the
  * argument intervals. Note that these operator on unbounded integers. If you

src/Deserialization.cpp

@@ -1035,6 +1035,15 @@ FuncSchedule Deserializer::deserialize_func_schedule(const Serialize::FuncSchedu
     const auto async = func_schedule->async();
     const auto ring_buffer = deserialize_expr(func_schedule->ring_buffer_type(), func_schedule->ring_buffer());
     const auto memoize_eviction_key = deserialize_expr(func_schedule->memoize_eviction_key_type(), func_schedule->memoize_eviction_key());
+    std::vector<std::pair<Expr, std::string>> type_change_checks;
+    if (func_schedule->type_change_checks() != nullptr) {
+        type_change_checks.reserve(func_schedule->type_change_checks()->size());
+        for (const auto *check : *func_schedule->type_change_checks()) {
+            type_change_checks.emplace_back(
+                deserialize_expr(check->condition_type(), check->condition()),
+                deserialize_string(check->message()));
+        }
+    }
     auto hl_func_schedule = FuncSchedule();
     hl_func_schedule.store_level() = store_level;
     hl_func_schedule.compute_level() = compute_level;
@@ -1048,6 +1057,7 @@ FuncSchedule Deserializer::deserialize_func_schedule(const Serialize::FuncSchedu
     hl_func_schedule.async() = async;
     hl_func_schedule.ring_buffer() = ring_buffer;
     hl_func_schedule.memoize_eviction_key() = memoize_eviction_key;
+    hl_func_schedule.type_change_checks() = std::move(type_change_checks);
     return hl_func_schedule;
 }