Aligned splits, take 2.

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/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.
@@ -62,10 +78,19 @@ vector<ApplySplitResult> apply_split(const Split &split, const string &prefix,
             // 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);
+            // of it.
+            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 {
+                // We don't also use the original loop min because
+                // it needlessly complicates the expressions and doesn't
+                // actually communicate anything new.
+                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 +101,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,36 +120,123 @@ 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 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 = inner < split.factor - shift_low;
+                Expr mask_high = 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 = 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 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 = inner < split.factor - shift_low;
+                Expr mask_high = 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 = 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);
         }
 
+        // Add align back in last, after all tail-strategy clamping/masking is
+        // done in terms of the unaligned base: this keeps align as a bare
+        // top-level addend in the final expressions (so e.g. it can still
+        // cancel algebraically against a matching subtraction elsewhere)
+        // rather than being smeared into a Max/Min-clamped expression, while
+        // letting the inner loop variable itself range over the simple,
+        // often-constant [0, factor) instead of [align, align + factor).
+        if (split.align.defined()) {
+            base = base + split.align;
+        }
+
         // Define the original variable as the base value computed above plus the inner loop variable.
         result.emplace_back(old_var_name, base_var + inner, ApplySplitResult::LetStmt);
         result.emplace_back(base_name, base, ApplySplitResult::LetStmt);
@@ -173,12 +283,19 @@ 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);
+        let_stmts.emplace_back(prefix + split.inner + ".loop_max", split.factor - 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.outer + ".loop_min", outer_min);
+            let_stmts.emplace_back(prefix + split.outer + ".loop_max", outer_max);
+        } else {
+            Expr outer_max = (old_var_max - old_var_min) / split.factor;
+            let_stmts.emplace_back(prefix + split.outer + ".loop_min", 0);
+            let_stmts.emplace_back(prefix + split.outer + ".loop_max", outer_max);
+        }
     } break;
     case Split::FuseVars: {
         // Define bounds on the fused var using the bounds on the inner and outer

src/Deserialization.cpp

@@ -1152,6 +1152,7 @@ Split Deserializer::deserialize_split(const Serialize::Split *split) {
     const auto exact = split->exact();
     const auto tail = deserialize_tail_strategy(split->tail());
     const auto split_type = deserialize_split_type(split->split_type());
+    const auto align = deserialize_expr(split->align_type(), split->align());
     auto hl_split = Split();
     hl_split.old_var = old_var;
     hl_split.outer = outer;
@@ -1160,6 +1161,7 @@ Split Deserializer::deserialize_split(const Serialize::Split *split) {
     hl_split.exact = exact;
     hl_split.tail = tail;
     hl_split.split_type = split_type;
+    hl_split.align = align;
     return hl_split;
 }
 

src/Func.cpp

@@ -1103,9 +1103,9 @@ Func Stage::rfactor(const vector<pair<RVar, Var>> &preserved) {
     return intm;
 }
 
-void Stage::split(const string &old, const string &outer, const string &inner, const Expr &factor_arg, bool exact, TailStrategy tail) {
+void Stage::split(const string &old, const string &outer, const string &inner, const Expr &factor_arg, const Expr &align_arg, bool exact, TailStrategy tail) {
     debug(4) << "In schedule for " << name() << ", split " << old << " into "
-             << outer << " and " << inner << " with factor of " << factor_arg << "\n";
+             << outer << " and " << inner << " with factor of " << factor_arg << " and align " << align_arg << "\n";
 
     user_assert(factor_arg.defined())
         << "In schedule for " << name() << ", split factor for splitting "
@@ -1115,6 +1115,14 @@ void Stage::split(const string &old, const string &outer, const string &inner, c
         << old << " has type " << factor_arg.type()
         << ", which is not representable as int32.\n";
     Expr factor = cast<int32_t>(factor_arg);
+    Expr align;
+    if (align_arg.defined()) {
+        user_assert(Int(32).can_represent(align_arg.type()))
+            << "In schedule for " << name() << ", split align for splitting "
+            << old << " has type " << align_arg.type()
+            << ", which is not representable as int32.\n";
+        align = cast<int32_t>(align_arg);
+    }
 
     vector<Dim> &dims = definition.schedule().dims();
 
@@ -1318,11 +1326,15 @@ void Stage::split(const string &old, const string &outer, const string &inner, c
     }
 
     // Add the split to the splits list
-    Split split = {old_name, outer_name, inner_name, factor, exact, tail, Split::SplitVar};
+    Split split = {old_name, outer_name, inner_name, factor, align, exact, tail, Split::SplitVar};
     definition.schedule().splits().push_back(split);
 }
 
-Stage &Stage::split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, TailStrategy tail) {
+void Stage::split(const std::string &old, const std::string &outer, const std::string &inner, const Expr &factor, bool exact, TailStrategy tail) {
+    split(old, outer, inner, factor, Expr(), exact, tail);
+}
+
+Stage &Stage::split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, const Expr &align, TailStrategy tail) {
     definition.schedule().touched() = true;
     if (old.is_rvar) {
         user_assert(outer.is_rvar) << "Can't split RVar " << old.name() << " into Var " << outer.name() << "\n";
@@ -1331,7 +1343,13 @@ Stage &Stage::split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVa
         user_assert(!outer.is_rvar) << "Can't split Var " << old.name() << " into RVar " << outer.name() << "\n";
         user_assert(!inner.is_rvar) << "Can't split Var " << old.name() << " into RVar " << inner.name() << "\n";
     }
-    split(old.name(), outer.name(), inner.name(), factor, old.is_rvar, tail);
+    split(old.name(), outer.name(), inner.name(), factor, align, old.is_rvar, tail);
+    return *this;
+}
+
+Stage &Stage::split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, TailStrategy tail) {
+    definition.schedule().touched() = true;
+    split(old.name(), outer.name(), inner.name(), factor, Expr(), old.is_rvar, tail);
     return *this;
 }
 
@@ -1413,7 +1431,7 @@ Stage &Stage::fuse(const VarOrRVar &inner, const VarOrRVar &outer, const VarOrRV
     set_dim_type(fused, dims[inner_pos].for_type);
 
     // Add the fuse to the splits list
-    Split split = {fused_name, outer_name, inner_name, Expr(), true, TailStrategy::RoundUp, Split::FuseVars};
+    Split split = {fused_name, outer_name, inner_name, Expr(), Expr(), true, TailStrategy::RoundUp, Split::FuseVars};
     definition.schedule().splits().push_back(split);
     return *this;
 }
@@ -1664,7 +1682,7 @@ Stage &Stage::rename(const VarOrRVar &old_var, const VarOrRVar &new_var) {
     }
 
     if (!found) {
-        Split split = {old_name, new_name, "", 1, old_var.is_rvar, TailStrategy::RoundUp, Split::RenameVar};
+        Split split = {old_name, new_name, "", 1, Expr(), old_var.is_rvar, TailStrategy::RoundUp, Split::RenameVar};
         definition.schedule().splits().push_back(split);
     }
 
@@ -2545,6 +2563,12 @@ Func &Func::split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar
     return *this;
 }
 
+Func &Func::split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, const Expr &align, TailStrategy tail) {
+    invalidate_cache();
+    Stage(func, func.definition(), 0).split(old, outer, inner, factor, align, tail);
+    return *this;
+}
+
 Func &Func::fuse(const VarOrRVar &inner, const VarOrRVar &outer, const VarOrRVar &fused) {
     invalidate_cache();
     Stage(func, func.definition(), 0).fuse(inner, outer, fused);