Possibility to align a split.

.gitignore

@@ -242,6 +242,7 @@ xcuserdata
 
 # NeoVim + clangd
 .cache
+.ccls-cache
 
 # Emacs
 tags

python_bindings/halide/src/halide_/PyScheduleMethods.h

@@ -29,6 +29,8 @@ HALIDE_NEVER_INLINE void add_schedule_methods(PythonClass &class_instance) {
 
         .def("split", (T & (T::*)(const VarOrRVar &, const VarOrRVar &, const VarOrRVar &, const Expr &, TailStrategy)) & T::split,
              py::arg("old"), py::arg("outer"), py::arg("inner"), py::arg("factor"), py::arg("tail") = TailStrategy::Auto)
+        .def("split", (T & (T::*)(const VarOrRVar &, const VarOrRVar &, const VarOrRVar &, const Expr &, const Expr &, TailStrategy)) & T::split,
+             py::arg("old"), py::arg("outer"), py::arg("inner"), py::arg("factor"), py::arg("align"), py::arg("tail") = TailStrategy::Auto)
 
         .def("fuse", &T::fuse,
              py::arg("inner"), py::arg("outer"), py::arg("fused"))

src/AllocationBoundsInference.cpp

@@ -1,5 +1,6 @@
 #include "AllocationBoundsInference.h"
 #include "Bounds.h"
+#include "BoundsTracker.h"
 #include "CSE.h"
 #include "ExternFuncArgument.h"
 #include "Function.h"
@@ -19,10 +20,6 @@ using std::vector;
 
 namespace {
 
-Expr cse_and_simplify(const Expr &x) {
-    return simplify(common_subexpression_elimination(x));
-}
-
 // Figure out the region touched of each buffer, and deposit them as
 // let statements outside of each realize node, or at the top level if
 // they're not internal allocations.
@@ -34,6 +31,25 @@ class AllocationInference : public IRMutator {
     const FuncValueBounds &func_bounds;
     set<string> touched_by_extern;
 
+    // Tracks the enclosing pure lets and for loops, so that box_touched's
+    // result -- computed with no knowledge of anything outside op->body --
+    // can be simplified as thoroughly as if it had been.
+    BoundsTracker tracker;
+
+    Expr cse_and_simplify(const Expr &x) {
+        return simplify(common_subexpression_elimination(tracker.simplify_with_context(x)));
+    }
+
+    Stmt visit(const LetStmt *op) override {
+        auto binding = tracker.push_let(op->name, op->value);
+        return IRMutator::visit(op);
+    }
+
+    Stmt visit(const For *op) override {
+        auto binding = tracker.push_for(op->name, op->min, op->max);
+        return IRMutator::visit(op);
+    }
+
     Stmt visit(const Realize *op) override {
         map<string, Function>::const_iterator iter = env.find(op->name);
         internal_assert(iter != env.end());

src/ApplySplit.cpp

@@ -23,10 +23,17 @@ vector<ApplySplitResult> apply_split(const Split &split, const string &prefix,
         Expr old_max = Variable::make(Int(32), prefix + split.old_var + ".loop_max");
         Expr old_min = Variable::make(Int(32), prefix + split.old_var + ".loop_min");
         Expr old_extent = (old_max - old_min) + 1;
+        Expr outer_min = Variable::make(Int(32), prefix + split.outer + ".loop_min");
 
         dim_extent_alignment[split.inner] = split.factor;
 
-        Expr base = outer * split.factor + old_min;
+        Expr base;
+        if (split.align.defined()) {
+            base = outer * split.factor;
+        } else {
+            base = outer * split.factor + old_min;
+        }
+
         string base_name = prefix + split.inner + ".base";
         Expr base_var = Variable::make(Int(32), base_name);
         string old_var_name = prefix + split.old_var;
@@ -38,8 +45,17 @@ vector<ApplySplitResult> apply_split(const Split &split, const string &prefix,
         internal_assert(tail != TailStrategy::Auto)
             << "An explicit tail strategy should exist at this point\n";
 
+        // When align is defined, tiles are anchored to align instead of to
+        // old_min, so knowing that the factor divides the extent is not
+        // enough to prove no boundary guard is needed: we additionally need
+        // the tiling anchored at align to line up with the tiling anchored
+        // at old_min, i.e. old_min and align must be congruent mod factor.
+        bool alignment_matches_old_min = !split.align.defined() ||
+                                         is_const_zero(simplify((old_min - split.align) % split.factor));
+
         if ((iter != dim_extent_alignment.end()) &&
-            is_const_zero(simplify(iter->second % split.factor))) {
+            is_const_zero(simplify(iter->second % split.factor)) &&
+            alignment_matches_old_min) {
             // We have proved that the split factor divides the
             // old extent. No need to adjust the base or add an if
             // statement.
@@ -58,14 +74,16 @@ vector<ApplySplitResult> apply_split(const Split &split, const string &prefix,
             // extent divides the factor. Use predication to guard
             // the calls and/or provides.
 
-            // Bounds inference has trouble exploiting an if
-            // condition. We'll directly tell it that the loop
-            // variable is bounded above by the original loop max by
-            // replacing the variable with a promise-clamped version
-            // of it. We don't also use the original loop min because
-            // it needlessly complicates the expressions and doesn't
-            // actually communicate anything new.
-            Expr guarded = promise_clamped(old_var, old_var, old_max);
+            Expr guarded;
+            if (split.align.defined()) {
+                // Because the un-rebased base block can start before old_min,
+                // we must clamp both the minimum and maximum boundaries.
+                guarded = promise_clamped(old_var, old_min, old_max);
+            } else {
+                // Legacy: structurally guaranteed to be >= old_min
+                guarded = promise_clamped(old_var, old_var, old_max);
+            }
+
             string guarded_var_name = prefix + split.old_var + ".guarded";
             Expr guarded_var = Variable::make(Int(32), guarded_var_name);
 
@@ -76,8 +94,6 @@ vector<ApplySplitResult> apply_split(const Split &split, const string &prefix,
                 predicate_type = ApplySplitResult::Predicate;
                 break;
             case TailStrategy::Predicate:
-                // This is identical to GuardWithIf, but maybe it makes
-                // sense to keep it anyways?
                 substitution_type = ApplySplitResult::Substitution;
                 predicate_type = ApplySplitResult::Predicate;
                 break;
@@ -97,31 +113,109 @@ vector<ApplySplitResult> apply_split(const Split &split, const string &prefix,
             // for the guarded version.
             result.emplace_back(prefix + split.old_var, guarded_var, substitution_type);
             result.emplace_back(guarded_var_name, guarded, ApplySplitResult::LetStmt);
-            result.emplace_back(likely(old_var <= old_max), predicate_type);
+
+            Expr guard_cond = likely(old_var <= old_max);
+            if (split.align.defined()) {
+                guard_cond = likely(old_var >= old_min && old_var <= old_max);
+            }
+            result.emplace_back(guard_cond, predicate_type);
 
         } else if (tail == TailStrategy::ShiftInwards) {
             // Adjust the base downwards to not compute off the
             // end of the realization.
 
-            // We'll only mark the base as likely (triggering a loop
-            // partition) if we're at or inside the innermost
-            // non-trivial loop.
             base = likely_if_innermost(base);
-            base = Min::make(base, old_max + (1 - split.factor));
+            if (split.align.defined()) {
+                base = Max::make(base, old_min - split.align);
+                base = Min::make(base, old_max + (1 - split.factor) - split.align);
+            } else {
+                base = Min::make(base, old_max + (1 - split.factor));
+            }
         } else if (tail == TailStrategy::ShiftInwardsAndBlend) {
+            // Unclamped base, saved before the Min/Max below adjust it. Used
+            // to figure out how much (if at all) the boundary tile got
+            // shifted, so we know which elements of it are redundant with a
+            // neighboring tile and must be masked out rather than
+            // recomputed (to avoid double-counting in a reduction).
             Expr old_base = base;
             base = likely(base);
-            base = Min::make(base, old_max + (1 - split.factor));
-            // Make a mask which will be a loop invariant if inner gets
-            // vectorized, and apply it if we're in the tail.
-            Expr unwanted_elems = (-old_extent) % split.factor;
-            Expr mask = inner >= unwanted_elems;
-            mask = select(base == old_base, likely(const_true()), mask);
+            Expr zero_based_inner = split.align.defined() ? (inner - split.align) : inner;
+            Expr mask;
+            if (split.align.defined()) {
+                // Because base is anchored to align instead of old_min, the
+                // boundary tile can now be shifted at either end (whereas
+                // without align only the max end is reachable, since base
+                // is structurally >= old_min already). Elements shifted in
+                // from the low end overlap the tile above (mask out the
+                // last shift_low of them); elements shifted in from the
+                // high end overlap the tile below (mask out the first
+                // shift_high of them).
+                Expr low_bound = old_min - split.align;
+                Expr high_bound = old_max + (1 - split.factor) - split.align;
+                Expr shift_low = low_bound - old_base;
+                Expr shift_high = old_base - high_bound;
+                base = Max::make(base, low_bound);
+                base = Min::make(base, high_bound);
+                Expr mask_low = zero_based_inner < split.factor - shift_low;
+                Expr mask_high = zero_based_inner >= shift_high;
+                mask = select(old_base < low_bound, mask_low,
+                              select(old_base > high_bound, mask_high, likely(const_true())));
+            } else {
+                // Without align, base is structurally >= old_min (outer
+                // starts at 0), so only the max end can ever be shifted.
+                base = Min::make(base, old_max + (1 - split.factor));
+                Expr unwanted_elems = (-old_extent) % split.factor;
+                mask = zero_based_inner >= unwanted_elems;
+                mask = select(base == old_base, likely(const_true()), mask);
+            }
             result.emplace_back(mask, ApplySplitResult::BlendProvides);
         } else if (tail == TailStrategy::RoundUpAndBlend) {
-            Expr unwanted_elems = (-old_extent) % split.factor;
-            Expr mask = inner < split.factor - unwanted_elems;
-            mask = select(outer < outer_max, likely(const_true()), mask);
+            Expr zero_based_inner = split.align.defined() ? (inner - split.align) : inner;
+            Expr mask;
+            if (split.align.defined()) {
+                // Unlike ShiftInwardsAndBlend, the max end is intentionally
+                // left unclamped here (RoundUp relies on padding, not on
+                // shifting, to handle overrun at the max end) -- but the min
+                // end still needs clamping: align can make the min-end tile
+                // start before old_min, and unlike ShiftInwards/blend at the
+                // max end, there's no padding below old_min to absorb an
+                // underrun into, so it has to be prevented outright.
+                //
+                // The mask below compares old_base (the unclamped base)
+                // against low_bound/high_bound directly, rather than
+                // comparing outer against outer_min/outer_max: the latter
+                // needs loop partitioning to split the loop into three
+                // pieces (prologue/steady-state/epilogue) to stay correct,
+                // and partition_loops doesn't reliably do that here when
+                // both boundaries are data-dependent, silently dropping the
+                // last tile. Comparing old_base against the bounds directly
+                // is correct regardless of how (or whether) the loop gets
+                // partitioned, matching the approach already proven correct
+                // above for ShiftInwardsAndBlend.
+                Expr old_base = base;
+                Expr low_bound = old_min - split.align;
+                Expr high_bound = old_max + (1 - split.factor) - split.align;
+                Expr shift_low = low_bound - old_base;
+                Expr shift_high = old_base - high_bound;
+                base = Max::make(likely(base), low_bound);
+                // The min end is clamped (shifted forward), so its overlap
+                // is with the tile *above* -- same geometry as
+                // ShiftInwardsAndBlend, mask out the trailing shift_low
+                // elements. The max end is left unclamped, so shift_high
+                // counts a genuine overrun past old_max with no
+                // neighboring tile to defer to -- mask out the trailing
+                // shift_high elements too (the opposite convention from
+                // ShiftInwardsAndBlend's clamped max end, which instead
+                // masks out the *leading* elements of a shifted-back tile).
+                Expr mask_low = zero_based_inner < split.factor - shift_low;
+                Expr mask_high = zero_based_inner < split.factor - shift_high;
+                mask = select(old_base < low_bound, mask_low,
+                              select(old_base > high_bound, mask_high, likely(const_true())));
+            } else {
+                Expr unwanted_elems = (-old_extent) % split.factor;
+                Expr fresh_high = zero_based_inner < split.factor - unwanted_elems;
+                mask = select(outer < outer_max, likely(const_true()), fresh_high);
+            }
             result.emplace_back(mask, ApplySplitResult::BlendProvides);
         } else {
             internal_assert(tail == TailStrategy::RoundUp);
@@ -173,12 +267,22 @@ vector<std::pair<string, Expr>> compute_loop_bounds_after_split(const Split &spl
     Expr old_var_min = Variable::make(Int(32), prefix + split.old_var + ".loop_min");
     switch (split.split_type) {
     case Split::SplitVar: {
-        Expr inner_extent = split.factor;
-        Expr outer_extent = (old_var_max - old_var_min + split.factor) / split.factor;
-        let_stmts.emplace_back(prefix + split.inner + ".loop_min", 0);
-        let_stmts.emplace_back(prefix + split.inner + ".loop_max", inner_extent - 1);
-        let_stmts.emplace_back(prefix + split.outer + ".loop_min", 0);
-        let_stmts.emplace_back(prefix + split.outer + ".loop_max", outer_extent - 1);
+        if (split.align.defined()) {
+            Expr align = split.align;
+            Expr outer_min = (old_var_min - align) / split.factor;
+            Expr outer_max = (old_var_max - align) / split.factor;
+            let_stmts.emplace_back(prefix + split.inner + ".loop_min", align);
+            let_stmts.emplace_back(prefix + split.inner + ".loop_max", align + split.factor - 1);
+            let_stmts.emplace_back(prefix + split.outer + ".loop_min", outer_min);
+            let_stmts.emplace_back(prefix + split.outer + ".loop_max", outer_max);
+        } else {
+            Expr inner_extent = split.factor;
+            Expr outer_extent = (old_var_max - old_var_min + split.factor) / split.factor;
+            let_stmts.emplace_back(prefix + split.inner + ".loop_min", 0);
+            let_stmts.emplace_back(prefix + split.inner + ".loop_max", inner_extent - 1);
+            let_stmts.emplace_back(prefix + split.outer + ".loop_min", 0);
+            let_stmts.emplace_back(prefix + split.outer + ".loop_max", outer_extent - 1);
+        }
     } break;
     case Split::FuseVars: {
         // Define bounds on the fused var using the bounds on the inner and outer