Allow hoisting invariant factors in rfactor

python_bindings/src/halide/halide_/PyFunc.cpp

@@ -219,6 +219,8 @@ void define_func(py::module &m) {
             .def("bound_storage", &Func::bound_storage)
             .def("memoize", &Func::memoize, py::arg("eviction_key") = EvictionKey())
             .def("compute_inline", &Func::compute_inline)
+            .def("eager_inline", &Func::eager_inline, py::arg("fs"))
+            .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)
 

python_bindings/src/halide/halide_/PyStage.cpp

@@ -23,6 +23,7 @@ void define_stage(py::module &m) {
                  py::arg("preserved"))
             .def("rfactor", static_cast<Func (Stage::*)(const RVar &, const Var &)>(&Stage::rfactor),
                  py::arg("r"), py::arg("v"))
+            .def("hoist_invariants", &Stage::hoist_invariants)
 
             .def("split_vars", [](const Stage &stage) -> py::list {
                 auto vars = stage.split_vars();

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

@@ -109,6 +109,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)); \
@@ -362,8 +367,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());
     }();
@@ -376,5 +380,23 @@ const vector<AssociativePattern> &get_ops_table(const vector<Expr> &exprs) {
     return table;
 }
 
+std::optional<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 std::nullopt;
+    }
+
+    const Expr &identity = table.front().identities.front();
+    for (const AssociativePattern &pattern : table) {
+        internal_assert(pattern.size() == 1);
+        if (!equal(pattern.identities.front(), identity)) {
+            return std::nullopt;
+        }
+    }
+    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. */
+std::optional<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