diff --git a/.gitignore b/.gitignore index 4a013d25b796..b198de54cd0b 100644 --- a/.gitignore +++ b/.gitignore @@ -242,6 +242,7 @@ xcuserdata # NeoVim + clangd .cache +.ccls-cache # Emacs tags diff --git a/python_bindings/halide/src/halide_/PyScheduleMethods.h b/python_bindings/halide/src/halide_/PyScheduleMethods.h index f528af886dff..df0dbaa085fa 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("aligned_split", (T & (T::*)(const VarOrRVar &, const VarOrRVar &, const VarOrRVar &, const Expr &, const Expr &, TailStrategy)) & T::aligned_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 e8d01e9741c7..c2f51810e91f 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. @@ -61,10 +77,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); @@ -90,36 +115,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); @@ -166,12 +278,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/Deinterleave.cpp b/src/Deinterleave.cpp index 0fe237adb5ad..d1fcc68c9a6a 100644 --- a/src/Deinterleave.cpp +++ b/src/Deinterleave.cpp @@ -520,6 +520,17 @@ class Interleaver : public IRMutator { return visit_let(op); } + // The largest stride we'll deinterleave a ramp by. This is what decides + // whether a vectorized loop over a repeating pattern -- a sensor's colour + // filter array, say -- gets split into one slice per phase, so that each + // slice has a phase known at compile time. Bayer needs two and X-Trans + // needs six. + // + // It stops below eight because a permute with a period of eight is often + // better done as a single shuffle than as eight slices: on Hexagon, for + // instance, deinterleaving (x/8)*9 + x%8 displaces a single vdelta. + static constexpr int max_deinterleave_stride = 7; + Expr visit(const Ramp *op) override { if (op->stride.type().is_vector() && is_const_one(op->stride) && @@ -536,7 +547,7 @@ class Interleaver : public IRMutator { Expr visit(const Mod *op) override { const Ramp *r = op->a.as(); - for (int i = 2; i <= 4; ++i) { + for (int i = 2; i <= max_deinterleave_stride; ++i) { if (r && is_const(op->b, i) && (r->type.lanes() % i) == 0) { @@ -550,7 +561,7 @@ class Interleaver : public IRMutator { Expr visit(const Div *op) override { const Ramp *r = op->a.as(); - for (int i = 2; i <= 4; ++i) { + for (int i = 2; i <= max_deinterleave_stride; ++i) { if (r && is_const(op->b, i) && (r->type.lanes() % i) == 0 && diff --git a/src/Deserialization.cpp b/src/Deserialization.cpp index dd1bed996c04..2f9050ddab00 100644 --- a/src/Deserialization.cpp +++ b/src/Deserialization.cpp @@ -1158,6 +1158,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; @@ -1166,6 +1167,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 5a4a57ee6d9c..56cb3b387dac 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -1380,9 +1380,9 @@ FuncVec Stage::hoist_invariants() { return intms; } -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 " @@ -1392,6 +1392,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(); @@ -1595,11 +1603,11 @@ 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) { +Stage &Stage::aligned_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"; @@ -1608,7 +1616,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; } @@ -1690,7 +1704,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; } @@ -1941,7 +1955,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); } @@ -2822,6 +2836,12 @@ Func &Func::split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar return *this; } +Func &Func::aligned_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).aligned_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 8ee2c8116d42..9f515aa18b85 100644 --- a/src/Func.h +++ b/src/Func.h @@ -82,7 +82,7 @@ class Stage { void set_dim_type(const VarOrRVar &var, Internal::ForType t); 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); + const Expr &factor, const Expr &align, bool exact, TailStrategy tail); void remove(const std::string &var); const std::vector &storage_dims() const { @@ -410,6 +410,7 @@ class Stage { // @{ Stage &split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, TailStrategy tail = TailStrategy::Auto); + Stage &aligned_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); @@ -1566,6 +1567,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.aligned_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 &aligned_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 8d5a99e46836..b1916910b1ae 100644 --- a/src/Schedule.cpp +++ b/src/Schedule.cpp @@ -341,6 +341,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()) { @@ -724,6 +727,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 3eb985a6b04a..986d01b0b65a 100644 --- a/src/Schedule.h +++ b/src/Schedule.h @@ -329,6 +329,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 53931a6e78a4..7acd42aa9a9f 100644 --- a/src/Serialization.cpp +++ b/src/Serialization.cpp @@ -1257,10 +1257,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..71dfbd909e49 100644 --- a/src/Simplify_Add.cpp +++ b/src/Simplify_Add.cpp @@ -21,6 +21,34 @@ Expr Simplify::visit(const Add *op, ExprInfo *info) { std::swap(a_info, b_info); } + // Reassociate/commute Param variables to the right. + auto is_param = [](const Expr &e) { + const auto *v = e.as(); + return v && v->param.defined(); + }; + if (!is_param(b) && !is_const(b)) { + if (const Add *add_a = a.as()) { + if (is_param(add_a->a)) { + // (p + x) + y => (x + y) + p + return mutate((add_a->b + b) + add_a->a, info); + } + if (is_param(add_a->b)) { + // (x + p) + y => (x + y) + p + return mutate((add_a->a + b) + add_a->b, info); + } + } + if (const Add *add_b = b.as()) { + if (is_param(add_b->a)) { + // x + (p + y) => (x + y) + p + return mutate((a + add_b->b) + add_b->a, info); + } + if (is_param(add_b->b)) { + // x + (y + p) => (x + y) + p + return mutate((a + add_b->a) + add_b->b, info); + } + } + } + auto rewrite = IRMatcher::rewriter(IRMatcher::add(a, b), op->type); if (rewrite(IRMatcher::Overflow() + x, a) || @@ -40,7 +68,15 @@ Expr Simplify::visit(const Add *op, ExprInfo *info) { if (EVAL_IN_LAMBDA // (rewrite(c0 + c1, fold(c0 + c1)) || + // Simple linear combinations of one Expr: rewrite(x + x, x * 2) || + rewrite(x * y + x, x * (y + 1)) || + rewrite(x + x * y, x * (y + 1)) || + rewrite(x * y + z * y, (x + z) * y) || + rewrite(x * y + y * z, (x + z) * y) || + rewrite(y * x + z * y, y * (x + z)) || + rewrite(y * x + y * z, y * (x + z)) || + rewrite(ramp(x, y, c0) + ramp(z, w, c0), ramp(x + z, y + w, c0)) || rewrite(ramp(x, y, c0) + broadcast(z, c0), ramp(x + z, y, c0)) || rewrite(broadcast(x, c0) + broadcast(y, c1), broadcast(x + broadcast(y, fold(c1 / c0)), c0), c1 % c0 == 0) || @@ -80,6 +116,14 @@ 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. + 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) || @@ -111,19 +155,11 @@ Expr Simplify::visit(const Add *op, ExprInfo *info) { rewrite(((0 - x) - y) + z, z - (x + y)) || rewrite(((c0 - x) - y) + c1, (fold(c0 + c1) - y) - x) || - rewrite(x * y + z * y, (x + z) * y) || - rewrite(x * y + y * z, (x + z) * y) || - rewrite(y * x + z * y, y * (x + z)) || - rewrite(y * x + y * z, y * (x + z)) || - rewrite((x * y) + (z - (w * x)), z + (x * (y - w))) || rewrite((x * y) + (z - (w * y)), z + (y * (x - w))) || rewrite((x * y) + (z - (x * w)), z + (x * (y - w))) || rewrite((x * y) + (z - (y * w)), z + (y * (x - w))) || - rewrite(x * c0 + y * c1, (x + y * fold(c1 / c0)) * c0, c1 % c0 == 0) || - rewrite(x * c0 + y * c1, (x * fold(c0 / c1) + y) * c1, c0 % c1 == 0) || - // Hoist shuffles. The Shuffle visitor wants to sink // extract_elements to the leaves, and those count as degenerate // slices, so only hoist shuffles that grab more than one lane. @@ -202,6 +238,9 @@ Expr Simplify::visit(const Add *op, ExprInfo *info) { 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) || + // Very specific form that shows up with working with aligned_splits: + rewrite(((c2 - x) / c0) + ((x % c0 + c1) / c0), fold(c2 / c0 + c1 / c0) - (x / c0), c0 > 0 && (c1 + 1) % c0 == 0 && c2 % c0 == 0) || + false)))) { return mutate(rewrite.result, info); } diff --git a/src/Simplify_Div.cpp b/src/Simplify_Div.cpp index 4098f6f027e7..dce8ae35c361 100644 --- a/src/Simplify_Div.cpp +++ b/src/Simplify_Div.cpp @@ -131,6 +131,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 @@ -242,7 +244,9 @@ Expr Simplify::visit(const Div *op, ExprInfo *info) { x * fold(c0 / c2) + fold(c1 / c2), c2 > 0 && c0 % c2 == 0) || // A very specific pattern that comes up in bounds in upsampling code. - rewrite((x % 2 + c0) / 2, x % 2 + fold(c0 / 2), c0 % 2 == 1))))) { + rewrite((x % 2 + c0) / 2, x % 2 + fold(c0 / 2), c0 % 2 == 1) || + rewrite((0 - x % 2) / 2, (0 - x % 2)) || + false)))) { return mutate(rewrite.result, info); } diff --git a/src/Simplify_Exprs.cpp b/src/Simplify_Exprs.cpp index c19fa2e7fed8..5ec95ef0e0fb 100644 --- a/src/Simplify_Exprs.cpp +++ b/src/Simplify_Exprs.cpp @@ -214,8 +214,19 @@ 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)) || + + // Same rules as above, but with a confounding w sub-Expr in the h_and. + // This lets us peel clauses one by one. + 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 +248,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_Internal.h b/src/Simplify_Internal.h index 94a50bebc644..b59be5f91f2c 100644 --- a/src/Simplify_Internal.h +++ b/src/Simplify_Internal.h @@ -425,18 +425,21 @@ class Simplify : public VariadicVisitor { bool should_commute(const Expr &a, const Expr &b) { if (a.node_type() < b.node_type()) { return true; - } - if (a.node_type() > b.node_type()) { + } else if (a.node_type() > b.node_type()) { return false; - } - - if (a.node_type() == IRNodeType::Variable) { + } else if (a.node_type() == IRNodeType::Variable) { const Variable *va = a.as(); const Variable *vb = b.as(); - return va->name.compare(vb->name) > 0; + if (va->param.defined() && !vb->param.defined()) { + return true; + } else if (!va->param.defined() && vb->param.defined()) { + return false; + } else { + return va->name.compare(vb->name) > 0; + } + } else { + return false; } - - return false; } std::set truths, falsehoods; diff --git a/src/Simplify_Let.cpp b/src/Simplify_Let.cpp index 63a25147b03e..c7fd6980d3d0 100644 --- a/src/Simplify_Let.cpp +++ b/src/Simplify_Let.cpp @@ -119,6 +119,18 @@ HALIDE_NEVER_INLINE Body Simplify::simplify_let_inner(const LetOrLetStmt *op, Ex debug(4) << "simplify let " << op->name << " = " << f.value << " in...\n"; + // A variable with no let binding of its own (e.g. a loop variable) + // scaled by a constant is as cheap to push into the uses as a bare + // variable: substituting it cannot pull in yet another let value. + auto free_var_times_const = [&](const Expr &e) { + const Mul *m = e.as(); + if (!m || !is_const(m->b)) { + return false; + } + const Variable *v = m->a.as(); + return v && !var_info.contains(v->name); + }; + while (true) { const Variable *var = f.new_value.template as(); const Add *add = f.new_value.template as(); @@ -148,20 +160,25 @@ HALIDE_NEVER_INLINE Body Simplify::simplify_let_inner(const LetOrLetStmt *op, Ex var_b = shuffle->vectors[1].as(); } - if (is_const(f.new_value)) { + if (is_const(f.new_value) || var) { replacement = substitute(f.new_name, f.new_value, replacement); f.new_value = Expr(); break; - } else if (var) { - replacement = substitute(f.new_name, var, replacement); - f.new_value = Expr(); - break; } else if (add && (is_const(add->b) || var_b)) { replacement = substitute(f.new_name, Add::make(new_var, add->b), replacement); f.new_value = add->a; } else if (add && var_a) { replacement = substitute(f.new_name, Add::make(add->a, new_var), replacement); f.new_value = add->b; + } else if (add && free_var_times_const(add->b)) { + replacement = substitute(f.new_name, Add::make(new_var, add->b), replacement); + f.new_value = add->a; + } else if (add && free_var_times_const(add->a)) { + replacement = substitute(f.new_name, Add::make(add->a, new_var), replacement); + f.new_value = add->b; + } else if (sub && free_var_times_const(sub->b)) { + replacement = substitute(f.new_name, Sub::make(new_var, sub->b), replacement); + f.new_value = sub->a; } else if (mul && (is_const(mul->b) || var_b)) { replacement = substitute(f.new_name, Mul::make(new_var, mul->b), replacement); f.new_value = mul->a; diff --git a/src/Simplify_Mod.cpp b/src/Simplify_Mod.cpp index 7e5232da0975..04b082483119 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) % 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(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/Simplify_Mul.cpp b/src/Simplify_Mul.cpp index 1639e063d004..3fb91b5bcfad 100644 --- a/src/Simplify_Mul.cpp +++ b/src/Simplify_Mul.cpp @@ -64,6 +64,7 @@ Expr Simplify::visit(const Mul *op, ExprInfo *info) { rewrite((x * c0 + c1) * (x * c2 + c3), x * (x * fold(c0 * c2) + fold(c0 * c3 + c1 * c2)) + fold(c1 * c3)))) || rewrite((x + c0) * c1, x * c1 + fold(c0 * c1), !overflows(c0 * c1)) || rewrite((c0 - x) * c1, x * fold(-c1) + fold(c0 * c1), !overflows(c0 * c1)) || + rewrite(x * -1, 0 - x) || rewrite((0 - x) * y, 0 - x * y) || rewrite(x * (0 - y), 0 - x * y) || rewrite((x - y) * c0, (y - x) * fold(-c0), c0 < 0 && -c0 > 0) || diff --git a/src/Simplify_Sub.cpp b/src/Simplify_Sub.cpp index cb3fc09a0a13..858e83cf0417 100644 --- a/src/Simplify_Sub.cpp +++ b/src/Simplify_Sub.cpp @@ -244,6 +244,22 @@ Expr Simplify::visit(const Sub *op, ExprInfo *info) { rewrite(min(y + x, z) - x, min(z - x, y)) || rewrite(min(z, x + y) - x, min(z - x, y)) || rewrite(min(z, y + x) - x, min(z - x, y)) || + // The same, where the subtracted sum appears with its terms in + // the other order inside the min/max. + rewrite(min(x + (y + z), w) - (z + y), min(x, w - (z + y))) || + rewrite(min(x, (y + z) + w) - (z + y), min(x - (y + z), w)) || + rewrite(min(x, y + (z + w)) - (w + z), min(x - (z + w), y)) || + rewrite(min((x + y) + z, w) - (y + x), min(z, (w - y) - x)) || + rewrite(max(x + (y + z), w) - (z + y), max(x, w - (z + y))) || + rewrite(max(x, (y + z) + w) - (z + y), max(x - (y + z), w)) || + rewrite(max(x, y + (z + w)) - (w + z), max(x - (z + w), y)) || + rewrite(max((x + y) + z, w) - (y + x), max(z, (w - y) - x)) || + // The clamped term hides a multiple of x that cancels against + // the subtracted one, leaving a constant upper bound. + rewrite(min(min(x, c0) * c1 + c2, y) - x * c1, min(min(y, fold(c0 * c1 + c2)) - x * c1, c2), c1 > 0) || + rewrite(min(y, min(x, c0) * c1 + c2) - x * c1, min(min(y, fold(c0 * c1 + c2)) - x * c1, c2), c1 > 0) || + rewrite(min(min(x, c0) * c1, y) - x * c1, min(min(y, fold(c0 * c1)) - x * c1, 0), c1 > 0) || + rewrite(min(y, min(x, c0) * c1) - x * c1, min(min(y, fold(c0 * c1)) - x * c1, 0), c1 > 0) || rewrite((min(x, (w + (y + z))) - z), min(x - z, w + y)) || rewrite((min(x, (w + (z + y))) - z), min(x - z, w + y)) || rewrite((min(x, ((y + z) + w)) - z), min(x - z, y + w)) || diff --git a/src/Solve.cpp b/src/Solve.cpp index 8b976fbcaabe..1c6e63459950 100644 --- a/src/Solve.cpp +++ b/src/Solve.cpp @@ -240,6 +240,16 @@ class SolveExpression : public IRMutator { } } + if (!expr.defined()) { + // Adding zero is a no-op. This comes up when terms cancel, or when + // a negation arrives here written as (0 - f(x)). + if (is_const_zero(b)) { + expr = a; + } else if (is_const_zero(a)) { + expr = b; + } + } + if (!expr.defined()) { if (a.same_as(op->a) && b.same_as(op->b)) { expr = op; 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 1a35d566d3a5..3743186b8a17 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -5,6 +5,13 @@ tests( SOURCES # keep-sorted start case=no align_bounds.cpp + aligned_split.cpp + aligned_split_2d.cpp + aligned_split_2d_3x3.cpp + aligned_split_mux_phase.cpp + aligned_split_nested.cpp + aligned_split_partition.cpp + aligned_split_reduction.cpp argmax.cpp arm_cpu_detect.cpp associativity.cpp @@ -491,6 +498,10 @@ tests( random.cpp reorder_rvars.cpp rfactor.cpp + rfactor_aligned_split.cpp + rfactor_aligned_split_2d.cpp + rfactor_aligned_split_nested.cpp + rfactor_aligned_split_phases.cpp ring_buffer.cpp stream_compaction.cpp thread_safety.cpp diff --git a/test/correctness/aligned_split.cpp b/test/correctness/aligned_split.cpp new file mode 100644 index 000000000000..aca19d78236b --- /dev/null +++ b/test/correctness/aligned_split.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 + .aligned_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/aligned_split_2d.cpp b/test/correctness/aligned_split_2d.cpp new file mode 100644 index 000000000000..00534e40ce4c --- /dev/null +++ b/test/correctness/aligned_split_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 + .aligned_split(x, xo, xi, 2, offset_x, Halide::TailStrategy::GuardWithIf) + .aligned_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/aligned_split_2d_3x3.cpp b/test/correctness/aligned_split_2d_3x3.cpp new file mode 100644 index 000000000000..3ad9456fed92 --- /dev/null +++ b/test/correctness/aligned_split_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 + .aligned_split(x, xo, xi, 3, offset_x, Halide::TailStrategy::GuardWithIf) + .aligned_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/aligned_split_mux_phase.cpp b/test/correctness/aligned_split_mux_phase.cpp new file mode 100644 index 000000000000..9db3028eefb8 --- /dev/null +++ b/test/correctness/aligned_split_mux_phase.cpp @@ -0,0 +1,115 @@ +#include "Halide.h" +#include + +// A minimal distillation of a demosaicer scheduled with aligned splits. A +// demosaicer selects which interpolation to apply from the position within the +// sensor's colour filter pattern -- a mux over (x - offset) % period, where the +// period is 2 for a Bayer sensor and 6 for an X-Trans one -- and is scheduled +// with a split aligned to that same offset precisely so that the phase is +// constant across every tile, letting the mux resolve at compile time. +// +// For that to happen the compiler has to see two things. The alignment the +// split adds into the reconstructed loop variable has to cancel against the +// subtraction at the use site, even though the split binds that variable to a +// LetStmt and the offset is a runtime value (see the aligned_split() rewrite +// rule in Simplify_Add.cpp). And the vectorized loop has to be deinterleaved +// by the period, so that each of the resulting slices has a phase of its own +// (see Deinterleave.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 test(int period) { + const int width = 1024; + // The tile has to be a whole number of periods for the phase to be the + // same in every tile, and the vector a whole number of periods for each + // deinterleaved slice to have a single phase. + const int tile = period * 64; + const int vec = period * 2; + + Var x{"x"}, xo{"xo"}, xi{"xi"}; + Param off{"off"}; + off.set_range(0, period - 1); + + // One phase per position in the pattern. + std::vector phases; + for (int i = 0; i < period; i++) { + phases.push_back(x * (i + 1) + i); + } + + Func prod{"prod"}; + prod(x) = mux((x - off) % period, phases); + + Func out{"out"}; + out(x) = prod(x); + + out.output_buffer().dim(0).set_min(0); + + // The consumer is tiled with a split aligned to the same offset, so every + // tile starts at a position congruent to off. + out.aligned_split(x, xo, xi, tile, off, TailStrategy::GuardWithIf) + .vectorize(xi) + .never_partition_all(); + + // The producer is computed per tile. align_bounds keeps its min congruent + // to off as well, so its phase really is constant, and then it is split + // again to vectorize. + prod.compute_at(out, xo) + .never_partition_all() + .align_bounds(x, period, off) + .vectorize(x, vec, TailStrategy::RoundUp); + + Pipeline p(out); + + // Every phase is known at compile time, so no mux should survive lowering. + Module m = p.compile_to_module({off}, "aligned_split_mux_phase"); + MuxCounter checker; + m.functions().front().body.accept(&checker); + if (checker.mux_count != 0) { + printf("Period %d: expected 0 muxes, got %d\n", period, checker.mux_count); + return 1; + } + + // And it still has to compute the right thing. + for (int o = 0; o < period; o++) { + off.set(o); + Buffer result = p.realize({width}); + for (int i = 0; i < width; i++) { + int phase = (((i - o) % period) + period) % period; + int correct = i * (phase + 1) + phase; + if (result(i) != correct) { + printf("Period %d, off = %d: result(%d) = %d instead of %d\n", + period, o, i, result(i), correct); + return 1; + } + } + } + + return 0; +} + +int main(int argc, char **argv) { + // Bayer, a three-phase pattern, and X-Trans. + for (int period : {2, 3, 6}) { + if (test(period) != 0) { + return 1; + } + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/aligned_split_nested.cpp b/test/correctness/aligned_split_nested.cpp new file mode 100644 index 000000000000..8a9b8e5d444e --- /dev/null +++ b/test/correctness/aligned_split_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 +// aligned_split.cpp): correctness must hold for all four combinations, and +// as in aligned_split.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.aligned_split(x, xo, xi, 4, p1, ts1) + .aligned_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/aligned_split_partition.cpp b/test/correctness/aligned_split_partition.cpp new file mode 100644 index 000000000000..f000a544aa15 --- /dev/null +++ b/test/correctness/aligned_split_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.aligned_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/aligned_split_reduction.cpp b/test/correctness/aligned_split_reduction.cpp new file mode 100644 index 000000000000..a9b5b314fa93 --- /dev/null +++ b/test/correctness/aligned_split_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_aligned_split_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) + .aligned_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; +} diff --git a/test/correctness/rfactor_aligned_split.cpp b/test/correctness/rfactor_aligned_split.cpp new file mode 100644 index 000000000000..a257779ed0b9 --- /dev/null +++ b/test/correctness/rfactor_aligned_split.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 aligned_split.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) + .aligned_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_aligned_split_2d.cpp b/test/correctness/rfactor_aligned_split_2d.cpp new file mode 100644 index 000000000000..d28a303dfa5b --- /dev/null +++ b/test/correctness/rfactor_aligned_split_2d.cpp @@ -0,0 +1,101 @@ +#include "Halide.h" +#include + +// A 2D companion to rfactor_aligned_split.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) + .aligned_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_aligned_split_nested.cpp b/test/correctness/rfactor_aligned_split_nested.cpp new file mode 100644 index 000000000000..77233c26227f --- /dev/null +++ b/test/correctness/rfactor_aligned_split_nested.cpp @@ -0,0 +1,114 @@ +#include "Halide.h" +#include + +// A companion to rfactor_aligned_split.cpp and aligned_split_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) + .aligned_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) + .aligned_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_aligned_split_phases.cpp b/test/correctness/rfactor_aligned_split_phases.cpp new file mode 100644 index 000000000000..955922675ec3 --- /dev/null +++ b/test/correctness/rfactor_aligned_split_phases.cpp @@ -0,0 +1,168 @@ +#include "Halide.h" +#include + +// A variant of rfactor_aligned_split.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) + .aligned_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) + .aligned_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..26b07d099d9d 100644 --- a/test/correctness/simplify.cpp +++ b/test/correctness/simplify.cpp @@ -224,6 +224,7 @@ void check_algebra() { check(x * y + z * x, (y + z) * x); check(y * x + x * z, (y + z) * x); check(y * x + z * x, (y + z) * x); + check((0 - x) * 3 + x, x * -2); check((x - y) * z + (y * z + w), x * z + w); check((x - y) * z + (w + y * z), x * z + w); check((x - y) * z + (y * z - w), x * z - w); @@ -313,9 +314,6 @@ void check_algebra() { check((x - y) * -2, (y - x) * 2); check((xf - yf) * -2.0f, (yf - xf) * 2.0f); - check(x * 3 + y * 9, (y * 3 + x) * 3); - check(x * 9 + y * 3, (x * 3 + y) * 3); - // Pull terms that are a multiple of the divisor out of a ternary expression check(((x * 4 + y) + z) / 2, (y + z) / 2 + x * 2); check(((x * 4 - y) + z) / 2, (z - y) / 2 + x * 2); @@ -446,6 +444,7 @@ void check_algebra() { check((x + 3) % 14 == -1, f); check((x + 3) % 14 == 20, f); check((x + 3) % -14 == -1, f); + check((0 - (x % 2)) / 2 * 2 + (x % 2), 0 - (x % 2)); // Check an optimization important for fusing dimensions check((x / 3) * 3 + x % 3, x); @@ -546,6 +545,10 @@ void check_algebra() { check(5 % x > 0, 5 % x != 0); // Test case with most negative 32-bit number, as constant to check that it is not negated. + // The two terms sharing that same coefficient do get collected into one + // multiply (sound: multiplication distributes over addition mod 2^32 + // regardless of intermediate overflow), but the coefficient itself must + // survive untouched, in particular never negated -- that's not representable. check(((x * (int32_t)0x80000000) + (z * (int32_t)0x80000000 + y)), ((x * (int32_t)0x80000000) + (z * (int32_t)0x80000000 + y))); @@ -1241,13 +1244,13 @@ void check_bounds() { check(max(min(x, 5), 1) == 5, 5 <= x); check(min((x * 32 + y) * 4, x * 128 + 127), min(y * 4, 127) + x * 128); - check(min((x * 32 + y) * 4, x * 128 + 4), (min(y, 1) + x * 32) * 4); + check(min((x * 32 + y) * 4, x * 128 + 4), min(y, 1) * 4 + x * 128); check(min((y + x * 32) * 4, x * 128 + 127), min(y * 4, 127) + x * 128); - check(min((y + x * 32) * 4, x * 128 + 4), (min(y, 1) + x * 32) * 4); + check(min((y + x * 32) * 4, x * 128 + 4), min(y, 1) * 4 + x * 128); check(max((x * 32 + y) * 4, x * 128 + 127), max(y * 4, 127) + x * 128); - check(max((x * 32 + y) * 4, x * 128 + 4), (max(y, 1) + x * 32) * 4); + check(max((x * 32 + y) * 4, x * 128 + 4), max(y, 1) * 4 + x * 128); check(max((y + x * 32) * 4, x * 128 + 127), max(y * 4, 127) + x * 128); - check(max((y + x * 32) * 4, x * 128 + 4), (max(y, 1) + x * 32) * 4); + check(max((y + x * 32) * 4, x * 128 + 4), max(y, 1) * 4 + x * 128); check((min(x + y, z) + w) - x, min(z - x, y) + w); check(min((x + y) + w, z) - x, min(z - x, w + y)); @@ -1767,6 +1770,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/solve.cpp b/test/correctness/solve.cpp index 606bfdcb1a64..3662a87961de 100644 --- a/test/correctness/solve.cpp +++ b/test/correctness/solve.cpp @@ -277,7 +277,7 @@ void test_simple_division_cases() { // Note: x + (x*2 - y)/2 == (x*4 - y)/2 for all integers (floor division); // the split form x*2 - y/2 is incorrect when y is odd. check_solve(x + (x * 2 - y) / 2, (x * 4 - y) / 2); - check_solve(x + (-(x * 2) / 2), x * 0 + 0); + check_solve(x + (-(x * 2) / 2), x * 0); check_solve(x + (-(x * 2 + -3)) / 2, x * 0 + 1); check_solve(x + (z - (x * 2 + -3)) / 2, x * 0 + (z - (-3)) / 2); check_solve(x + (y * 16 + (z - (x * 2 + -1))) / 2,