Add hoist_invariants() directive for semiring reductions

python_bindings/halide/src/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("eager_inline", (Stage & (Stage::*)(const std::vector<Func> &)) & Stage::eager_inline, py::arg("fs"))
             .def("eager_inline", [](Stage &stage, const py::args &args) -> Stage & {

src/Func.h

@@ -18,6 +18,7 @@
 #include "Var.h"
 
 #include <map>
+#include <type_traits>
 #include <utility>
 
 namespace Halide {
@@ -213,6 +214,44 @@ class Stage {
     eager_inline(const Func &first, Args &&...args);
     // @}
 
+    /** Hoist a loop-invariant factor out of an associative reduction by applying
+     * the distributive law of a semiring. Like rfactor(), this must be called on
+     * an update definition; it splits the update into an intermediate that
+     * accumulates the factor-free reduction over all of the update's RVars and a
+     * write-back that applies the hoisted factor once. The intermediate Func is
+     * returned.
+     *
+     * A factor is hoistable if it does not depend on any RVar being reduced. It
+     * may be nested at any depth of an associative/commutative chain. The valid
+     * hoistings are:
+     *
+     *   Outer op    Inner combine   Law
+     *   ---------   -------------   ---
+     *   + (sum)     *               sum_k(s * x_k) = s * sum_k(x_k)
+     *   min         +               min_k(c + x_k) = c + min_k(x_k)
+     *   max         +               max_k(c + x_k) = c + max_k(x_k)
+     *   || (bool)   &&              or_k(p && x_k) = p && or_k(x_k)
+     *   && (bool)   ||              and_k(p || x_k) = p || and_k(x_k)
+     *
+     * For example, hoist_invariants() rewrites a pipeline like this:
+     * \code
+     * f(x) = 0;
+     * f(x) += s(x) * g(x, r);
+     * \endcode
+     * into a pipeline like this:
+     * \code
+     * f_intm(x) = 0;
+     * f_intm(x) += g(x, r);
+     *
+     * f(x) = 0;
+     * f(x) += s(x) * f_intm(x);
+     * \endcode
+     *
+     * This reduces the number of factor applications from |R| to one per pure
+     * point. It is an error if no distributable invariant factor is found.
+     */
+    Func hoist_invariants();
+
     /** Schedule the iteration over this stage to be fused with another
      * stage 's' from outermost loop to a given LoopLevel. 'this' stage will
      * be computed AFTER 's' in the innermost fused dimension. There should not

test/correctness/CMakeLists.txt

@@ -196,6 +196,7 @@ tests(
     hexagon_scatter.cpp
     histogram.cpp
     histogram_equalize.cpp
+    hoist_invariants.cpp
     hoist_loop_invariant_if_statements.cpp
     hoist_storage.cpp
     host_alignment.cpp