diff --git a/python_bindings/halide/src/halide_/PyScheduleMethods.h b/python_bindings/halide/src/halide_/PyScheduleMethods.h index f528af886dff..7e585690e1e3 100644 --- a/python_bindings/halide/src/halide_/PyScheduleMethods.h +++ b/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")) diff --git a/src/ApplySplit.cpp b/src/ApplySplit.cpp index ddb9bc1098c5..14c77ae772d4 100644 --- a/src/ApplySplit.cpp +++ b/src/ApplySplit.cpp @@ -23,10 +23,17 @@ vector 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 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 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 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 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> 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 diff --git a/src/Deserialization.cpp b/src/Deserialization.cpp index f7f8566326db..be75ec77d8da 100644 --- a/src/Deserialization.cpp +++ b/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; } diff --git a/src/Func.cpp b/src/Func.cpp index 468188530c67..6391eb49297a 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -1103,9 +1103,9 @@ Func Stage::rfactor(const vector> &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(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(align_arg); + } vector &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); diff --git a/src/Func.h b/src/Func.h index 4df562e272ca..dc9dcf62a04e 100644 --- a/src/Func.h +++ b/src/Func.h @@ -81,6 +81,8 @@ class Stage { void set_dim_device_api(const VarOrRVar &var, DeviceAPI device_api); void split(const std::string &old, const std::string &outer, const std::string &inner, const Expr &factor, bool exact, TailStrategy tail); + void split(const std::string &old, const std::string &outer, const std::string &inner, + const Expr &factor, const Expr &align, bool exact, TailStrategy tail); void remove(const std::string &var); const std::vector &storage_dims() const { @@ -365,6 +367,7 @@ class Stage { // @{ Stage &split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, TailStrategy tail = TailStrategy::Auto); + Stage &split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, const Expr &align, TailStrategy tail = TailStrategy::Auto); Stage &fuse(const VarOrRVar &inner, const VarOrRVar &outer, const VarOrRVar &fused); Stage &serial(const VarOrRVar &var); Stage ¶llel(const VarOrRVar &var); @@ -1519,6 +1522,41 @@ class Func { * factor does not provably divide the extent. */ Func &split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, TailStrategy tail = TailStrategy::Auto); + /** A version of split() that additionally takes a runtime-valued + * 'align' Expr. This version anchors the inner-loop's iterations + * to absolute coordinates, instead of to Halide's inferred loop + * bounds. + * + * As such, in absolute coordinates, the inner-loop boundaries fall + * at ``align``, ``align + factor``, ``align + 2*factor``, and so on. + * The inner dimension still iterates over ``[0, factor-1]``, same as + * an unaligned split. The difference is how the original loop Var + * is reconstructed from the outer and inner loop Vars. + * This may increase the number of iterations over the outer + * loop by 1 compared to an unaligned split. + * + * This is useful when an algorithm selects between cases using an + * expression like ``(x - offset) % factor``, where 'offset' is a + * value only known at runtime (e.g. a Param). Passing that same + * 'offset' as 'align' makes ``(x - offset) % factor`` a + * compile-time constant on each unrolled iteration of the inner + * loop, so that a mux() indexed by it can be resolved statically + * instead of compiling to a runtime select: + \code + Var x, xo, xi; + Param offset; + f(x) = mux((x - offset) % 4, {a(x), b(x), c(x), d(x)}); + f.split(x, xo, xi, 4, offset, TailStrategy::GuardWithIf) + .unroll(xi); + \endcode + * Without 'align', the compiler can't tell at compile time which of + * the four mux() cases applies to a given unrolled value of 'xi', + * because that depends on the runtime value of 'offset'. With it, + * ``(x - offset) % 4`` simplifies to a distinct compile-time + * constant for each unrolled value of 'xi', and each mux() call + * collapses to its selected case. */ + Func &split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, const Expr &align, TailStrategy tail = TailStrategy::Auto); + /** Join two dimensions into a single fused dimension. The fused dimension * covers the product of the extents of the inner and outer dimensions * given. The loop type (e.g. parallel, vectorized) of the resulting fused diff --git a/src/Schedule.cpp b/src/Schedule.cpp index 948233112b7c..77f27d8e89a6 100644 --- a/src/Schedule.cpp +++ b/src/Schedule.cpp @@ -340,6 +340,9 @@ struct StageScheduleContents { if (s.factor.defined()) { s.factor = mutator(s.factor); } + if (s.align.defined()) { + s.align = mutator(s.align); + } } for (PrefetchDirective &p : prefetches) { if (p.offset.defined()) { @@ -702,6 +705,9 @@ void StageSchedule::accept(IRVisitor *visitor) const { if (s.factor.defined()) { s.factor.accept(visitor); } + if (s.align.defined()) { + s.align.accept(visitor); + } } for (const PrefetchDirective &p : prefetches()) { if (p.offset.defined()) { diff --git a/src/Schedule.h b/src/Schedule.h index ba3d1eea5ca3..7bfa92981ac1 100644 --- a/src/Schedule.h +++ b/src/Schedule.h @@ -334,6 +334,8 @@ struct ReductionVariable; struct Split { std::string old_var, outer, inner; Expr factor; + Expr align; // If defined, the inner var loops over [align, + // align + factor - 1] instead of [0, factor - 1]. bool exact; // Is it required that the factor divides the extent // of the old var. True for splits of RVars. Forces // tail strategy to be GuardWithIf. diff --git a/src/Serialization.cpp b/src/Serialization.cpp index 2dd7bf4f33aa..36ea9d84984f 100644 --- a/src/Serialization.cpp +++ b/src/Serialization.cpp @@ -1259,10 +1259,12 @@ Offset Serializer::serialize_split(FlatBufferBuilder &builder, const auto exact = split.exact; const auto tail_serialized = serialize_tail_strategy(split.tail); const auto split_type_serialized = serialize_split_type(split.split_type); + const auto align_serialized = serialize_expr(builder, split.align); return Serialize::CreateSplit(builder, old_var_serialized, outer_serialized, inner_serialized, factor_serialized.first, factor_serialized.second, - exact, tail_serialized, split_type_serialized); + exact, tail_serialized, split_type_serialized, + align_serialized.first, align_serialized.second); } Offset Serializer::serialize_dim(FlatBufferBuilder &builder, const Dim &dim) { diff --git a/src/SimplifyCorrelatedDifferences.cpp b/src/SimplifyCorrelatedDifferences.cpp index f7d1049ea4f0..0c41b498292d 100644 --- a/src/SimplifyCorrelatedDifferences.cpp +++ b/src/SimplifyCorrelatedDifferences.cpp @@ -53,6 +53,13 @@ class PartiallyCancelDifferences : public IRMutator { rewrite(min(x, y) - max(x, z), min(min(x, y) - max(x, z), 0)) || rewrite(max(x, y) - min(x, z), max(max(x, y) - min(x, z), 0)) || + // min(x + c0, y) - max(x, z) + // = min(x, y - c0) - max(x, z) + c0 + // = min(min(x, y - c0) - max(x, z), 0) + c0 + // = min(min(x, y - c0) - max(x, z) + c0, c0) + // = min(min(x + c0, y) - max(x, z), c0) + rewrite(min(x + c0, y) - max(x, z), min(min(x, y) - max(x, z), c0)) || + rewrite(min(x + c0, y) - select(z, min(x, y) + c1, x), select(z, (max(min(y - x, c0), 0) - c1), min(y - x, c0)), c0 > 0) || rewrite(min(y, x + c0) - select(z, min(y, x) + c1, x), select(z, (max(min(y - x, c0), 0) - c1), min(y - x, c0)), c0 > 0) || diff --git a/src/Simplify_Add.cpp b/src/Simplify_Add.cpp index a07ad1b4464b..8ab0fb2c61c0 100644 --- a/src/Simplify_Add.cpp +++ b/src/Simplify_Add.cpp @@ -80,6 +80,16 @@ Expr Simplify::visit(const Add *op, ExprInfo *info) { rewrite((c0 - x) + y, (y - x) + c0) || rewrite(max(x, y * c0 + z) + (u - y) * c0, max(x - y * c0, z) + u * c0) || + // Collect a repeated term across a nested sum, so that facts about + // it become visible to the rules and to modulus_remainder. For + // example, x + y + y is even in y, but only once written as x + y*2. + rewrite((x + y) + y, x + y * 2) || + rewrite((y + x) + y, x + y * 2) || + rewrite((x + (y + z)) + z, z * 2 + (x + y)) || + rewrite((x + (z + y)) + z, z * 2 + (x + y)) || + rewrite(((y + z) + x) + z, z * 2 + (x + y)) || + rewrite(((z + y) + x) + z, z * 2 + (x + y)) || + rewrite((x - y) + y, x) || rewrite(x + (y - x), y) || @@ -198,9 +208,12 @@ Expr Simplify::visit(const Add *op, ExprInfo *info) { rewrite((x / w) * w + (z + x % w), select(w == 0, 0, x) + z) || rewrite(x / 2 + x % 2, (x + 1) / 2) || + rewrite((0 - (x % 2)) / 2 * 2 + (x % 2), 0 - (x % 2)) || + rewrite(x + ((c0 - x) / c1) * c1, c0 - ((c0 - x) % c1), c1 > 0) || rewrite(x + ((c0 - x) / c1 + y) * c1, y * c1 - ((c0 - x) % c1) + c0, c1 > 0) || rewrite(x + (y + (c0 - x) / c1) * c1, y * c1 - ((c0 - x) % c1) + c0, c1 > 0) || + rewrite(((0 - x) / c0) + ((x % c0 + c1) / c0), fold(c1 / c0) - (x / c0), c0 > 0 && (c1 + 1) % c0 == 0) || false)))) { return mutate(rewrite.result, info); diff --git a/src/Simplify_Div.cpp b/src/Simplify_Div.cpp index 4098f6f027e7..c042abd1f036 100644 --- a/src/Simplify_Div.cpp +++ b/src/Simplify_Div.cpp @@ -95,6 +95,7 @@ Expr Simplify::visit(const Div *op, ExprInfo *info) { rewrite(max((x * c0), c1) / c2, max(x * fold(c0 / c2), fold(c1 / c2)), c0 % c2 == 0 && c2 > 0) || rewrite((x * c0 + y) / c1, y / c1 + x * fold(c0 / c1), c0 % c1 == 0 && c1 > 0) || + rewrite((x * c0 - y) / c1, (0 - y) / c1 + x * fold(c0 / c1), c0 % c1 == 0 && c1 > 0) || rewrite((x * c0 - y) / c0, x + (0 - y) / c0) || rewrite((x * c1 - y) / c0, (0 - y) / c0 - x, c0 + c1 == 0) || rewrite((y + x * c0) / c1, y / c1 + x * fold(c0 / c1), c0 % c1 == 0 && c1 > 0) || @@ -131,6 +132,8 @@ Expr Simplify::visit(const Div *op, ExprInfo *info) { rewrite((w + (z + (x * c0 + y))) / c1, (y + z + w) / c1 + x * fold(c0 / c1), c0 % c1 == 0 && c1 > 0) || rewrite((w + (z + (y + x * c0))) / c1, (y + z + w) / c1 + x * fold(c0 / c1), c0 % c1 == 0 && c1 > 0) || + rewrite(((0 - x) - x) / 2, -x) || + /** In (x + c0) / c1, when can we pull the constant addition out of the numerator? An obvious answer is the constant is a multiple of the denominator, but diff --git a/src/Simplify_Exprs.cpp b/src/Simplify_Exprs.cpp index c19fa2e7fed8..07512e9186d0 100644 --- a/src/Simplify_Exprs.cpp +++ b/src/Simplify_Exprs.cpp @@ -214,8 +214,28 @@ Expr Simplify::visit(const VectorReduce *op, ExprInfo *info) { x + max(y * (arg_lanes - 1), 0) <= z) || rewrite(h_and(broadcast(x, arg_lanes) < ramp(y, z, arg_lanes), 1), x < y + min(z * (arg_lanes - 1), 0)) || - rewrite(h_and(broadcast(x, arg_lanes) < ramp(y, z, arg_lanes), 1), + rewrite(h_and(broadcast(x, arg_lanes) <= ramp(y, z, arg_lanes), 1), x <= y + min(z * (arg_lanes - 1), 0)) || + + // The "all lanes of a ramp lie within [lo, hi]" check loop + // partitioning builds is a chain of these same ramp/broadcast + // comparisons ANDed together (a lower-bound comparison against a + // ramp's minimum lane, an upper-bound comparison against its + // maximum lane, for one or more ramps sharing a stride). Rather + // than special-case every clause count, peel one comparison off + // the tail of the && chain at a time: the four rules above already + // reduce that comparison to a scalar; h_and(w, 1) on what's left + // recurses into this same case, so it either peels the next + // clause or -- once w is down to a single comparison -- hits one + // of the four rules above directly. + rewrite(h_and(w && (ramp(x, y, arg_lanes) < broadcast(z, arg_lanes)), 1), + h_and(w, 1) && (x + max(y * (arg_lanes - 1), 0) < z)) || + rewrite(h_and(w && (ramp(x, y, arg_lanes) <= broadcast(z, arg_lanes)), 1), + h_and(w, 1) && (x + max(y * (arg_lanes - 1), 0) <= z)) || + rewrite(h_and(w && (broadcast(x, arg_lanes) < ramp(y, z, arg_lanes)), 1), + h_and(w, 1) && (x < y + min(z * (arg_lanes - 1), 0))) || + rewrite(h_and(w && (broadcast(x, arg_lanes) <= ramp(y, z, arg_lanes)), 1), + h_and(w, 1) && (x <= y + min(z * (arg_lanes - 1), 0))) || false) { return mutate(rewrite.result, info); } @@ -237,7 +257,7 @@ Expr Simplify::visit(const VectorReduce *op, ExprInfo *info) { x + min(y * (arg_lanes - 1), 0) <= z) || rewrite(h_or(broadcast(x, arg_lanes) < ramp(y, z, arg_lanes), 1), x < y + max(z * (arg_lanes - 1), 0)) || - rewrite(h_or(broadcast(x, arg_lanes) < ramp(y, z, arg_lanes), 1), + rewrite(h_or(broadcast(x, arg_lanes) <= ramp(y, z, arg_lanes), 1), x <= y + max(z * (arg_lanes - 1), 0)) || false) { return mutate(rewrite.result, info); diff --git a/src/Simplify_Mod.cpp b/src/Simplify_Mod.cpp index 7e5232da0975..649073494163 100644 --- a/src/Simplify_Mod.cpp +++ b/src/Simplify_Mod.cpp @@ -59,6 +59,14 @@ Expr Simplify::visit(const Mod *op, ExprInfo *info) { rewrite((x * c0 - y) % c1, (-y) % c1, c0 % c1 == 0) || rewrite((y - x * c0) % c1, y % c1, c0 % c1 == 0) || rewrite((x - y) % 2, (x + y) % 2) || // Addition and subtraction are the same modulo 2, because -1 == 1 + rewrite((((x * c0) + y) - z) % c1, (y - z) % c1, c0 % c1 == 0) || + rewrite((((x * c0) + y) + z) % c1, (y + z) % c1, c0 % c1 == 0) || + rewrite((((x * c0) - y) - z) % c1, (-y - z) % c1, c0 % c1 == 0) || + rewrite((((x * c0) - y) + z) % c1, (z - y) % c1, c0 % c1 == 0) || + + rewrite((x + y + y) % 2, x % 2) || + rewrite(((y + x) + y) % 2, x % 2) || + rewrite((((z + y) + x) + y) % 2, (z + x) % 2) || rewrite(ramp(x, c0, c2) % broadcast(c1, c2), broadcast(x, c2) % broadcast(c1, c2), (c0 % c1 == 0)) || rewrite(ramp(x, c0, lanes) % broadcast(c1, lanes), ramp(x % c1, c0, lanes), diff --git a/src/halide_ir.fbs b/src/halide_ir.fbs index 4bba4bb79a8f..3e81a44fb3b1 100644 --- a/src/halide_ir.fbs +++ b/src/halide_ir.fbs @@ -568,6 +568,7 @@ table Split { exact: bool; tail: TailStrategy; split_type: SplitType; + align: Expr; } enum DimType: ubyte { diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index d88c13fa177f..93048b3bed21 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -339,6 +339,12 @@ tests( specialize_to_gpu.cpp specialize_trim_condition.cpp spirv_ir.cpp + split_aligned.cpp + split_aligned_2d.cpp + split_aligned_2d_3x3.cpp + split_aligned_nested.cpp + split_aligned_partition.cpp + split_aligned_reduction.cpp split_by_non_factor.cpp split_factor_type.cpp split_fuse_rvar.cpp @@ -481,6 +487,10 @@ tests( random.cpp reorder_rvars.cpp rfactor.cpp + rfactor_split_aligned.cpp + rfactor_split_aligned_2d.cpp + rfactor_split_aligned_nested.cpp + rfactor_split_aligned_phases.cpp ring_buffer.cpp stream_compaction.cpp thread_safety.cpp diff --git a/test/correctness/rfactor_split_aligned.cpp b/test/correctness/rfactor_split_aligned.cpp new file mode 100644 index 000000000000..19f51dfa63bc --- /dev/null +++ b/test/correctness/rfactor_split_aligned.cpp @@ -0,0 +1,94 @@ +#include "Halide.h" +#include + +// rfactor() eagerly applies any splits present on the RVar(s) it's given (see +// Stage::rfactor / project_rdom in Func.cpp), so it needs to tolerate splits +// that carry an alignment (Stage::split's 'align' argument) just as well as +// ordinary ones. This test factors the *outer* half of an aligned split of +// the reduction variable out into a parallel-reducible intermediate Func, +// while unrolling the *inner* (aligned) half in the reducing computation. +// Because the inner half is not itself preserved by rfactor(), it keeps the +// exact loop bounds computed by compute_loop_bounds_after_split (rather than +// being re-derived by general bounds inference), so unrolling it still lets +// the compiler resolve the runtime-offset mux() to a compile-time constant +// per lane, exactly as it does without rfactor in split_aligned.cpp. + +using namespace Halide; +using namespace Halide::Internal; + +class MuxCounter : public IRVisitor { + using IRVisitor::visit; + + void visit(const Call *op) override { + IRVisitor::visit(op); + if (op->is_intrinsic(Call::IntrinsicOp::mux)) { + mux_count++; + } + } + +public: + int mux_count{0}; +}; + +int main(int argc, char **argv) { + Var x{"x"}; + Func f{"f"}; + RDom r(0, 40, "r"); + Param offset{"offset"}; + offset.set_range(0, 3); + + f(x) = 0; + f(x) += mux((r - offset) % 4, {r + x, r * r + x, 2 * r + x, -r * (r + 1) + x}); + + RVar ro{"ro"}, ri{"ri"}; + f.update(0) + .split(r, ro, ri, 4, offset, TailStrategy::GuardWithIf) + .unroll(ri); + + Var u{"u"}; + Func intm = f.update(0).rfactor(ro, u); + intm.compute_root(); + intm.update(0).parallel(u); + + Module module = f.compile_to_module({offset}); + MuxCounter checker; + for (const LoweredFunc &lf : module.functions()) { + lf.body.accept(&checker); + } + if (checker.mux_count != 0) { + printf("Expected 0 muxes (the aligned+unrolled inner split var should " + "resolve the mux at compile time even after rfactor): %d\n", + checker.mux_count); + return 1; + } + + for (int off = 0; off < 4; off++) { + printf("Testing runtime alignment: %d\n", off); + offset.set(off); + Buffer im = f.realize({10}); + for (int x = 0; x < 10; x++) { + int expected = 0; + for (int r = 0; r < 40; r++) { + int selector = (4 + r - off) % 4; + int term; + if (selector == 0) { + term = r + x; + } else if (selector == 1) { + term = r * r + x; + } else if (selector == 2) { + term = 2 * r + x; + } else { + term = -r * (r + 1) + x; + } + expected += term; + } + if (im(x) != expected) { + printf("im(%d) = %d instead of %d (offset: %d)\n", x, im(x), expected, off); + return 1; + } + } + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/rfactor_split_aligned_2d.cpp b/test/correctness/rfactor_split_aligned_2d.cpp new file mode 100644 index 000000000000..9120c26a4578 --- /dev/null +++ b/test/correctness/rfactor_split_aligned_2d.cpp @@ -0,0 +1,101 @@ +#include "Halide.h" +#include + +// A 2D companion to rfactor_split_aligned.cpp. Here rfactor() is applied to +// an RVar (r.x) that is unrelated to the one carrying the aligned split +// (r.y), which is the more common pattern in practice: factor out one +// reduction dimension for parallel/vector reduction while a separate +// dimension is scheduled with an alignment-aware split so a +// runtime-offset-dependent mux() can be resolved statically once its half of +// the split is unrolled. Since r.y's split is entirely unrelated to the +// preserved var, both halves of the split remain ordinary (non-preserved) +// reduction variables of the intermediate Func, retaining their exact +// compile-time loop bounds and so still collapsing the mux to nothing. + +using namespace Halide; +using namespace Halide::Internal; + +class MuxCounter : public IRVisitor { + using IRVisitor::visit; + + void visit(const Call *op) override { + IRVisitor::visit(op); + if (op->is_intrinsic(Call::IntrinsicOp::mux)) { + mux_count++; + } + } + +public: + int mux_count{0}; +}; + +int main(int argc, char **argv) { + Var x{"x"}, y{"y"}; + Func f{"f"}; + RDom r(0, 20, 0, 16, "r"); + Param offset{"offset"}; + offset.set_range(0, 3); + + f(x, y) = 0; + f(x, y) += mux((r.y - offset) % 4, + {r.x + r.y + x + y, + r.x * r.y + x - y, + 2 * r.x - r.y + x, + -r.x * (r.y + 1) + y}) * + select(r.x % 2 == 0, 1, -1); + + RVar ryo{"ryo"}, ryi{"ryi"}; + f.update(0) + .split(r.y, ryo, ryi, 4, offset, TailStrategy::GuardWithIf) + .unroll(ryi); + + Var u{"u"}; + Func intm = f.update(0).rfactor(r.x, u); + intm.compute_root(); + intm.update(0).parallel(u); + + Module module = f.compile_to_module({offset}); + MuxCounter checker; + for (const LoweredFunc &lf : module.functions()) { + lf.body.accept(&checker); + } + if (checker.mux_count != 0) { + printf("Expected 0 muxes: %d\n", checker.mux_count); + return 1; + } + + for (int off = 0; off < 4; off++) { + printf("Testing runtime alignment: %d\n", off); + offset.set(off); + Buffer im = f.realize({6, 6}); + for (int y = 0; y < 6; y++) { + for (int x = 0; x < 6; x++) { + int expected = 0; + for (int rx = 0; rx < 20; rx++) { + for (int ry = 0; ry < 16; ry++) { + int selector = (4 + ry - off) % 4; + int term; + if (selector == 0) { + term = rx + ry + x + y; + } else if (selector == 1) { + term = rx * ry + x - y; + } else if (selector == 2) { + term = 2 * rx - ry + x; + } else { + term = -rx * (ry + 1) + y; + } + term *= (rx % 2 == 0) ? 1 : -1; + expected += term; + } + } + if (im(x, y) != expected) { + printf("im(%d, %d) = %d instead of %d (offset: %d)\n", x, y, im(x, y), expected, off); + return 1; + } + } + } + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/rfactor_split_aligned_nested.cpp b/test/correctness/rfactor_split_aligned_nested.cpp new file mode 100644 index 000000000000..58c30b7056d5 --- /dev/null +++ b/test/correctness/rfactor_split_aligned_nested.cpp @@ -0,0 +1,114 @@ +#include "Halide.h" +#include + +// A companion to rfactor_split_aligned.cpp and split_aligned_nested.cpp: +// after r's aligned split (factor 4, aligned to offset) is rfactored on its +// outer half into a preserved pure var u, u is itself split again with a +// second, independent alignment (p2), tried with GuardWithIf, +// RoundUpAndBlend, and ShiftInwardsAndBlend. +// +// GuardWithIf and Predicate are the only tail strategies Stage::split allows +// on an RVar (splitting r itself), because RoundUp/ShiftInwards-family +// strategies would change the meaning of a reduction by recomputing or +// overrunning it -- but u is an ordinary pure Var of the intermediate +// Func's own update definition, so RoundUpAndBlend/ShiftInwardsAndBlend +// (the update-definition-safe counterparts of RoundUp/ShiftInwards) are +// legal there, and are exactly the tail strategies meant for vectorizing +// an update like this one. +// +// This combination exercises boundary handling in ApplySplit.cpp +// (apply_split's ShiftInwardsAndBlend/RoundUpAndBlend branches) that plain, +// unnested aligned splits don't: u's own old_min is not a compile-time +// constant (it comes from r's split, a function of the runtime offset +// Param), so both the low and high boundary tiles of u's split can only be +// distinguished from the interior at runtime. + +using namespace Halide; +using namespace Halide::Internal; + +class MuxCounter : public IRVisitor { + using IRVisitor::visit; + + void visit(const Call *op) override { + IRVisitor::visit(op); + if (op->is_intrinsic(Call::IntrinsicOp::mux)) { + mux_count++; + } + } + +public: + int mux_count{0}; +}; + +int main(int argc, char **argv) { + for (auto ts : {TailStrategy::GuardWithIf, TailStrategy::RoundUpAndBlend, TailStrategy::ShiftInwardsAndBlend}) { + printf("Testing tail strategy: %d\n", (int)ts); + + Var x{"x"}; + Func f{"f"}; + RDom r(0, 40, "r"); + Param offset{"offset"}; + offset.set_range(0, 3); + + f(x) = 0; + f(x) += mux((r - offset) % 4, {r + x, r * r + x, 2 * r + x, -r * (r + 1) + x}); + + RVar ro{"ro"}, ri{"ri"}; + f.update(0) + .split(r, ro, ri, 4, offset, TailStrategy::GuardWithIf) + .unroll(ri); + + Var u{"u"}, uo{"uo"}, ui{"ui"}; + Param p2{"p2"}; + p2.set_range(0, 1); + + Func intm = f.update(0).rfactor(ro, u); + intm.compute_root(); + intm.update(0) + .split(u, uo, ui, 2, p2, ts) + .vectorize(ui); + + Module module = f.compile_to_module({offset, p2}); + MuxCounter checker; + for (const LoweredFunc &lf : module.functions()) { + lf.body.accept(&checker); + } + if (checker.mux_count != 0) { + printf("Expected 0 muxes: %d\n", checker.mux_count); + return 1; + } + + for (int off = 0; off < 4; off++) { + for (int a2 = 0; a2 < 2; a2++) { + offset.set(off); + p2.set(a2); + Buffer im = f.realize({10}); + for (int x = 0; x < 10; x++) { + int expected = 0; + for (int r = 0; r < 40; r++) { + int selector = (4 + r - off) % 4; + int term; + if (selector == 0) { + term = r + x; + } else if (selector == 1) { + term = r * r + x; + } else if (selector == 2) { + term = 2 * r + x; + } else { + term = -r * (r + 1) + x; + } + expected += term; + } + if (im(x) != expected) { + printf("im(%d) = %d instead of %d (offset: %d, p2: %d, ts: %d)\n", + x, im(x), expected, off, a2, (int)ts); + return 1; + } + } + } + } + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/rfactor_split_aligned_phases.cpp b/test/correctness/rfactor_split_aligned_phases.cpp new file mode 100644 index 000000000000..2f1cfd277205 --- /dev/null +++ b/test/correctness/rfactor_split_aligned_phases.cpp @@ -0,0 +1,168 @@ +#include "Halide.h" +#include + +// A variant of rfactor_split_aligned.cpp that preserves the *inner* (aligned, +// unrolled) half of the split via rfactor() instead of the outer half, +// turning it into four separate per-phase partial-sum accumulators that get +// combined at the end. rfactor() must still produce correct results here: +// this is precisely the "does rfactor tolerate splits with an alignment" +// question, exercised in the case where the aligned split is the one being +// preserved (and therefore promoted from an RVar with exact, +// compute_loop_bounds_after_split-derived bounds to an ordinary pure Var of +// the intermediate Func, whose bounds are instead re-derived by general +// bounds inference). That promotion means the compiler can no longer read +// off the new pure var's range directly from the split; it has to prove it +// symbolically from the surrounding min/max clamps instead, which is what +// the mux_count checks below are exercising. +// +// The second case additionally makes the RDom's own extent a runtime Param +// rather than a compile-time constant, so the split's "factor provably +// divides the extent" fast path (see apply_split in ApplySplit.cpp) can't +// fire either, and everything -- the boundary guard, the alignment, and the +// mux resolution -- has to come out of the general GuardWithIf path instead. + +using namespace Halide; +using namespace Halide::Internal; + +class MuxCounter : public IRVisitor { + using IRVisitor::visit; + + void visit(const Call *op) override { + IRVisitor::visit(op); + if (op->is_intrinsic(Call::IntrinsicOp::mux)) { + mux_count++; + } + } + +public: + int mux_count{0}; +}; + +int expected_value(int r, int x, int off) { + int selector = (4 + r - off) % 4; + if (selector == 0) { + return r + x; + } else if (selector == 1) { + return r * r + x; + } else if (selector == 2) { + return 2 * r + x; + } else { + return -r * (r + 1) + x; + } +} + +int test_fixed_extent() { + Var x{"x"}; + Func f{"f"}; + RDom r(0, 40, "r"); + Param offset{"offset"}; + offset.set_range(0, 3); + + f(x) = 0; + f(x) += mux((r - offset) % 4, {r + x, r * r + x, 2 * r + x, -r * (r + 1) + x}); + + RVar ro{"ro"}, ri{"ri"}; + f.update(0) + .split(r, ro, ri, 4, offset, TailStrategy::GuardWithIf) + .unroll(ri); + + Var u{"u"}; + Func intm = f.update(0).rfactor(ri, u); + intm.compute_root(); + intm.update(0).unroll(u); + + Module module = f.compile_to_module({offset}); + MuxCounter checker; + for (const LoweredFunc &lf : module.functions()) { + lf.body.accept(&checker); + } + if (checker.mux_count != 0) { + printf("Expected 0 muxes: %d\n", checker.mux_count); + return 1; + } + + for (int off = 0; off < 4; off++) { + printf("Testing runtime alignment: %d\n", off); + offset.set(off); + Buffer im = f.realize({10}); + for (int x = 0; x < 10; x++) { + int expected = 0; + for (int r = 0; r < 40; r++) { + expected += expected_value(r, x, off); + } + if (im(x) != expected) { + printf("im(%d) = %d instead of %d (offset: %d)\n", x, im(x), expected, off); + return 1; + } + } + } + + return 0; +} + +int test_param_extent() { + Var x{"x"}; + Func f{"f"}; + Param extent{"extent"}; + RDom r(0, extent, "r"); + Param offset{"offset"}; + offset.set_range(0, 3); + + f(x) = 0; + f(x) += mux((r - offset) % 4, {r + x, r * r + x, 2 * r + x, -r * (r + 1) + x}); + + RVar ro{"ro"}, ri{"ri"}; + f.update(0) + .split(r, ro, ri, 4, offset, TailStrategy::GuardWithIf) + .unroll(ri); + + Var u{"u"}; + Func intm = f.update(0).rfactor(ri, u); + intm.compute_root(); + intm.update(0).unroll(u); + + Module module = f.compile_to_module({extent, offset}); + MuxCounter checker; + for (const LoweredFunc &lf : module.functions()) { + lf.body.accept(&checker); + } + if (checker.mux_count != 0) { + printf("Expected 0 muxes (with a Param extent): %d\n", checker.mux_count); + return 1; + } + + // 40 is a multiple of the split factor; 37 is not, so it also exercises + // the tail of the RDom's own range. + for (int ext : {40, 37}) { + for (int off = 0; off < 4; off++) { + printf("Testing runtime extent %d, alignment %d\n", ext, off); + extent.set(ext); + offset.set(off); + Buffer im = f.realize({10}); + for (int x = 0; x < 10; x++) { + int expected = 0; + for (int r = 0; r < ext; r++) { + expected += expected_value(r, x, off); + } + if (im(x) != expected) { + printf("im(%d) = %d instead of %d (extent: %d, offset: %d)\n", x, im(x), expected, ext, off); + return 1; + } + } + } + } + + return 0; +} + +int main(int argc, char **argv) { + if (test_fixed_extent()) { + return 1; + } + if (test_param_extent()) { + return 1; + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/simplify.cpp b/test/correctness/simplify.cpp index 981a00e6f0ee..538ef1982303 100644 --- a/test/correctness/simplify.cpp +++ b/test/correctness/simplify.cpp @@ -1767,6 +1767,43 @@ void check_boolean() { check(ramp(x * 8 + 5, -1, 4) < broadcast(y * 8, 4), broadcast(x < y, 4)); check(ramp(x * 8 - 1, -1, 4) < broadcast(y * 8, 4), broadcast(x < y + 1, 4)); + // A horizontal AND/OR of a single ramp/broadcast comparison collapses to + // a plain scalar comparison on the ramp's endpoints, for both orderings + // of ramp vs broadcast and both '<' and '<='. + check(VectorReduce::make(VectorReduce::And, ramp(x, y, 4) < broadcast(z, 4), 1), + max(y, 0) * 3 + x < z); + check(VectorReduce::make(VectorReduce::And, ramp(x, y, 4) <= broadcast(z, 4), 1), + max(y, 0) * 3 + x <= z); + check(VectorReduce::make(VectorReduce::And, broadcast(x, 4) < ramp(y, z, 4), 1), + x < min(z, 0) * 3 + y); + check(VectorReduce::make(VectorReduce::And, broadcast(x, 4) <= ramp(y, z, 4), 1), + x <= min(z, 0) * 3 + y); + + check(VectorReduce::make(VectorReduce::Or, ramp(x, y, 4) < broadcast(z, 4), 1), + min(y, 0) * 3 + x < z); + check(VectorReduce::make(VectorReduce::Or, ramp(x, y, 4) <= broadcast(z, 4), 1), + min(y, 0) * 3 + x <= z); + check(VectorReduce::make(VectorReduce::Or, broadcast(x, 4) < ramp(y, z, 4), 1), + x < max(z, 0) * 3 + y); + check(VectorReduce::make(VectorReduce::Or, broadcast(x, 4) <= ramp(y, z, 4), 1), + x <= max(z, 0) * 3 + y); + + // The "all lanes of a ramp lie within [lo, hi]" shape loop partitioning + // builds -- a lower-bound comparison ANDed with an upper-bound + // comparison, both against the same stride -- fuses to a plain And of + // two scalar comparisons, regardless of clause order. + { + Expr u = Var("u"); + check(VectorReduce::make(VectorReduce::And, + (broadcast(x, 4) <= ramp(y, z, 4)) && (ramp(w, z, 4) <= broadcast(u, 4)), + 1), + (x <= min(z, 0) * 3 + y) && (max(z, 0) * 3 + w <= u)); + check(VectorReduce::make(VectorReduce::And, + (ramp(w, z, 4) <= broadcast(u, 4)) && (broadcast(x, 4) <= ramp(y, z, 4)), + 1), + (max(z, 0) * 3 + w <= u) && (x <= min(z, 0) * 3 + y)); + } + // Check anded conditions apply to the then case only check(IfThenElse::make(x == 4 && y == 5, not_no_op(z + x + y), diff --git a/test/correctness/split_aligned.cpp b/test/correctness/split_aligned.cpp new file mode 100644 index 000000000000..5801185db446 --- /dev/null +++ b/test/correctness/split_aligned.cpp @@ -0,0 +1,95 @@ +#include "Halide.h" +#include + +using namespace Halide; +using namespace Halide::Internal; + +class MuxCounter : public IRVisitor { + using IRVisitor::visit; + + void visit(const Call *op) override { + IRVisitor::visit(op); + if (op->is_intrinsic(Call::IntrinsicOp::mux)) { + mux_count++; + } + } + + void visit(const For *op) override { + IRVisitor::visit(op); + for_count++; + } + +public: + int for_count{0}; + int mux_count{0}; +}; + +int main(int argc, char **argv) { + Var x{"x"}, xo{"xo"}, xi{"xi"}; + for (auto ts : {TailStrategy::ShiftInwards, TailStrategy::GuardWithIf}) { + Func f; + Param offset{"offset"}; + offset.set_range(0, 3); + f(x) = mux((x - offset) % 4, {x, x * x, 2 * x, -x * (x + 1)}); + f.output_buffer().dim(0).set_min(0); + f + .split(x, xo, xi, 4, offset, ts) + .unroll(xi); + + Module module = f.compile_to_module({offset}); + MuxCounter checker; + for (const LoweredFunc &f : module.functions()) { + f.body.accept(&checker); + } + + for (int i = 0; i < 4; i++) { + printf("Testing runtime alignment: %d\n", i); + offset.set(i); + Buffer im = f.realize({32}); + f.realize(im, get_target_from_environment()); + + for (int x = 0; x < 32; x++) { + int selector = (4 + x - offset.get()) % 4; + int expected; + if (selector == 0) { + expected = x; + } else if (selector == 1) { + expected = x * x; + } else if (selector == 2) { + expected = 2 * x; + } else { + expected = -x * (x + 1); + } + if (im(x) != expected) { + printf("im(%d) = %d instead of %d (selector: %d)\n", x, im(x), expected, selector); + return 1; + } + } + } + + if (ts == Halide::TailStrategy::ShiftInwards) { + if (checker.mux_count != 8) { + std::printf("Expected 8 muxes: %d\n", checker.mux_count); + return 1; + } + if (checker.for_count != 1) { + // The head and tail are reduced to a single iteration, so the loop is stripped. + std::printf("Expected one for loop: %d\n", checker.for_count); + return 1; + } + } else if (ts == Halide::TailStrategy::GuardWithIf) { + if (checker.mux_count != 0) { + std::printf("Expected 0 muxes: %d\n", checker.mux_count); + return 1; + } + if (checker.for_count != 2) { + // The head and tail are reduced to a single iteration, so the loop is stripped. + std::printf("Expected one for loop: %d\n", checker.for_count); + return 1; + } + } + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/split_aligned_2d.cpp b/test/correctness/split_aligned_2d.cpp new file mode 100644 index 000000000000..99ac9743d8c9 --- /dev/null +++ b/test/correctness/split_aligned_2d.cpp @@ -0,0 +1,90 @@ +#include "Halide.h" +#include + +using namespace Halide; +using namespace Halide::Internal; + +class MuxCounter : public IRVisitor { + using IRVisitor::visit; + + void visit(const Call *op) override { + IRVisitor::visit(op); + if (op->is_intrinsic(Call::IntrinsicOp::mux)) { + mux_count++; + } + } + + void visit(const For *op) override { + IRVisitor::visit(op); + for_count++; + } + +public: + int for_count{0}; + int mux_count{0}; +}; + +int main(int argc, char **argv) { + Var x{"x"}, xo{"xo"}, xi{"xi"}; + Var y{"y"}, yo{"yo"}, yi{"yi"}; + Func f; + Param offset_x{"offset_x"}, offset_y{"offset_y"}; + offset_x.set_range(0, 1); + offset_y.set_range(0, 1); + auto idx = [](const auto &x, const auto &y, const auto &offset_x, const auto &offset_y) { + return (2 * ((y - offset_y) % 2)) + ((x - offset_x) % 2); + }; + auto a = [](const auto &x, const auto &y) { return x * x; }; + auto b = [](const auto &x, const auto &y) { return x * y; }; + auto c = [](const auto &x, const auto &y) { return y * y; }; + auto d = [](const auto &x, const auto &y) { return x + y; }; + f(x, y) = mux(idx(x, y, offset_x, offset_y), {a(x, y), b(x, y), c(x, y), d(x, y)}); + f.output_buffer().dim(0).set_min(0); + f.output_buffer().dim(1).set_min(0); + + f + .split(x, xo, xi, 2, offset_x, Halide::TailStrategy::GuardWithIf) + .split(y, yo, yi, 2, offset_y, Halide::TailStrategy::GuardWithIf) + .never_partition_all() + .reorder(xi, yi, xo, yo) + .unroll(xi) + .unroll(yi) + .parallel(yo); + + Module module = f.compile_to_module({offset_x, offset_y}); + MuxCounter checker; + for (const LoweredFunc &f : module.functions()) { + f.body.accept(&checker); + } + + for (int i = 0; i < 4; i++) { + printf("Testing runtime alignment: x=%d y=%d\n", i / 2, i % 2); + offset_x.set(i / 2); + offset_y.set(i % 2); + Buffer im = f.realize({32, 32}); + f.realize(im, get_target_from_environment()); + + for (int y = 0; y < 32; y++) { + for (int x = 0; x < 32; x++) { + int selector = idx(2 + x, 2 + y, offset_x.get(), offset_y.get()); + int expected = std::vector>{a, b, c, d}[selector](x, y); + if (im(x, y) != expected) { + printf("im(%d, %d) = %d instead of %d (selector: %d)\n", x, y, im(x, y), expected, selector); + return 1; + } + } + } + } + + if (checker.mux_count != 0) { + std::printf("Expected 0 muxes: %d\n", checker.mux_count); + return 1; + } + if (checker.for_count != 1) { + std::printf("Expected 3 for loops: %d\n", checker.for_count); + return 1; + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/split_aligned_2d_3x3.cpp b/test/correctness/split_aligned_2d_3x3.cpp new file mode 100644 index 000000000000..37e060481141 --- /dev/null +++ b/test/correctness/split_aligned_2d_3x3.cpp @@ -0,0 +1,139 @@ +#include "Halide.h" +#include + +using namespace Halide; +using namespace Halide::Internal; + +class MuxCounter : public IRVisitor { + using IRVisitor::visit; + + void visit(const Call *op) override { + IRVisitor::visit(op); + if (op->is_intrinsic(Call::IntrinsicOp::mux)) { + mux_count++; + } + } + + void visit(const For *op) override { + IRVisitor::visit(op); + for_count++; + } + +public: + int for_count{0}; + int mux_count{0}; +}; + +template +T produce_mux_argument(int i, const T &x, const T &y) { + return x * (i % 3) * 3 + y * (i % 3); +} + +int main(int argc, char **argv) { + Var c{"c"}; + Var x{"x"}, xo{"xo"}, xi{"xi"}; + Var y{"y"}, yo{"yo"}, yi{"yi"}; + Func f("f"), R("R"), G("G"), B("B"); + Param offset_x{"offset_x"}, offset_y{"offset_y"}; + offset_x.set_range(0, 2); + offset_y.set_range(0, 2); + auto idx = [](const auto &x, const auto &y, const auto &offset_x, const auto &offset_y) { + return (3 * ((y - offset_y) % 3)) + ((x - offset_x) % 3); + }; + std::vector ways; + ways.reserve(9); + for (int i = 0; i < 9; ++i) { + ways.push_back(produce_mux_argument(i, x, y)); + } + R(x, y) = mux(idx(x, y, offset_x, offset_y), ways); + G(x, y) = mux(idx(x, y, offset_x, offset_y), ways); + B(x, y) = mux(idx(x, y, offset_x, offset_y), ways); + f(x, y, c) = mux(c, {R(x, y), G(x, y), B(x, y)}); + f.output_buffer().dim(0).set_min(0); + f.output_buffer().dim(1).set_min(0); + + // Split both dimensions so that the inner loops iterate over exactly one + // 3x3 tile of the repeating pattern, anchored at (offset_x, offset_y). + // Unrolling those inner loops should give each mux a constant index, so + // every mux folds away to the single way it selects. + f + .split(x, xo, xi, 3, offset_x, Halide::TailStrategy::GuardWithIf) + .split(y, yo, yi, 3, offset_y, Halide::TailStrategy::GuardWithIf) + .never_partition_all() + .reorder(c, xi, yi, xo, yo) + .unroll(xi) + .unroll(yi) + .bound(c, 0, 3) + .unroll(c); + + for (Func *channel : {&R, &G, &B}) { + channel->compute_at(f, xo) + .never_partition_all() + // We want to compute a *whole* tile of the channel: + .bound_storage(x, 3) + .bound_extent(x, 3) + .align_bounds(x, 3, offset_x) + .bound_storage(y, 3) + .bound_extent(y, 3) + .align_bounds(y, 3, offset_y) + // In principle, the .align_bounds() should be enough, but the simplifier in Halide + // fails to determine that the producer tile perfectly overlaps with the consumer + // tile. Both have extent 3, but the producer is not getting simplified. So here, + // we can help it and tell it to use an extent of 3, after which the simplifier + // manages to prove that 3 satisfies the minimum required. Proving the expression + // is <= 3 is easier than proving it's == 3. + // The simplifier gap manifests itself twice: once for the compute extent, + // and once for the the storage extent. So we specify both. Unfortunately, + // the scheduling order here is important: .bound_extent() must precede + // .align_bounds() for the bounds inference to do the correct thing. + // + // Finally, we unroll to get rid of the muxes: + .unroll(x) + .unroll(y); + } + + Module module = f.compile_to_module({offset_x, offset_y}); + MuxCounter checker; + for (const LoweredFunc &lf : module.functions()) { + lf.body.accept(&checker); + } + + const int W = 32, H = 32; + for (int oy = 0; oy < 3; oy++) { + for (int ox = 0; ox < 3; ox++) { + printf("Testing runtime alignment: x=%d y=%d\n", ox, oy); + offset_x.set(ox); + offset_y.set(oy); + Buffer im = f.realize({W, H, 3}); + + for (int cc = 0; cc < 3; cc++) { + for (int y = 0; y < H; y++) { + for (int x = 0; x < W; x++) { + // Bias by 3 so the operands of % stay non-negative, + // where C++'s truncated % agrees with Halide's + // Euclidean %. + int selector = idx(3 + x, 3 + y, ox, oy); + int expected = produce_mux_argument(selector, x, y); + if (im(x, y, cc) != expected) { + printf("im(%d, %d, %d) = %d instead of %d (selector: %d)\n", + x, y, cc, im(x, y, cc), expected, selector); + return 1; + } + } + } + } + } + } + + if (checker.mux_count != 0) { + printf("Expected 0 muxes, got: %d\n", checker.mux_count); + return 1; + } + if (checker.for_count != 2) { + printf("Expected 2 for loops, got: %d\n", checker.for_count); + return 1; + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/split_aligned_nested.cpp b/test/correctness/split_aligned_nested.cpp new file mode 100644 index 000000000000..59b411c614f6 --- /dev/null +++ b/test/correctness/split_aligned_nested.cpp @@ -0,0 +1,100 @@ +#include "Halide.h" +#include + +// Nests two aligned splits: x is split into (xo, xi) aligned to p1, and then +// the resulting outer var xo is itself split into (xoo, xoi) aligned to a +// second, independent runtime Param p2. This exercises the aligned-split +// machinery (ApplySplit.cpp's apply_split/compute_loop_bounds_after_split) +// on a var whose own loop_min is not a compile-time constant (it comes from +// the first split's outer bound, which is a function of p1), stacked with a +// second, unrelated alignment. The mux selector only depends on p1, so this +// is primarily a correctness test of composing aligned splits -- the +// reconstruction of x from xoo, xoi, and xi has to be correct for every +// combination of the two independently-varying runtime alignments. +// +// Both splits are tried with both GuardWithIf and ShiftInwards (as in +// split_aligned.cpp): correctness must hold for all four combinations, and +// as in split_aligned.cpp the mux only fully resolves at compile time (0 +// muxes) when the split that carries the selector's alignment (the first +// one, on x) uses GuardWithIf; ShiftInwards leaves 8 muxes unresolved +// because the clamped base is no longer a compile-time-constant offset from +// the unrolled lane on every iteration. The tail strategy of the second, +// unrelated split (on xo) doesn't affect that count either way. + +using namespace Halide; +using namespace Halide::Internal; + +class MuxCounter : public IRVisitor { + using IRVisitor::visit; + + void visit(const Call *op) override { + IRVisitor::visit(op); + if (op->is_intrinsic(Call::IntrinsicOp::mux)) { + mux_count++; + } + } + +public: + int mux_count{0}; +}; + +int main(int argc, char **argv) { + for (auto ts1 : {TailStrategy::GuardWithIf, TailStrategy::ShiftInwards}) { + for (auto ts2 : {TailStrategy::GuardWithIf, TailStrategy::ShiftInwards}) { + printf("Testing tail strategies: ts1=%d ts2=%d\n", (int)ts1, (int)ts2); + + Var x{"x"}, xo{"xo"}, xi{"xi"}, xoo{"xoo"}, xoi{"xoi"}; + Func f{"f"}; + Param p1{"p1"}, p2{"p2"}; + p1.set_range(0, 3); + p2.set_range(0, 2); + + f(x) = mux((x - p1) % 4, {x, x * x, 2 * x, -x * (x + 1)}); + f.output_buffer().dim(0).set_min(0); + + f.split(x, xo, xi, 4, p1, ts1) + .split(xo, xoo, xoi, 3, p2, ts2) + .unroll(xi); + + Module module = f.compile_to_module({p1, p2}); + MuxCounter checker; + for (const LoweredFunc &lf : module.functions()) { + lf.body.accept(&checker); + } + int expected_mux_count = (ts1 == TailStrategy::GuardWithIf) ? 0 : 8; + if (checker.mux_count != expected_mux_count) { + printf("Expected %d muxes: %d\n", expected_mux_count, checker.mux_count); + return 1; + } + + for (int a1 = 0; a1 < 4; a1++) { + for (int a2 = 0; a2 < 3; a2++) { + p1.set(a1); + p2.set(a2); + Buffer im = f.realize({61}); + for (int x = 0; x < 61; x++) { + int selector = (4 + x - a1) % 4; + int expected; + if (selector == 0) { + expected = x; + } else if (selector == 1) { + expected = x * x; + } else if (selector == 2) { + expected = 2 * x; + } else { + expected = -x * (x + 1); + } + if (im(x) != expected) { + printf("im(%d) = %d instead of %d (p1: %d, p2: %d, ts1: %d, ts2: %d)\n", + x, im(x), expected, a1, a2, (int)ts1, (int)ts2); + return 1; + } + } + } + } + } + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/split_aligned_partition.cpp b/test/correctness/split_aligned_partition.cpp new file mode 100644 index 000000000000..6e457b70dc2b --- /dev/null +++ b/test/correctness/split_aligned_partition.cpp @@ -0,0 +1,116 @@ +#include "Halide.h" +#include + +using namespace Halide; +using namespace Halide::Internal; + +namespace { + +// Note: this test is built with NDEBUG, so assert() compiles to nothing. +bool check(bool ok, const char *msg) { + if (!ok) { + printf("Failed: %s\n", msg); + } + return ok; +} + +class LoopExtents : public IRVisitor { + using IRVisitor::visit; + + void visit(const For *op) override { + extents.push_back(simplify(op->extent())); + IRVisitor::visit(op); + } + +public: + std::vector extents; +}; + +class CountMod : public IRVisitor { + using IRVisitor::visit; + + void visit(const Mod *op) override { + count++; + IRVisitor::visit(op); + } + +public: + int count{0}; +}; + +} // namespace + +int main(int argc, char **argv) { + // A loop of eight elements whose first and last iterations are special, + // and whose interior is periodic with period two. Unrolling the interior + // by two turns the % into a constant, but only if the unrolled pairs line + // up with the periodicity -- which means the tiles have to start at x=1, + // where the interior begins, not at x=0. + // + // An aligned split expresses exactly that: split by two, anchored at one. + // Loop partitioning then peels the one iteration at each end that the + // likely() marks as not-steady-state, leaving x in [1, 6]. That's six + // iterations, or three of the unrolled-by-two loop. + // + // Without the alignment the tiles start at x=0 instead, the interior + // doesn't fill a whole number of them, and partitioning has to peel two + // iterations at each end rather than one -- leaving a steady-state loop + // of two rather than three. + Var x{"x"}, xo{"xo"}, xi{"xi"}; + Func f{"f"}; + f(x) = select(x <= 0, 100, + x < 7, likely(x % 2), + 200); + f.bound(x, 0, 8); + f.split(x, xo, xi, 2, 1, TailStrategy::GuardWithIf) + .always_partition(xo) + .unroll(xi); + + Module m = f.compile_to_module({}, "f"); + + LoopExtents loops; + CountMod mods; + for (const LoweredFunc &lf : m.functions()) { + lf.body.accept(&loops); + lf.body.accept(&mods); + } + + printf("Loops:"); + for (const Expr &e : loops.extents) { + std::cout << " " << e; + } + printf("\n"); + + // The two peeled iterations are single elements, so they come out as + // straight-line code rather than loops. What's left is the steady state. + if (!check(loops.extents.size() == 1, "expected exactly one remaining loop")) { + return 1; + } + if (!check(is_const(loops.extents[0], 3), + "expected the steady-state loop to run three times (six " + "elements, unrolled by two)")) { + return 1; + } + + // The whole point of unrolling the interior was to fold away the %. + if (!check(mods.count == 0, "expected the modulo to fold away")) { + return 1; + } + + Buffer out = f.realize({8}); + for (int i = 0; i < 8; i++) { + int expected = 200; + if (i <= 0) { + expected = 100; + } else if (i < 7) { + expected = i % 2; + } + if (out(i) != expected) { + printf("out(%d) = %d instead of %d\n", i, out(i), expected); + return 1; + } + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/split_aligned_reduction.cpp b/test/correctness/split_aligned_reduction.cpp new file mode 100644 index 000000000000..65c1fc66a9e8 --- /dev/null +++ b/test/correctness/split_aligned_reduction.cpp @@ -0,0 +1,66 @@ +#include "Halide.h" +#include + +// A simple reduction (no rfactor) with a single aligned split, tried with +// GuardWithIf, RoundUpAndBlend, and ShiftInwardsAndBlend. +// +// The split here is of the pure var x, not of the RDom's r: Stage::split +// only allows GuardWithIf or Predicate when splitting an RVar itself (see +// Func.cpp), since RoundUp/ShiftInwards-family strategies would change the +// meaning of the reduction by recomputing or overrunning it. Splitting a +// pure var of an update definition doesn't have that restriction, and +// RoundUpAndBlend/ShiftInwardsAndBlend are exactly the tail strategies +// meant for vectorizing an update like this one (see their doc comments in +// Schedule.h). +// +// This is the same boundary-handling code in ApplySplit.cpp's +// ShiftInwardsAndBlend/RoundUpAndBlend branches exercised by +// rfactor_split_aligned_nested.cpp, but without rfactor's extra layer of +// indirection (splitting a var that's already itself the result of an +// aligned split) -- here x's own bounds are simple compile-time constants, +// so this isolates the aligned-split-plus-blend mechanics on their own. + +using namespace Halide; + +int main(int argc, char **argv) { + for (auto ts : {TailStrategy::GuardWithIf, TailStrategy::RoundUpAndBlend, TailStrategy::ShiftInwardsAndBlend}) { + printf("Testing tail strategy: %d\n", (int)ts); + + Var x{"x"}, xo{"xo"}, xi{"xi"}; + Func h{"h"}; + RDom r(0, 5, "r"); + Param p{"p"}; + p.set_range(0, 3); + + h(x) = 0; + h(x) += x + r; + h.compute_root(); + + h.update(0) + .split(x, xo, xi, 4, p, ts) + .vectorize(xi); + + // h is read through a further Func rather than realized directly, + // so that RoundUpAndBlend/ShiftInwardsAndBlend get an + // internally-allocated (and thus paddable) buffer to blend into, + // instead of a caller-provided one of a fixed, non-factor-multiple + // size. + Func out{"out"}; + out(x) = h(x); + + for (int a = 0; a < 4; a++) { + p.set(a); + Buffer im = out.realize({37}); + for (int x = 0; x < 37; x++) { + int expected = 5 * x + 10; + if (im(x) != expected) { + printf("im(%d) = %d instead of %d (p: %d, ts: %d)\n", x, im(x), expected, a, (int)ts); + return 1; + } + } + } + } + + printf("Success!\n"); + return 0; +}