Allow hoisting invariant factors in rfactor

python_bindings/src/halide/halide_/PyEnums.cpp

@@ -60,6 +60,10 @@ void define_enums(py::module &m) {
         .value("GuardWithIf", PrefetchBoundStrategy::GuardWithIf)
         .value("NonFaulting", PrefetchBoundStrategy::NonFaulting);
 
+    py::enum_<RFactorOptions>(m, "RFactorOptions")
+        .value("None", RFactorOptions::None)
+        .value("HoistInvariantFactor", RFactorOptions::HoistInvariantFactor);
+
     py::enum_<StmtOutputFormat>(m, "StmtOutputFormat")
         .value("Text", StmtOutputFormat::Text)
         .value("HTML", StmtOutputFormat::HTML);

python_bindings/src/halide/halide_/PyStage.cpp

@@ -14,10 +14,10 @@ void define_stage(py::module &m) {
             .def("dump_argument_list", &Stage::dump_argument_list)
             .def("name", &Stage::name)
 
-            .def("rfactor", static_cast<Func (Stage::*)(const std::vector<std::pair<RVar, Var>> &)>(&Stage::rfactor),
-                 py::arg("preserved"))
-            .def("rfactor", static_cast<Func (Stage::*)(const RVar &, const Var &)>(&Stage::rfactor),
-                 py::arg("r"), py::arg("v"))
+            .def("rfactor", static_cast<Func (Stage::*)(const std::vector<std::pair<RVar, Var>> &, RFactorOptions)>(&Stage::rfactor),
+                 py::arg("preserved"), py::arg("options") = RFactorOptions::None)
+            .def("rfactor", static_cast<Func (Stage::*)(const RVar &, const Var &, RFactorOptions)>(&Stage::rfactor),
+                 py::arg("r"), py::arg("v"), py::arg("options") = RFactorOptions::None)
 
             .def("split_vars", [](const Stage &stage) -> py::list {
                 auto vars = stage.split_vars();

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