diff --git a/python_bindings/src/halide/halide_/PyScheduleMethods.h b/python_bindings/src/halide/halide_/PyScheduleMethods.h index f528af886dff..7e585690e1e3 100644 --- a/python_bindings/src/halide/halide_/PyScheduleMethods.h +++ b/python_bindings/src/halide/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..6df71d634b9e 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. @@ -58,14 +74,16 @@ vector 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 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 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> 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 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..2de86684a8e7 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,35 @@ 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 + * phase, 'align', which need not be known at compile time. Instead + * of the inner dimension always iterating over [0, factor-1], it + * iterates over [align, align+factor-1]. 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/Lower.cpp b/src/Lower.cpp index e179376e05c9..58201af7126c 100644 --- a/src/Lower.cpp +++ b/src/Lower.cpp @@ -438,13 +438,13 @@ void lower_impl(const vector &output_funcs, if (t.has_feature(Target::Profile) || t.has_feature(Target::ProfileByTimer)) { debug(1) << "Injecting profiling...\n"; s = inject_profiling(s, pipeline_name, env, t); - s = simplify(s); log("Lowering after injecting profiling:", s); } debug(1) << "Finding intrinsics...\n"; // Must be run after the last simplification, because it turns // divisions into shifts, which the simplifier reverses. + s = simplify(s); s = find_intrinsics(s); log("Lowering after finding intrinsics:", s); @@ -458,9 +458,6 @@ void lower_impl(const vector &output_funcs, log("Lowering after stripping asserts:", s); } - debug(1) << "Lowering after final simplification:\n" - << s << "\n\n"; - if (!custom_passes.empty()) { for (size_t i = 0; i < custom_passes.size(); i++) { debug(1) << "Running custom lowering pass " << i << "...\n"; @@ -472,6 +469,8 @@ void lower_impl(const vector &output_funcs, // Make a copy of the Stmt code, before we lower anything to less human-readable code. result_module.set_conceptual_code_stmt(s); + debug(1) << "Lowering after reaching conceptual Stmt:\n" + << s << "\n\n"; if (t.arch != Target::Hexagon && t.has_feature(Target::HVX)) { debug(1) << "Splitting off Hexagon offload...\n"; 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/Simplify_Add.cpp b/src/Simplify_Add.cpp index a07ad1b4464b..a2298c2e019c 100644 --- a/src/Simplify_Add.cpp +++ b/src/Simplify_Add.cpp @@ -201,6 +201,7 @@ Expr Simplify::visit(const Add *op, ExprInfo *info) { 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), (c1 / c0) - (x / c0), c0 > 0 && (c1 + 1) % c0 == 0) || false)))) { return mutate(rewrite.result, info); diff --git a/src/Simplify_Exprs.cpp b/src/Simplify_Exprs.cpp index c19fa2e7fed8..b3ce824f0cf1 100644 --- a/src/Simplify_Exprs.cpp +++ b/src/Simplify_Exprs.cpp @@ -214,8 +214,23 @@ 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 (a lower-bound comparison ANDed with an + // upper-bound comparison, both against the same stride, e.g. + // (0 <= ramp(b0, s, n)) && (ramp(b1, s, n) <= extent)) + rewrite(h_and((broadcast(x, arg_lanes) <= ramp(y, z, arg_lanes)) && + (ramp(w, z, arg_lanes) <= broadcast(u, arg_lanes)), + 1), + (x <= y + min(z * (arg_lanes - 1), 0)) && + (w + max(z * (arg_lanes - 1), 0) <= u)) || + rewrite(h_and((ramp(w, z, arg_lanes) <= broadcast(u, arg_lanes)) && + (broadcast(x, arg_lanes) <= ramp(y, z, arg_lanes)), + 1), + (w + max(z * (arg_lanes - 1), 0) <= u) && + (x <= y + min(z * (arg_lanes - 1), 0))) || false) { return mutate(rewrite.result, info); } @@ -237,7 +252,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..0bbbddb4ec34 100644 --- a/src/Simplify_Mod.cpp +++ b/src/Simplify_Mod.cpp @@ -59,6 +59,10 @@ 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) % c0, (y - z) % c0) || + rewrite((((x * c0) + y) + z) % c0, (y + z) % c0) || + rewrite((((x * c0) - y) - z) % c0, (-y - z) % c0) || + rewrite((((x * c0) - y) + z) % c0, (z - y) % c0) || 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 c82d5d5ca513..62105207ec9d 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -336,6 +336,10 @@ tests( specialize_to_gpu.cpp specialize_trim_condition.cpp spirv_ir.cpp + split_aligned.cpp + split_aligned_2d.cpp + split_aligned_nested.cpp + split_aligned_reduction.cpp split_by_non_factor.cpp split_factor_type.cpp split_fuse_rvar.cpp @@ -476,6 +480,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/async_copy_chain.cpp b/test/correctness/async_copy_chain.cpp index 238b1eb8821b..efab90f153cd 100644 --- a/test/correctness/async_copy_chain.cpp +++ b/test/correctness/async_copy_chain.cpp @@ -5,7 +5,8 @@ using namespace Halide; Var x, y; void check(Func f) { - Buffer out = f.realize({256, 256}); + Target target = get_jit_target_from_environment().with_feature(Target::EnableBacktraces); + Buffer out = f.realize({256, 256}, target); out.for_each_element([&](int x, int y) { if (out(x, y) != x + y) { printf("out(%d, %d) = %d instead of %d\n", x, y, out(x, y), x + y); 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 a03b37d531fb..924e8fadc500 100644 --- a/test/correctness/simplify.cpp +++ b/test/correctness/simplify.cpp @@ -1751,6 +1751,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_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_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; +}