diff --git a/Makefile b/Makefile index c66ac327781f..d3680426ae4e 100644 --- a/Makefile +++ b/Makefile @@ -524,6 +524,7 @@ SOURCE_FILES = \ ImageParam.cpp \ InferArguments.cpp \ InjectHostDevBufferCopies.cpp \ + InjectModuloVars.cpp \ Inline.cpp \ InlineReductions.cpp \ IntegerDivisionTable.cpp \ @@ -730,6 +731,7 @@ HEADER_FILES = \ ImageParam.h \ InferArguments.h \ InjectHostDevBufferCopies.h \ + InjectModuloVars.h \ Inline.h \ InlineReductions.h \ IntegerDivisionTable.h \ diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 02215600ef8e..26add5d97e41 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -136,6 +136,7 @@ target_sources( ImageParam.h InferArguments.h InjectHostDevBufferCopies.h + InjectModuloVars.h Inline.h InlineReductions.h IntegerDivisionTable.h @@ -319,6 +320,7 @@ target_sources( ImageParam.cpp InferArguments.cpp InjectHostDevBufferCopies.cpp + InjectModuloVars.cpp Inline.cpp InlineReductions.cpp IntegerDivisionTable.cpp 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/InjectModuloVars.cpp b/src/InjectModuloVars.cpp new file mode 100644 index 000000000000..22a1d73a6282 --- /dev/null +++ b/src/InjectModuloVars.cpp @@ -0,0 +1,387 @@ +#include "InjectModuloVars.h" + +#include "IREquality.h" +#include "IRMutator.h" +#include "IROperator.h" +#include "IRVisitor.h" +#include "Scope.h" +#include "Simplify.h" + +#include +#include +#include +#include +#include + +namespace Halide { +namespace Internal { + +namespace { + +// Reducing modulo m is only exact on integer types where overflow is +// undefined. On a type that wraps, a value and its reduction agree modulo m +// only if m divides the wrap-around, which isn't worth chasing. +bool suitable_type(Type t) { + return t.is_int() && t.bits() >= 32; +} + +// Note the check against zero: x % 0 is defined to be 0 rather than x, so it +// tells us nothing about x. +bool is_multiple_of(const Expr &e, int64_t m) { + auto c = as_const_int(e); + return c && *c != 0 && mod_imp(*c, m) == 0; +} + +// Call f on every Variable sitting in a position where only its value modulo +// m matters. Addition, subtraction and multiplication all commute with +// reduction, a ramp's base and stride each contribute linearly, and a value +// already reduced modulo a multiple of m is congruent to the original. +// Division is not on the list, and anything else is opaque. +template +void for_each_congruent_var(const Expr &e, int64_t m, Fn f) { + if (const Add *op = e.as()) { + for_each_congruent_var(op->a, m, f); + for_each_congruent_var(op->b, m, f); + } else if (const Sub *op = e.as()) { + for_each_congruent_var(op->a, m, f); + for_each_congruent_var(op->b, m, f); + } else if (const Mul *op = e.as()) { + for_each_congruent_var(op->a, m, f); + for_each_congruent_var(op->b, m, f); + } else if (const Ramp *op = e.as()) { + for_each_congruent_var(op->base, m, f); + for_each_congruent_var(op->stride, m, f); + } else if (const Broadcast *op = e.as()) { + for_each_congruent_var(op->value, m, f); + } else if (const Mod *op = e.as()) { + if (is_multiple_of(op->b, m)) { + for_each_congruent_var(op->a, m, f); + } + } else if (const Variable *op = e.as()) { + f(op); + } +} + +// Rebuild e with the variables in those same positions replaced. +Expr substitute_congruent_vars(const Expr &e, int64_t m, + const std::map &replacements) { + if (const Add *op = e.as()) { + return Add::make(substitute_congruent_vars(op->a, m, replacements), + substitute_congruent_vars(op->b, m, replacements)); + } else if (const Sub *op = e.as()) { + return Sub::make(substitute_congruent_vars(op->a, m, replacements), + substitute_congruent_vars(op->b, m, replacements)); + } else if (const Mul *op = e.as()) { + return Mul::make(substitute_congruent_vars(op->a, m, replacements), + substitute_congruent_vars(op->b, m, replacements)); + } else if (const Ramp *op = e.as()) { + return Ramp::make(substitute_congruent_vars(op->base, m, replacements), + substitute_congruent_vars(op->stride, m, replacements), + op->lanes); + } else if (const Broadcast *op = e.as()) { + return Broadcast::make(substitute_congruent_vars(op->value, m, replacements), + op->lanes); + } else if (const Mod *op = e.as()) { + if (is_multiple_of(op->b, m)) { + return Mod::make(substitute_congruent_vars(op->a, m, replacements), op->b); + } + } else if (const Variable *op = e.as()) { + auto it = replacements.find(op->name); + if (it != replacements.end() && it->second.type() == op->type) { + return it->second; + } + } + return e; +} + +template +const V *find(const std::map &m, const std::string &key) { + auto it = m.find(key); + return it == m.end() ? nullptr : &it->second; +} + +// Which let-bound variables are needed modulo which constants, and what the +// lets bind them to. +class FindModuloUses : public IRVisitor { + using IRVisitor::visit; + + void note(const Expr &e, int64_t m) { + for_each_congruent_var(e, m, [&](const Variable *v) { + if (let_values.count(v->name)) { + required[v->name].insert(m); + } + }); + } + + void visit(const Mod *op) override { + if (suitable_type(op->type.element_of())) { + if (auto m = as_const_int(op->b); m && *m > 1) { + note(op->a, *m); + } + } + IRVisitor::visit(op); + } + + void visit(const Let *op) override { + if (suitable_type(op->value.type().element_of())) { + let_values[op->name] = op->value; + } + IRVisitor::visit(op); + } + + void visit(const LetStmt *op) override { + if (suitable_type(op->value.type().element_of())) { + let_values[op->name] = op->value; + } + IRVisitor::visit(op); + } + +public: + std::map> required; + std::map let_values; + + // A value that must be known modulo m needs the things it is built from + // known modulo m too. Chase that to a fixed point. + void close_over_let_values() { + std::vector> worklist; + for (const auto &[name, moduli] : required) { + for (int64_t m : moduli) { + worklist.emplace_back(name, m); + } + } + while (!worklist.empty()) { + // Not a structured binding: the lambda below captures the modulus, + // and capturing a structured binding is C++20. + const std::pair item = worklist.back(); + worklist.pop_back(); + const int64_t m = item.second; + for_each_congruent_var(let_values[item.first], m, [&](const Variable *v) { + if (let_values.count(v->name) && required[v->name].insert(m).second) { + worklist.emplace_back(v->name, m); + } + }); + } + } +}; + +// The bounds a scalar Parameter is constrained to. add_parameter_checks +// asserts these on entry, so the simplifier may rely on them, but by this +// point in lowering the constrained copy of the parameter is long gone. +class FindParamBounds : public IRVisitor { + using IRVisitor::visit; + + void visit(const Variable *op) override { + if (op->param.defined() && !op->param.is_buffer() && + !bounds.contains(op->name)) { + Expr lo = op->param.min_value(), hi = op->param.max_value(); + if (lo.defined() || hi.defined()) { + bounds.push(op->name, Interval(lo.defined() ? lo : Interval::neg_inf(), + hi.defined() ? hi : Interval::pos_inf())); + } + } + } + +public: + Scope bounds; +}; + +class InjectModuloVars : public IRMutator { + using IRMutator::visit; + + const std::map> &required; + const std::map &let_values; + const Scope ¶m_bounds; + + // How deep we're willing to look through lets while building a reduced + // value, and how big we'll let the result get before giving up. Both only + // bound work done to decide whether a reduction collapses; nothing this + // large is ever added to the IR. + static constexpr int max_inline_depth = 4; + static constexpr int max_nodes = 1000; + int nodes_built = 0; + + // Conditions asserted by statements that dominate the one we're in. + // A require() on a parameter lands here, which is often what makes a + // reduction collapse: nothing else tells the simplifier that an offset + // is small enough to be its own remainder. + std::vector assumptions; + + // For a variable needed modulo m, the constant or variable it is congruent + // to, keyed by the name of the variable we bound that to. Reductions that + // didn't collapse aren't in here, and aren't bound at all. + std::map> congruent_value; + std::map> congruent_var; + + static std::string mod_var_name(const std::string &name, int64_t m) { + return name + ".mod." + std::to_string(m); + } + + // Build something congruent to e modulo m, in the hope that it collapses. + // A variable in a congruence-preserving position becomes whatever we + // already know it to be congruent to, or failing that its own let value, + // so that terms which cancel get the chance to meet. + Expr reduce(const Expr &e, int64_t m, int depth) { + if (nodes_built++ > max_nodes) { + return e; + } + if (const Add *op = e.as()) { + return Add::make(reduce(op->a, m, depth), reduce(op->b, m, depth)); + } else if (const Sub *op = e.as()) { + return Sub::make(reduce(op->a, m, depth), reduce(op->b, m, depth)); + } else if (const Mul *op = e.as()) { + return Mul::make(reduce(op->a, m, depth), reduce(op->b, m, depth)); + } else if (const Ramp *op = e.as()) { + return Ramp::make(reduce(op->base, m, depth), reduce(op->stride, m, depth), op->lanes); + } else if (const Broadcast *op = e.as()) { + return Broadcast::make(reduce(op->value, m, depth), op->lanes); + } else if (const Mod *op = e.as()) { + if (is_multiple_of(op->b, m)) { + return Mod::make(reduce(op->a, m, depth), op->b); + } + } else if (const Variable *op = e.as()) { + if (const Expr *known = find(congruent_value[m], op->name)) { + return *known; + } + if (depth < max_inline_depth) { + if (const Expr *value = find(let_values, op->name)) { + if (value->type() == op->type) { + return reduce(*value, m, depth + 1); + } + } + } + } + return e; + } + + // The reductions to bind just inside a let of this name and value, in + // order. Records what they are congruent to so later reductions and use + // sites can pick them up. + std::vector> reductions_for(const std::string &name, + const Expr &value) { + std::vector> result; + auto it = required.find(name); + if (it == required.end()) { + return result; + } + for (int64_t m : it->second) { + nodes_built = 0; + Expr reduced = reduce(value, m, 0); + reduced = simplify(Mod::make(reduced, make_const(value.type(), m)), + param_bounds, Scope::empty_scope(), assumptions); + // Only worth binding if it collapsed to something the simplifier + // will substitute back in at the use site. + if (!is_const(reduced) && !reduced.as()) { + continue; + } + std::string new_name = mod_var_name(name, m); + result.emplace_back(new_name, reduced); + congruent_value[m][name] = reduced; + congruent_var[m][name] = Variable::make(value.type(), new_name); + } + return result; + } + + void forget(const std::string &name, + const std::vector> &reductions) { + for (const auto &[new_name, value] : reductions) { + // Recover the modulus from the name we just made. + for (auto &[m, names] : congruent_value) { + if (mod_var_name(name, m) == new_name) { + names.erase(name); + congruent_var[m].erase(name); + } + } + } + } + + Stmt visit(const Block *op) override { + Stmt first = mutate(op->first); + bool pushed = false; + if (const AssertStmt *a = first.as()) { + if (is_pure(a->condition)) { + assumptions.push_back(a->condition); + pushed = true; + } + } + Stmt rest = mutate(op->rest); + if (pushed) { + assumptions.pop_back(); + } + if (first.same_as(op->first) && rest.same_as(op->rest)) { + return op; + } + return Block::make(first, rest); + } + + Expr visit(const Mod *op) override { + Expr a = mutate(op->a), b = mutate(op->b); + if (suitable_type(op->type.element_of())) { + if (auto m = as_const_int(b); m && *m > 1) { + auto it = congruent_var.find(*m); + if (it != congruent_var.end()) { + a = substitute_congruent_vars(a, *m, it->second); + } + } + } + if (a.same_as(op->a) && b.same_as(op->b)) { + return op; + } + return Mod::make(a, b); + } + + Expr visit(const Let *op) override { + Expr value = mutate(op->value); + auto reductions = reductions_for(op->name, value); + Expr body = mutate(op->body); + forget(op->name, reductions); + for (auto it = reductions.rbegin(); it != reductions.rend(); it++) { + body = Let::make(it->first, it->second, body); + } + if (value.same_as(op->value) && body.same_as(op->body)) { + return op; + } + return Let::make(op->name, value, body); + } + + Stmt visit(const LetStmt *op) override { + Expr value = mutate(op->value); + auto reductions = reductions_for(op->name, value); + Stmt body = mutate(op->body); + forget(op->name, reductions); + for (auto it = reductions.rbegin(); it != reductions.rend(); it++) { + body = LetStmt::make(it->first, it->second, body); + } + if (value.same_as(op->value) && body.same_as(op->body)) { + return op; + } + return LetStmt::make(op->name, value, body); + } + +public: + InjectModuloVars(const std::map> &required, + const std::map &let_values, + const Scope ¶m_bounds) + : required(required), let_values(let_values), param_bounds(param_bounds) { + } + + using IRMutator::mutate; +}; + +} // namespace + +Stmt inject_modulo_vars(const Stmt &s) { + FindModuloUses finder; + s.accept(&finder); + if (finder.required.empty()) { + return s; + } + finder.close_over_let_values(); + + FindParamBounds params; + s.accept(¶ms); + + return InjectModuloVars(finder.required, finder.let_values, params.bounds).mutate(s); +} + +} // namespace Internal +} // namespace Halide diff --git a/src/InjectModuloVars.h b/src/InjectModuloVars.h new file mode 100644 index 000000000000..346865d017c3 --- /dev/null +++ b/src/InjectModuloVars.h @@ -0,0 +1,34 @@ +#ifndef HALIDE_INJECT_MODULO_VARS_H +#define HALIDE_INJECT_MODULO_VARS_H + +/** \file + * Defines a pass that makes the value of a variable modulo a constant visible + * at the places where only that much of it is needed. + */ + +#include "Expr.h" + +namespace Halide { +namespace Internal { + +/** Find variables used only modulo a constant -- x % 5, (x + 3) % 5, and + * (x + 2*y) % 3 all need no more of x than x % 5 -- and bind that reduced + * value to a new variable next to the let that defines the original. Uses of + * the original in those positions are then replaced by the new variable. + * + * Reducing a value where it is defined often collapses it. A loop variable + * reconstructed by an aligned split looks like xo * 16 + offset, which modulo + * two is just offset. The simplifier can't discover that at the use site: a + * LetStmt hides the value behind a name, and it won't substitute a whole + * expression back in. It will substitute a variable, so giving the use site a + * name bound to the reduced value is enough. (x - offset) % 2 becomes + * (x.mod.2 - offset) % 2, and with x.mod.2 bound to offset that folds to zero. + * + * Only reductions that collapse to a constant or a single variable are kept, + * so this never grows the IR. */ +Stmt inject_modulo_vars(const Stmt &s); + +} // namespace Internal +} // namespace Halide + +#endif diff --git a/src/Lower.cpp b/src/Lower.cpp index 21acfff8df2d..82f8b5419323 100644 --- a/src/Lower.cpp +++ b/src/Lower.cpp @@ -41,6 +41,7 @@ #include "IRPrinter.h" #include "InferArguments.h" #include "InjectHostDevBufferCopies.h" +#include "InjectModuloVars.h" #include "Inline.h" #include "LICM.h" #include "LoopCarry.h" @@ -377,6 +378,16 @@ void lower_impl(const vector &output_funcs, s = simplify(s); log("Lowering after rewriting vector interleavings:", s); + debug(1) << "Injecting variables for values needed modulo a constant...\n"; + { + // Only worth re-simplifying if it bound anything. + Stmt t = inject_modulo_vars(s); + if (!t.same_as(s)) { + s = simplify(t); + log("Lowering after injecting modulo variables:", s); + } + } + debug(1) << "Partitioning loops to simplify boundary conditions...\n"; s = partition_loops(s); s = simplify(s); diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index 93048b3bed21..b73879a5d543 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -207,6 +207,7 @@ tests( in_place.cpp indexing_access_undef.cpp infer_arguments.cpp + inject_modulo_vars.cpp inline_reduction.cpp inlined_generator.cpp input_func_name_unique.cpp @@ -342,6 +343,7 @@ tests( split_aligned.cpp split_aligned_2d.cpp split_aligned_2d_3x3.cpp + split_aligned_mux_phase.cpp split_aligned_nested.cpp split_aligned_partition.cpp split_aligned_reduction.cpp diff --git a/test/correctness/inject_modulo_vars.cpp b/test/correctness/inject_modulo_vars.cpp new file mode 100644 index 000000000000..edb7cdc88d04 --- /dev/null +++ b/test/correctness/inject_modulo_vars.cpp @@ -0,0 +1,111 @@ +#include "Halide.h" +#include + +// Tests the lowering pass that binds a variable's value modulo a constant next +// to the let that defines it. See InjectModuloVars.h. + +using namespace Halide; +using namespace Halide::Internal; + +namespace { + +// Pull the argument back out of the sink() call the statements below wrap +// around the expression under test. +class FindSinkArg : public IRVisitor { + using IRVisitor::visit; + + void visit(const Call *op) override { + IRVisitor::visit(op); + if (op->name == "sink") { + arg = op->args[0]; + } + } + +public: + Expr arg; +}; + +Expr run(const Expr &e, const std::vector &asserts, const Expr &let_value) { + Expr x = Variable::make(Int(32), "x"); + // Mention x a second time so the simplifier won't just inline the let and + // make the whole exercise moot -- that is exactly the situation this pass + // exists for. + Stmt s = Evaluate::make(Call::make(Int(32), "sink", {e, x}, Call::Extern)); + s = LetStmt::make("x", let_value, s); + for (auto it = asserts.rbegin(); it != asserts.rend(); it++) { + s = Block::make(AssertStmt::make(*it, 0), s); + } + s = simplify(inject_modulo_vars(s)); + FindSinkArg finder; + s.accept(&finder); + return finder.arg; +} + +int check(const char *name, const Expr &got, const Expr &expected) { + if (!equal(got, expected)) { + std::cerr << name << ": got " << got << ", expected " << expected << "\n"; + return 1; + } + return 0; +} + +} // namespace + +int main(int argc, char **argv) { + Expr x = Variable::make(Int(32), "x"); + Expr xo = Variable::make(Int(32), "xo"); + Expr off = Variable::make(Int(32), "off"); + std::vector off_is_small{0 <= off && off <= 1}; + + // An aligned split's loop variable, used with the alignment subtracted + // back off. x is congruent to off modulo two, so this is zero. The extra + // term keeps the simplifier from peeling the "+ off" off the let by + // itself, which is what it does when the value ends in a bare variable. + Expr yo = Variable::make(Int(32), "yo"); + Expr aligned = xo * 16 + (yo * 2 + off); + int result = 0; + result |= check("cancel", run((x - off) % 2, off_is_small, aligned), 0); + // The same thing written the other way round, and with an offset. + result |= check("cancel, added", run((x + off) % 2, off_is_small, aligned), 0); + result |= check("cancel, +1", run((x - off + 1) % 2, off_is_small, aligned), 1); + + // The variable reached through a multiply is still only needed mod 2. + result |= check("through a multiply", run((3 * x - 3 * off) % 2, off_is_small, aligned), 0); + + // Without the assert there's nothing to say off is its own remainder, so + // this must be left alone rather than rewritten into something wrong. + { + Expr got = run((x - off) % 2, {}, aligned); + if (is_const(got)) { + std::cerr << "unbounded offset: should not have folded, got " << got << "\n"; + result = 1; + } + } + + // A congruence modulo 2 says nothing about the value modulo 3 here, since + // 16 is not a multiple of 3. + { + Expr got = run((x - off) % 3, off_is_small, aligned); + if (is_const(got)) { + std::cerr << "wrong modulus: should not have folded, got " << got << "\n"; + result = 1; + } + } + + // Division is not congruence-preserving, so nothing may be substituted + // through it. + { + Expr got = run((x / 2) % 2, off_is_small, aligned); + if (is_const(got)) { + std::cerr << "through a divide: should not have folded, got " << got << "\n"; + result = 1; + } + } + + if (result) { + return 1; + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/split_aligned_mux_phase.cpp b/test/correctness/split_aligned_mux_phase.cpp new file mode 100644 index 000000000000..5f3932325dc9 --- /dev/null +++ b/test/correctness/split_aligned_mux_phase.cpp @@ -0,0 +1,114 @@ +#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 reduce_expr_modulo_symbolic). +// 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.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}, "split_aligned_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; +}