<commit_before>#include "HexagonOptimize.h"
#include "ConciseCasts.h"
#include "IRMutator.h"
#include "IROperator.h"
#include "IRMatch.h"
#include "IREquality.h"
#include "ExprUsesVar.h"
#include "CSE.h"
#include "Simplify.h"
#include "Substitute.h"
#include "Scope.h"
#include "Bounds.h"
#include "Lerp.h"

namespace Halide {
namespace Internal {

using std::set;
using std::vector;
using std::string;
using std::pair;

using namespace Halide::ConciseCasts;

Expr native_interleave(Expr x) {
    string fn;
    switch (x.type().bits()) {
    case 8: fn = "halide.hexagon.interleave.vb"; break;
    case 16: fn = "halide.hexagon.interleave.vh"; break;
    case 32: fn = "halide.hexagon.interleave.vw"; break;
    default: internal_error << "Cannot interleave native vectors of type " << x.type() << "\n";
    }
    return Call::make(x.type(), fn, {x}, Call::PureExtern);
}

Expr native_deinterleave(Expr x) {
    string fn;
    switch (x.type().bits()) {
    case 8: fn = "halide.hexagon.deinterleave.vb"; break;
    case 16: fn = "halide.hexagon.deinterleave.vh"; break;
    case 32: fn = "halide.hexagon.deinterleave.vw"; break;
    default: internal_error << "Cannot deinterleave native vectors of type " << x.type() << "\n";
    }
    return Call::make(x.type(), fn, {x}, Call::PureExtern);
}

bool is_native_interleave_op(Expr x, const char *name) {
    const Call *c = x.as<Call>();
    if (!c || c->args.size() != 1) return false;
    return starts_with(c->name, name);
}

bool is_native_interleave(Expr x) {
    return is_native_interleave_op(x, "halide.hexagon.interleave");
}

bool is_native_deinterleave(Expr x) {
    return is_native_interleave_op(x, "halide.hexagon.deinterleave");
}

namespace {

// Broadcast to an unknown number of lanes, for making patterns.
Expr bc(Expr x) { return Broadcast::make(x, 0); }

// Flatten a binary operator tree into a list of operands.
template <typename BinOp>
void flatten_binary_op(const BinOp *op, vector<Expr> &operands) {
    const BinOp *a = op->a.template as<BinOp>();
    const BinOp *b = op->b.template as<BinOp>();
    if (a) {
        flatten_binary_op(a, operands);
    } else {
        operands.push_back(op->a);
    }
    if (b) {
        flatten_binary_op(b, operands);
    } else {
        operands.push_back(op->b);
    }
}

// Reassemble a list of binary operator operands into a (balanced)
// binary tree.
template <typename BinOp>
Expr make_binary_op(const vector<Expr> &operands) {
    if (operands.empty()) {
        return Expr();
    }
    if (operands.size() == 1) {
        return operands[0];
    }
    vector<Expr> ops;
    for (size_t i = 0; i < operands.size(); i += 2) {
        if (i + 1 < operands.size()) {
            ops.push_back(BinOp::make(operands[i], operands[i + 1]));
        } else {
            ops.push_back(operands[i]);
        }
    }
    return make_binary_op<BinOp>(ops);
}

// This mutator rewrites patterns with an unknown number of lanes to
// have the specified number of lanes.
class WithLanes : public IRMutator {
    using IRMutator::visit;

    int lanes;

    Type with_lanes(Type t) { return t.with_lanes(lanes); }

    void visit(const Cast *op) {
        if (op->type.lanes() != lanes) {
            expr = Cast::make(with_lanes(op->type), mutate(op->value));
        } else {
            IRMutator::visit(op);
        }
    }

    void visit(const Variable *op) {
        if (op->type.lanes() != lanes) {
            expr = Variable::make(with_lanes(op->type), op->name);
        } else {
            expr = op;
        }
    }

    void visit(const Broadcast *op) {
        if (op->type.lanes() != lanes) {
            expr = Broadcast::make(op->value, lanes);
        } else {
            IRMutator::visit(op);
        }
    }

public:
    WithLanes(int lanes) : lanes(lanes) {}
};

Expr with_lanes(Expr x, int lanes) {
    return WithLanes(lanes).mutate(x);
}

struct Pattern {
    enum Flags {
        InterleaveResult = 1 << 0,  // After evaluating the pattern, interleave native vectors of the result.
        SwapOps01 = 1 << 1,  // Swap operands 0 and 1 prior to substitution.
        SwapOps12 = 1 << 2,  // Swap operands 1 and 2 prior to substitution.
        ExactLog2Op1 = 1 << 3, // Replace operand 1 with its log base 2, if the log base 2 is exact.
        ExactLog2Op2 = 1 << 4, // Save as above, but for operand 2.

        BeginExactLog2Op = 1,   // BeginExactLog2Op and EndExactLog2Op ensure that we check only op1 and op2
        EndExactLog2Op = 3,     // for ExactLog2Op

        DeinterleaveOp0 = 1 << 5,  // Prior to evaluating the pattern, deinterleave native vectors of operand 0.
        DeinterleaveOp1 = 1 << 6,  // Same as above, but for operand 1.
        DeinterleaveOp2 = 1 << 7,
        DeinterleaveOps = DeinterleaveOp0 | DeinterleaveOp1 | DeinterleaveOp2,

        BeginDeinterleaveOp = 0, // BeginDeinterleaveOp and EndDeinterleaveOp ensure that we check only three
        EndDeinterleaveOp = 3,   // deinterleave Op0, 1 and 2.
        // Many patterns are instructions that widen only
        // operand 0, which need to both deinterleave operand 0, and then
        // re-interleave the result.
        ReinterleaveOp0 = InterleaveResult | DeinterleaveOp0,

        NarrowOp0 = 1 << 10,  // Replace operand 0 with its half-width equivalent.
        NarrowOp1 = 1 << 11,  // Same as above, but for operand 1.
        NarrowOp2 = 1 << 12,
        NarrowOp3 = 1 << 13,
        NarrowOps = NarrowOp0 | NarrowOp1 | NarrowOp2 | NarrowOp3,

        NarrowUnsignedOp0 = 1 << 15,  // Similar to the above, but narrow to an unsigned half width type.
        NarrowUnsignedOp1 = 1 << 16,
        NarrowUnsignedOp2 = 1 << 17,
        NarrowUnsignedOps = NarrowUnsignedOp0 | NarrowUnsignedOp1 | NarrowUnsignedOp2,

   };

    string intrin;        // Name of the intrinsic
    Expr pattern;         // The pattern to match against
    int flags;

    Pattern() {}
    Pattern(const string &intrin, Expr p, int flags = 0)
        : intrin(intrin), pattern(p), flags(flags) {}
};

Expr wild_u8 = Variable::make(UInt(8), "*");
Expr wild_u16 = Variable::make(UInt(16), "*");
Expr wild_u32 = Variable::make(UInt(32), "*");
Expr wild_u64 = Variable::make(UInt(64), "*");
Expr wild_i8 = Variable::make(Int(8), "*");
Expr wild_i16 = Variable::make(Int(16), "*");
Expr wild_i32 = Variable::make(Int(32), "*");
Expr wild_i64 = Variable::make(Int(64), "*");

Expr wild_u8x = Variable::make(Type(Type::UInt, 8, 0), "*");
Expr wild_u16x = Variable::make(Type(Type::UInt, 16, 0), "*");
Expr wild_u32x = Variable::make(Type(Type::UInt, 32, 0), "*");
Expr wild_u64x = Variable::make(Type(Type::UInt, 64, 0), "*");
Expr wild_i8x = Variable::make(Type(Type::Int, 8, 0), "*");
Expr wild_i16x = Variable::make(Type(Type::Int, 16, 0), "*");
Expr wild_i32x = Variable::make(Type(Type::Int, 32, 0), "*");
Expr wild_i64x = Variable::make(Type(Type::Int, 64, 0), "*");

// Attempt to apply one of the patterns to x. If a match is
// successful, the expression is replaced with a call using the
// matched operands. Prior to substitution, the matches are mutated
// with op_mutator.
Expr apply_patterns(Expr x, const vector<Pattern> &patterns, IRMutator *op_mutator) {
    vector<Expr> matches;
    for (const Pattern &p : patterns) {
        if (expr_match(p.pattern, x, matches)) {
            // The Pattern::Narrow*Op* flags are ordered such that
            // the operand corresponds to the bit (with operand 0
            // corresponding to the least significant bit), so we
            // can check for them all in a loop.
            bool is_match = true;
            for (size_t i = 0; i < matches.size() && is_match; i++) {
                Type t = matches[i].type();
                Type target_t = t.with_bits(t.bits()/2);
                if (p.flags & (Pattern::NarrowOp0 << i)) {
                    matches[i] = lossless_cast(target_t, matches[i]);
                } else if (p.flags & (Pattern::NarrowUnsignedOp0 << i)) {
                    matches[i] = lossless_cast(target_t.with_code(Type::UInt), matches[i]);
                }
                if (!matches[i].defined()) is_match = false;
            }
            if (!is_match) continue;

            for (size_t i = Pattern::BeginExactLog2Op; i < Pattern::EndExactLog2Op && is_match; i++) {
                // This flag is mainly to capture shifts. When the
                // operand of a div or mul is a power of 2, we can use
                // a shift instead.
                if (p.flags & (Pattern::ExactLog2Op1 << (i - Pattern::BeginExactLog2Op))) {
                    int pow;
                    if (is_const_power_of_two_integer(matches[i], &pow)) {
                        matches[i] = cast(matches[i].type().with_lanes(1), pow);
                    } else {
                        is_match = false;
                    }
                }
            }
            if (!is_match) continue;

            for (size_t i = Pattern::BeginDeinterleaveOp; i < Pattern::EndDeinterleaveOp; i++) {
                if (p.flags &
                    (Pattern::DeinterleaveOp0 << (i - Pattern::BeginDeinterleaveOp))) {
                    internal_assert(matches[i].type().is_vector());
                    matches[i] = native_deinterleave(matches[i]);
                }
            }
            if (p.flags & Pattern::SwapOps01) {
                internal_assert(matches.size() >= 2);
                std::swap(matches[0], matches[1]);
            }
            if (p.flags & Pattern::SwapOps12) {
                internal_assert(matches.size() >= 3);
                std::swap(matches[1], matches[2]);
            }
            // Mutate the operands with the given mutator.
            for (Expr &op : matches) {
                op = op_mutator->mutate(op);
            }
            x = Call::make(x.type(), p.intrin, matches, Call::PureExtern);
            if (p.flags & Pattern::InterleaveResult) {
                // The pattern wants us to interleave the result.
                x = native_interleave(x);
            }
            return x;
        }
    }
    return x;
}

// Replace x with a negated version of x, if it can be done without
// overflow.
Expr lossless_negate(Expr x) {
    const Mul *m = x.as<Mul>();
    if (m) {
        Expr a = lossless_negate(m->a);
        if (a.defined()) {
            return Mul::make(a, m->b);
        }
        Expr b = lossless_negate(m->b);
        if (b.defined()) {
            return Mul::make(m->a, b);
        }
    }
    if (is_negative_negatable_const(x) || is_positive_const(x)) {
        return simplify(-x);
    }
    return Expr();
}

template <typename T>
Expr apply_commutative_patterns(const T *op, const vector<Pattern> &patterns, IRMutator *mutator) {
    Expr ret = apply_patterns(op, patterns, mutator);
    if (!ret.same_as(op)) return ret;

    // Try commuting the op
    Expr commuted = T::make(op->b, op->a);
    ret = apply_patterns(commuted, patterns, mutator);
    if (!ret.same_as(commuted)) return ret;

    return op;
}

// Perform peephole optimizations on the IR, adding appropriate
// interleave and deinterleave calls.
class OptimizePatterns : public IRMutator {
private:
    using IRMutator::visit;

    void visit(const Mul *op) {
        static vector<Pattern> scalar_muls = {
            // Vector by scalar widening multiplies.
            { "halide.hexagon.mpy.vub.ub", wild_u16x*bc(wild_u16), Pattern::InterleaveResult | Pattern::NarrowOps },
            { "halide.hexagon.mpy.vub.b",  wild_i16x*bc(wild_i16), Pattern::InterleaveResult | Pattern::NarrowUnsignedOp0 | Pattern::NarrowOp1 },
            { "halide.hexagon.mpy.vuh.uh", wild_u32x*bc(wild_u32), Pattern::InterleaveResult | Pattern::NarrowOps },
            { "halide.hexagon.mpy.vh.h",   wild_i32x*bc(wild_i32), Pattern::InterleaveResult | Pattern::NarrowOps },

            // Multiplication by powers of 2.
            { "halide.hexagon.shl.vub.ub", wild_u8x*bc(wild_u8), Pattern::ExactLog2Op1 },
            { "halide.hexagon.shl.vuh.uh", wild_u16x*bc(wild_u16), Pattern::ExactLog2Op1 },
            { "halide.hexagon.shl.vuw.uw", wild_u32x*bc(wild_u32), Pattern::ExactLog2Op1 },
            { "halide.hexagon.shl.vb.b", wild_i8x*bc(wild_i8), Pattern::ExactLog2Op1 },
            { "halide.hexagon.shl.vh.h", wild_i16x*bc(wild_i16), Pattern::ExactLog2Op1 },
            { "halide.hexagon.shl.vw.w", wild_i32x*bc(wild_i32), Pattern::ExactLog2Op1 },

            // Non-widening scalar multiplication.
            { "halide.hexagon.mul.vh.b", wild_i16x*bc(wild_i16), Pattern::NarrowOp1 },
            { "halide.hexagon.mul.vw.h", wild_i32x*bc(wild_i32), Pattern::NarrowOp1 },
            // TODO: There's also mul.vw.b. We currently generate mul.vw.h
            // instead. I'm not sure mul.vw.b is faster, it might even be
            // slower due to the extra step in broadcasting the scalar up to
            // 32 bits.
        };

        static vector<Pattern> muls = {
            // Widening multiplication
            { "halide.hexagon.mpy.vub.vub", wild_u16x*wild_u16x, Pattern::InterleaveResult | Pattern::NarrowOps },
            { "halide.hexagon.mpy.vuh.vuh", wild_u32x*wild_u32x, Pattern::InterleaveResult | Pattern::NarrowOps },
            { "halide.hexagon.mpy.vb.vb",   wild_i16x*wild_i16x, Pattern::InterleaveResult | Pattern::NarrowOps },
            { "halide.hexagon.mpy.vh.vh",   wild_i32x*wild_i32x, Pattern::InterleaveResult | Pattern::NarrowOps },

            { "halide.hexagon.mpy.vub.vb",  wild_i16x*wild_i16x, Pattern::InterleaveResult | Pattern::NarrowUnsignedOp0 | Pattern::NarrowOp1 },
            { "halide.hexagon.mpy.vh.vuh",  wild_i32x*wild_i32x, Pattern::InterleaveResult | Pattern::NarrowOp0 | Pattern::NarrowUnsignedOp1 },
            // We need to check for the commuted versions of these patterns
            // before the more general patterns below catch these ops. The
            // other fix for this would be to break this into a third group of
            // multiply patterns, so the commuted versions of these would get
            // matched first.
            { "halide.hexagon.mpy.vub.vb",  wild_i16x*wild_i16x, Pattern::InterleaveResult | Pattern::NarrowOp0 | Pattern::NarrowUnsignedOp1 | Pattern::SwapOps01 },
            { "halide.hexagon.mpy.vh.vuh",  wild_i32x*wild_i32x, Pattern::InterleaveResult | Pattern::NarrowUnsignedOp0 | Pattern::NarrowOp1 | Pattern::SwapOps01 },

            // One operand widening multiplication.
            { "halide.hexagon.mul.vw.vh", wild_i32x*wild_i32x, Pattern::ReinterleaveOp0 | Pattern::NarrowOp1 },
            { "halide.hexagon.mul.vw.vuh", wild_i32x*wild_i32x, Pattern::ReinterleaveOp0 | Pattern::NarrowUnsignedOp1 },
            { "halide.hexagon.mul.vuw.vuh", wild_u32x*wild_u32x, Pattern::ReinterleaveOp0 | Pattern::NarrowUnsignedOp1 },
        };

        if (op->type.is_vector()) {
            expr = apply_commutative_patterns(op, scalar_muls, this);
            if (!expr.same_as(op)) return;

            expr = apply_commutative_patterns(op, muls, this);
            if (!expr.same_as(op)) return;
        }
        IRMutator::visit(op);
    }

    // Helpers to generate horizontally reducing multiply operations.
    static Expr halide_hexagon_add_2mpy(string suffix, Expr v0, Expr v1, Expr c0, Expr c1) {
        Type t = v0.type();
        Type result_type = Int(t.bits()*2).with_lanes(t.lanes());
        Expr call = Call::make(result_type, "halide.hexagon.add_2mpy" + suffix, {v0, v1, c0, c1}, Call::PureExtern);
        return native_interleave(call);
    }

    static Expr halide_hexagon_add_2mpy(string suffix, Expr v01, Expr c01) {
        Type t = v01.type();
        Type result_type = t.with_bits(t.bits()*2).with_lanes(t.lanes()/2);
        if (c01.type().is_int()) {
            result_type = result_type.with_code(Type::Int);
        }
        return Call::make(result_type, "halide.hexagon.add_2mpy" + suffix, {v01, c01}, Call::PureExtern);
    }

    static Expr halide_hexagon_add_4mpy(string suffix, Expr v01, Expr c01) {
        Type t = v01.type();
        Type result_type = t.with_bits(t.bits()*4).with_lanes(t.lanes()/4);
        if (c01.type().is_int()) {
            result_type = result_type.with_code(Type::Int);
        }
        return Call::make(result_type, "halide.hexagon.add_4mpy" + suffix, {v01, c01}, Call::PureExtern);
    }

    typedef pair<Expr, Expr> MulExpr;

    // In the operands vector, find elements of the operands list that
    // are widening multiplies from types a_ty and b_ty. If mpy_count
    // multiplies aren't found (but at least one is), it will attempt
    // to multiply remaining operands by 1 to match the target
    // widening multiply.
    static pair<vector<MulExpr>, Expr> find_mpy_ops(const vector<Expr> &operands, Type a_ty, Type b_ty, int mpy_count) {
        Type add_ty = operands.front().type();
        int mpy_bits = std::max(a_ty.bits(), b_ty.bits())*2;
        Expr a_wild = Variable::make(a_ty.with_bits(mpy_bits).with_code(add_ty.code()), "*");
        Expr b_wild = Variable::make(b_ty.with_bits(mpy_bits).with_code(add_ty.code()), "*");
        if (a_wild.type().lanes() == 1) {
            a_wild = Broadcast::make(a_wild, b_ty.lanes());
        }
        if (b_wild.type().lanes() == 1) {
            b_wild = Broadcast::make(b_wild, a_ty.lanes());
        }

        // Try to match both the multiply and its commutation.
        Expr pattern = a_wild*b_wild;
        Expr commuted_pattern = b_wild*a_wild;
        if (add_ty.bits() != pattern.type().bits()) {
            pattern = Cast::make(add_ty, pattern);
            commuted_pattern = Cast::make(add_ty, commuted_pattern);
        }

        // Go through the operands, putting the matching multiplies in
        // matched, and the rest in not_matched.
        vector<MulExpr> matched;
        vector<Expr> not_matched;
        for (Expr i : operands) {
            if ((int)matched.size() >= mpy_count) {
                not_matched.emplace_back(i);
                continue;
            }
            Expr a, b;
            vector<Expr> matches;
            if (expr_match(pattern, i, matches)) {
                a = matches[0];
                b = matches[1];
            } else if (expr_match(commuted_pattern, i, matches)) {
                a = matches[1];
                b = matches[0];
            }

            if (a.defined()) {
                a = lossless_cast(a_ty.with_lanes(a.type().lanes()), a);
            }
            if (b.defined()) {
                b = lossless_cast(b_ty.with_lanes(b.type().lanes()), b);
            }

            if (a.defined() && b.defined()) {
                matched.emplace_back(a, b);
            } else {
                not_matched.emplace_back(i);
            }
        }

        if (matched.empty()) {
            // Don't fake a bunch of mpy ops if there are no real ones.
            return {matched, make_binary_op<Add>(not_matched)};
        }

        // If we didn't find mpy_count matches, go through the
        // not_matched list and try to find elements that match one of
        // the operand types, and set the other operand to 1.
        vector<Expr> rest;
        for (Expr i : not_matched) {
            if ((int)matched.size() < mpy_count) {
                Expr as_a = lossless_cast(a_ty.with_lanes(i.type().lanes()), i);
                Expr as_b = lossless_cast(b_ty.with_lanes(i.type().lanes()), i);
                if (as_a.defined()) {
                    matched.emplace_back(as_a, make_one(b_ty));
                } else if (as_b.defined()) {
                    matched.emplace_back(make_one(a_ty), as_b);
                } else {
                    rest.push_back(i);
                }
            } else {
                rest.push_back(i);
            }
        }

        internal_assert((int)matched.size() <= mpy_count);

        return {matched, make_binary_op<Add>(rest)};
    }

    void visit(const Add *op) {
        vector<Expr> matches;

        // vmpa, vdmpy, and vrmpy instructions are hard to match with
        // patterns, do it manually here.
        vector<Expr> operands;
        flatten_binary_op(op, operands);

        // Try to find vrmpy opportunities first, which consume 4 operands.
        if (op->type.is_vector() && operands.size() >= 4) {
            int lanes = op->type.lanes();
            vector<MulExpr> mpys;
            Expr rest;
            string suffix;

            // Try to find a vector*scalar multiply first, which will
            // match a subset of the expressions that vector*vector
            // matches.
            if (op->type.is_uint()) {
                std::tie(mpys, rest) = find_mpy_ops(operands, UInt(8).with_lanes(lanes), UInt(8), 4);
                suffix = ".vub.ub";
            } else {
                std::tie(mpys, rest) = find_mpy_ops(operands, UInt(8).with_lanes(lanes), Int(8), 4);
                suffix = ".vub.b";
            }
            if (mpys.size() == 4) {
                // TODO: It's possible that permuting the order of the
                // multiply operands can simplify the shuffle away.
                Expr a0123 = Shuffle::make_interleave({mpys[0].first, mpys[1].first, mpys[2].first, mpys[3].first});
                Expr b0123 = Shuffle::make_interleave({mpys[0].second, mpys[1].second, mpys[2].second, mpys[3].second});
                a0123 = simplify(a0123);
                b0123 = simplify(b0123);
                b0123 = reinterpret(Type(b0123.type().code(), 32, 1), b0123);
                expr = halide_hexagon_add_4mpy(suffix, a0123, b0123);
                if (rest.defined()) {
                    expr = Add::make(expr, rest);
                }
                expr = mutate(expr);
                return;
            }

            // Now try to match vector*vector vrmpy expressions.
            if (op->type.is_uint()) {
                std::tie(mpys, rest) = find_mpy_ops(operands, UInt(8).with_lanes(lanes), UInt(8).with_lanes(lanes), 4);
                suffix = ".vub.vub";
            } else {
                std::tie(mpys, rest) = find_mpy_ops(operands, UInt(8).with_lanes(lanes), Int(8).with_lanes(lanes), 4);
                suffix = ".vub.vb";
            }
            // TODO: suffix = ".vb.vb"
            if (mpys.size() == 4) {
                // TODO: It's possible that permuting the order of the
                // multiply operands can simplify the shuffle away.
                Expr a0123 = Shuffle::make_interleave({mpys[0].first, mpys[1].first, mpys[2].first, mpys[3].first});
                Expr b0123 = Shuffle::make_interleave({mpys[0].second, mpys[1].second, mpys[2].second, mpys[3].second});
                a0123 = simplify(a0123);
                b0123 = simplify(b0123);
                b0123 = reinterpret(Type(b0123.type().code(), 32, 1), b0123);
                expr = halide_hexagon_add_4mpy(suffix, a0123, b0123);
                if (rest.defined()) {
                    expr = Add::make(expr, rest);
                }
                expr = mutate(expr);
                return;
            }
        }

        // Find opportunities vdmpy or vmpa.
        if (op->type.is_vector() && operands.size() >= 2) {
            int lanes = op->type.lanes();

            vector<MulExpr> mpys;
            Expr rest;
            string vmpa_suffix;
            string vdmpy_suffix;
<commit_msg>Don't reinterpret vector operands as scalars.<commit_after>#include "HexagonOptimize.h"
#include "ConciseCasts.h"
#include "IRMutator.h"
#include "IROperator.h"
#include "IRMatch.h"
#include "IREquality.h"
#include "ExprUsesVar.h"
#include "CSE.h"
#include "Simplify.h"
#include "Substitute.h"
#include "Scope.h"
#include "Bounds.h"
#include "Lerp.h"

namespace Halide {
namespace Internal {

using std::set;
using std::vector;
using std::string;
using std::pair;

using namespace Halide::ConciseCasts;

Expr native_interleave(Expr x) {
    string fn;
    switch (x.type().bits()) {
    case 8: fn = "halide.hexagon.interleave.vb"; break;
    case 16: fn = "halide.hexagon.interleave.vh"; break;
    case 32: fn = "halide.hexagon.interleave.vw"; break;
    default: internal_error << "Cannot interleave native vectors of type " << x.type() << "\n";
    }
    return Call::make(x.type(), fn, {x}, Call::PureExtern);
}

Expr native_deinterleave(Expr x) {
    string fn;
    switch (x.type().bits()) {
    case 8: fn = "halide.hexagon.deinterleave.vb"; break;
    case 16: fn = "halide.hexagon.deinterleave.vh"; break;
    case 32: fn = "halide.hexagon.deinterleave.vw"; break;
    default: internal_error << "Cannot deinterleave native vectors of type " << x.type() << "\n";
    }
    return Call::make(x.type(), fn, {x}, Call::PureExtern);
}

bool is_native_interleave_op(Expr x, const char *name) {
    const Call *c = x.as<Call>();
    if (!c || c->args.size() != 1) return false;
    return starts_with(c->name, name);
}

bool is_native_interleave(Expr x) {
    return is_native_interleave_op(x, "halide.hexagon.interleave");
}

bool is_native_deinterleave(Expr x) {
    return is_native_interleave_op(x, "halide.hexagon.deinterleave");
}

namespace {

// Broadcast to an unknown number of lanes, for making patterns.
Expr bc(Expr x) { return Broadcast::make(x, 0); }

// Flatten a binary operator tree into a list of operands.
template <typename BinOp>
void flatten_binary_op(const BinOp *op, vector<Expr> &operands) {
    const BinOp *a = op->a.template as<BinOp>();
    const BinOp *b = op->b.template as<BinOp>();
    if (a) {
        flatten_binary_op(a, operands);
    } else {
        operands.push_back(op->a);
    }
    if (b) {
        flatten_binary_op(b, operands);
    } else {
        operands.push_back(op->b);
    }
}

// Reassemble a list of binary operator operands into a (balanced)
// binary tree.
template <typename BinOp>
Expr make_binary_op(const vector<Expr> &operands) {
    if (operands.empty()) {
        return Expr();
    }
    if (operands.size() == 1) {
        return operands[0];
    }
    vector<Expr> ops;
    for (size_t i = 0; i < operands.size(); i += 2) {
        if (i + 1 < operands.size()) {
            ops.push_back(BinOp::make(operands[i], operands[i + 1]));
        } else {
            ops.push_back(operands[i]);
        }
    }
    return make_binary_op<BinOp>(ops);
}

// This mutator rewrites patterns with an unknown number of lanes to
// have the specified number of lanes.
class WithLanes : public IRMutator {
    using IRMutator::visit;

    int lanes;

    Type with_lanes(Type t) { return t.with_lanes(lanes); }

    void visit(const Cast *op) {
        if (op->type.lanes() != lanes) {
            expr = Cast::make(with_lanes(op->type), mutate(op->value));
        } else {
            IRMutator::visit(op);
        }
    }

    void visit(const Variable *op) {
        if (op->type.lanes() != lanes) {
            expr = Variable::make(with_lanes(op->type), op->name);
        } else {
            expr = op;
        }
    }

    void visit(const Broadcast *op) {
        if (op->type.lanes() != lanes) {
            expr = Broadcast::make(op->value, lanes);
        } else {
            IRMutator::visit(op);
        }
    }

public:
    WithLanes(int lanes) : lanes(lanes) {}
};

Expr with_lanes(Expr x, int lanes) {
    return WithLanes(lanes).mutate(x);
}

struct Pattern {
    enum Flags {
        InterleaveResult = 1 << 0,  // After evaluating the pattern, interleave native vectors of the result.
        SwapOps01 = 1 << 1,  // Swap operands 0 and 1 prior to substitution.
        SwapOps12 = 1 << 2,  // Swap operands 1 and 2 prior to substitution.
        ExactLog2Op1 = 1 << 3, // Replace operand 1 with its log base 2, if the log base 2 is exact.
        ExactLog2Op2 = 1 << 4, // Save as above, but for operand 2.

        BeginExactLog2Op = 1,   // BeginExactLog2Op and EndExactLog2Op ensure that we check only op1 and op2
        EndExactLog2Op = 3,     // for ExactLog2Op

        DeinterleaveOp0 = 1 << 5,  // Prior to evaluating the pattern, deinterleave native vectors of operand 0.
        DeinterleaveOp1 = 1 << 6,  // Same as above, but for operand 1.
        DeinterleaveOp2 = 1 << 7,
        DeinterleaveOps = DeinterleaveOp0 | DeinterleaveOp1 | DeinterleaveOp2,

        BeginDeinterleaveOp = 0, // BeginDeinterleaveOp and EndDeinterleaveOp ensure that we check only three
        EndDeinterleaveOp = 3,   // deinterleave Op0, 1 and 2.
        // Many patterns are instructions that widen only
        // operand 0, which need to both deinterleave operand 0, and then
        // re-interleave the result.
        ReinterleaveOp0 = InterleaveResult | DeinterleaveOp0,

        NarrowOp0 = 1 << 10,  // Replace operand 0 with its half-width equivalent.
        NarrowOp1 = 1 << 11,  // Same as above, but for operand 1.
        NarrowOp2 = 1 << 12,
        NarrowOp3 = 1 << 13,
        NarrowOps = NarrowOp0 | NarrowOp1 | NarrowOp2 | NarrowOp3,

        NarrowUnsignedOp0 = 1 << 15,  // Similar to the above, but narrow to an unsigned half width type.
        NarrowUnsignedOp1 = 1 << 16,
        NarrowUnsignedOp2 = 1 << 17,
        NarrowUnsignedOps = NarrowUnsignedOp0 | NarrowUnsignedOp1 | NarrowUnsignedOp2,

   };

    string intrin;        // Name of the intrinsic
    Expr pattern;         // The pattern to match against
    int flags;

    Pattern() {}
    Pattern(const string &intrin, Expr p, int flags = 0)
        : intrin(intrin), pattern(p), flags(flags) {}
};

Expr wild_u8 = Variable::make(UInt(8), "*");
Expr wild_u16 = Variable::make(UInt(16), "*");
Expr wild_u32 = Variable::make(UInt(32), "*");
Expr wild_u64 = Variable::make(UInt(64), "*");
Expr wild_i8 = Variable::make(Int(8), "*");
Expr wild_i16 = Variable::make(Int(16), "*");
Expr wild_i32 = Variable::make(Int(32), "*");
Expr wild_i64 = Variable::make(Int(64), "*");

Expr wild_u8x = Variable::make(Type(Type::UInt, 8, 0), "*");
Expr wild_u16x = Variable::make(Type(Type::UInt, 16, 0), "*");
Expr wild_u32x = Variable::make(Type(Type::UInt, 32, 0), "*");
Expr wild_u64x = Variable::make(Type(Type::UInt, 64, 0), "*");
Expr wild_i8x = Variable::make(Type(Type::Int, 8, 0), "*");
Expr wild_i16x = Variable::make(Type(Type::Int, 16, 0), "*");
Expr wild_i32x = Variable::make(Type(Type::Int, 32, 0), "*");
Expr wild_i64x = Variable::make(Type(Type::Int, 64, 0), "*");

// Attempt to apply one of the patterns to x. If a match is
// successful, the expression is replaced with a call using the
// matched operands. Prior to substitution, the matches are mutated
// with op_mutator.
Expr apply_patterns(Expr x, const vector<Pattern> &patterns, IRMutator *op_mutator) {
    vector<Expr> matches;
    for (const Pattern &p : patterns) {
        if (expr_match(p.pattern, x, matches)) {
            // The Pattern::Narrow*Op* flags are ordered such that
            // the operand corresponds to the bit (with operand 0
            // corresponding to the least significant bit), so we
            // can check for them all in a loop.
            bool is_match = true;
            for (size_t i = 0; i < matches.size() && is_match; i++) {
                Type t = matches[i].type();
                Type target_t = t.with_bits(t.bits()/2);
                if (p.flags & (Pattern::NarrowOp0 << i)) {
                    matches[i] = lossless_cast(target_t, matches[i]);
                } else if (p.flags & (Pattern::NarrowUnsignedOp0 << i)) {
                    matches[i] = lossless_cast(target_t.with_code(Type::UInt), matches[i]);
                }
                if (!matches[i].defined()) is_match = false;
            }
            if (!is_match) continue;

            for (size_t i = Pattern::BeginExactLog2Op; i < Pattern::EndExactLog2Op && is_match; i++) {
                // This flag is mainly to capture shifts. When the
                // operand of a div or mul is a power of 2, we can use
                // a shift instead.
                if (p.flags & (Pattern::ExactLog2Op1 << (i - Pattern::BeginExactLog2Op))) {
                    int pow;
                    if (is_const_power_of_two_integer(matches[i], &pow)) {
                        matches[i] = cast(matches[i].type().with_lanes(1), pow);
                    } else {
                        is_match = false;
                    }
                }
            }
            if (!is_match) continue;

            for (size_t i = Pattern::BeginDeinterleaveOp; i < Pattern::EndDeinterleaveOp; i++) {
                if (p.flags &
                    (Pattern::DeinterleaveOp0 << (i - Pattern::BeginDeinterleaveOp))) {
                    internal_assert(matches[i].type().is_vector());
                    matches[i] = native_deinterleave(matches[i]);
                }
            }
            if (p.flags & Pattern::SwapOps01) {
                internal_assert(matches.size() >= 2);
                std::swap(matches[0], matches[1]);
            }
            if (p.flags & Pattern::SwapOps12) {
                internal_assert(matches.size() >= 3);
                std::swap(matches[1], matches[2]);
            }
            // Mutate the operands with the given mutator.
            for (Expr &op : matches) {
                op = op_mutator->mutate(op);
            }
            x = Call::make(x.type(), p.intrin, matches, Call::PureExtern);
            if (p.flags & Pattern::InterleaveResult) {
                // The pattern wants us to interleave the result.
                x = native_interleave(x);
            }
            return x;
        }
    }
    return x;
}

// Replace x with a negated version of x, if it can be done without
// overflow.
Expr lossless_negate(Expr x) {
    const Mul *m = x.as<Mul>();
    if (m) {
        Expr a = lossless_negate(m->a);
        if (a.defined()) {
            return Mul::make(a, m->b);
        }
        Expr b = lossless_negate(m->b);
        if (b.defined()) {
            return Mul::make(m->a, b);
        }
    }
    if (is_negative_negatable_const(x) || is_positive_const(x)) {
        return simplify(-x);
    }
    return Expr();
}

template <typename T>
Expr apply_commutative_patterns(const T *op, const vector<Pattern> &patterns, IRMutator *mutator) {
    Expr ret = apply_patterns(op, patterns, mutator);
    if (!ret.same_as(op)) return ret;

    // Try commuting the op
    Expr commuted = T::make(op->b, op->a);
    ret = apply_patterns(commuted, patterns, mutator);
    if (!ret.same_as(commuted)) return ret;

    return op;
}

// Perform peephole optimizations on the IR, adding appropriate
// interleave and deinterleave calls.
class OptimizePatterns : public IRMutator {
private:
    using IRMutator::visit;

    void visit(const Mul *op) {
        static vector<Pattern> scalar_muls = {
            // Vector by scalar widening multiplies.
            { "halide.hexagon.mpy.vub.ub", wild_u16x*bc(wild_u16), Pattern::InterleaveResult | Pattern::NarrowOps },
            { "halide.hexagon.mpy.vub.b",  wild_i16x*bc(wild_i16), Pattern::InterleaveResult | Pattern::NarrowUnsignedOp0 | Pattern::NarrowOp1 },
            { "halide.hexagon.mpy.vuh.uh", wild_u32x*bc(wild_u32), Pattern::InterleaveResult | Pattern::NarrowOps },
            { "halide.hexagon.mpy.vh.h",   wild_i32x*bc(wild_i32), Pattern::InterleaveResult | Pattern::NarrowOps },

            // Multiplication by powers of 2.
            { "halide.hexagon.shl.vub.ub", wild_u8x*bc(wild_u8), Pattern::ExactLog2Op1 },
            { "halide.hexagon.shl.vuh.uh", wild_u16x*bc(wild_u16), Pattern::ExactLog2Op1 },
            { "halide.hexagon.shl.vuw.uw", wild_u32x*bc(wild_u32), Pattern::ExactLog2Op1 },
            { "halide.hexagon.shl.vb.b", wild_i8x*bc(wild_i8), Pattern::ExactLog2Op1 },
            { "halide.hexagon.shl.vh.h", wild_i16x*bc(wild_i16), Pattern::ExactLog2Op1 },
            { "halide.hexagon.shl.vw.w", wild_i32x*bc(wild_i32), Pattern::ExactLog2Op1 },

            // Non-widening scalar multiplication.
            { "halide.hexagon.mul.vh.b", wild_i16x*bc(wild_i16), Pattern::NarrowOp1 },
            { "halide.hexagon.mul.vw.h", wild_i32x*bc(wild_i32), Pattern::NarrowOp1 },
            // TODO: There's also mul.vw.b. We currently generate mul.vw.h
            // instead. I'm not sure mul.vw.b is faster, it might even be
            // slower due to the extra step in broadcasting the scalar up to
            // 32 bits.
        };

        static vector<Pattern> muls = {
            // Widening multiplication
            { "halide.hexagon.mpy.vub.vub", wild_u16x*wild_u16x, Pattern::InterleaveResult | Pattern::NarrowOps },
            { "halide.hexagon.mpy.vuh.vuh", wild_u32x*wild_u32x, Pattern::InterleaveResult | Pattern::NarrowOps },
            { "halide.hexagon.mpy.vb.vb",   wild_i16x*wild_i16x, Pattern::InterleaveResult | Pattern::NarrowOps },
            { "halide.hexagon.mpy.vh.vh",   wild_i32x*wild_i32x, Pattern::InterleaveResult | Pattern::NarrowOps },

            { "halide.hexagon.mpy.vub.vb",  wild_i16x*wild_i16x, Pattern::InterleaveResult | Pattern::NarrowUnsignedOp0 | Pattern::NarrowOp1 },
            { "halide.hexagon.mpy.vh.vuh",  wild_i32x*wild_i32x, Pattern::InterleaveResult | Pattern::NarrowOp0 | Pattern::NarrowUnsignedOp1 },
            // We need to check for the commuted versions of these patterns
            // before the more general patterns below catch these ops. The
            // other fix for this would be to break this into a third group of
            // multiply patterns, so the commuted versions of these would get
            // matched first.
            { "halide.hexagon.mpy.vub.vb",  wild_i16x*wild_i16x, Pattern::InterleaveResult | Pattern::NarrowOp0 | Pattern::NarrowUnsignedOp1 | Pattern::SwapOps01 },
            { "halide.hexagon.mpy.vh.vuh",  wild_i32x*wild_i32x, Pattern::InterleaveResult | Pattern::NarrowUnsignedOp0 | Pattern::NarrowOp1 | Pattern::SwapOps01 },

            // One operand widening multiplication.
            { "halide.hexagon.mul.vw.vh", wild_i32x*wild_i32x, Pattern::ReinterleaveOp0 | Pattern::NarrowOp1 },
            { "halide.hexagon.mul.vw.vuh", wild_i32x*wild_i32x, Pattern::ReinterleaveOp0 | Pattern::NarrowUnsignedOp1 },
            { "halide.hexagon.mul.vuw.vuh", wild_u32x*wild_u32x, Pattern::ReinterleaveOp0 | Pattern::NarrowUnsignedOp1 },
        };

        if (op->type.is_vector()) {
            expr = apply_commutative_patterns(op, scalar_muls, this);
            if (!expr.same_as(op)) return;

            expr = apply_commutative_patterns(op, muls, this);
            if (!expr.same_as(op)) return;
        }
        IRMutator::visit(op);
    }

    // Helpers to generate horizontally reducing multiply operations.
    static Expr halide_hexagon_add_2mpy(string suffix, Expr v0, Expr v1, Expr c0, Expr c1) {
        Type t = v0.type();
        Type result_type = Int(t.bits()*2).with_lanes(t.lanes());
        Expr call = Call::make(result_type, "halide.hexagon.add_2mpy" + suffix, {v0, v1, c0, c1}, Call::PureExtern);
        return native_interleave(call);
    }

    static Expr halide_hexagon_add_2mpy(string suffix, Expr v01, Expr c01) {
        Type t = v01.type();
        Type result_type = t.with_bits(t.bits()*2).with_lanes(t.lanes()/2);
        if (c01.type().is_int()) {
            result_type = result_type.with_code(Type::Int);
        }
        return Call::make(result_type, "halide.hexagon.add_2mpy" + suffix, {v01, c01}, Call::PureExtern);
    }

    static Expr halide_hexagon_add_4mpy(string suffix, Expr v01, Expr c01) {
        Type t = v01.type();
        Type result_type = t.with_bits(t.bits()*4).with_lanes(t.lanes()/4);
        if (c01.type().is_int()) {
            result_type = result_type.with_code(Type::Int);
        }
        return Call::make(result_type, "halide.hexagon.add_4mpy" + suffix, {v01, c01}, Call::PureExtern);
    }

    typedef pair<Expr, Expr> MulExpr;

    // In the operands vector, find elements of the operands list that
    // are widening multiplies from types a_ty and b_ty. If mpy_count
    // multiplies aren't found (but at least one is), it will attempt
    // to multiply remaining operands by 1 to match the target
    // widening multiply.
    static pair<vector<MulExpr>, Expr> find_mpy_ops(const vector<Expr> &operands, Type a_ty, Type b_ty, int mpy_count) {
        Type add_ty = operands.front().type();
        int mpy_bits = std::max(a_ty.bits(), b_ty.bits())*2;
        Expr a_wild = Variable::make(a_ty.with_bits(mpy_bits).with_code(add_ty.code()), "*");
        Expr b_wild = Variable::make(b_ty.with_bits(mpy_bits).with_code(add_ty.code()), "*");
        if (a_wild.type().lanes() == 1) {
            a_wild = Broadcast::make(a_wild, b_ty.lanes());
        }
        if (b_wild.type().lanes() == 1) {
            b_wild = Broadcast::make(b_wild, a_ty.lanes());
        }

        // Try to match both the multiply and its commutation.
        Expr pattern = a_wild*b_wild;
        Expr commuted_pattern = b_wild*a_wild;
        if (add_ty.bits() != pattern.type().bits()) {
            pattern = Cast::make(add_ty, pattern);
            commuted_pattern = Cast::make(add_ty, commuted_pattern);
        }

        // Go through the operands, putting the matching multiplies in
        // matched, and the rest in not_matched.
        vector<MulExpr> matched;
        vector<Expr> not_matched;
        for (Expr i : operands) {
            if ((int)matched.size() >= mpy_count) {
                not_matched.emplace_back(i);
                continue;
            }
            Expr a, b;
            vector<Expr> matches;
            if (expr_match(pattern, i, matches)) {
                a = matches[0];
                b = matches[1];
            } else if (expr_match(commuted_pattern, i, matches)) {
                a = matches[1];
                b = matches[0];
            }

            if (a.defined()) {
                a = lossless_cast(a_ty.with_lanes(a.type().lanes()), a);
            }
            if (b.defined()) {
                b = lossless_cast(b_ty.with_lanes(b.type().lanes()), b);
            }

            if (a.defined() && b.defined()) {
                matched.emplace_back(a, b);
            } else {
                not_matched.emplace_back(i);
            }
        }

        if (matched.empty()) {
            // Don't fake a bunch of mpy ops if there are no real ones.
            return {matched, make_binary_op<Add>(not_matched)};
        }

        // If we didn't find mpy_count matches, go through the
        // not_matched list and try to find elements that match one of
        // the operand types, and set the other operand to 1.
        vector<Expr> rest;
        for (Expr i : not_matched) {
            if ((int)matched.size() < mpy_count) {
                Expr as_a = lossless_cast(a_ty.with_lanes(i.type().lanes()), i);
                Expr as_b = lossless_cast(b_ty.with_lanes(i.type().lanes()), i);
                if (as_a.defined()) {
                    matched.emplace_back(as_a, make_one(b_ty));
                } else if (as_b.defined()) {
                    matched.emplace_back(make_one(a_ty), as_b);
                } else {
                    rest.push_back(i);
                }
            } else {
                rest.push_back(i);
            }
        }

        internal_assert((int)matched.size() <= mpy_count);

        return {matched, make_binary_op<Add>(rest)};
    }

    void visit(const Add *op) {
        vector<Expr> matches;

        // vmpa, vdmpy, and vrmpy instructions are hard to match with
        // patterns, do it manually here.
        vector<Expr> operands;
        flatten_binary_op(op, operands);

        // Try to find vrmpy opportunities first, which consume 4 operands.
        if (op->type.is_vector() && operands.size() >= 4) {
            int lanes = op->type.lanes();
            vector<MulExpr> mpys;
            Expr rest;
            string suffix;

            // Try to find a vector*scalar multiply first, which will
            // match a subset of the expressions that vector*vector
            // matches.
            if (op->type.is_uint()) {
                std::tie(mpys, rest) = find_mpy_ops(operands, UInt(8).with_lanes(lanes), UInt(8), 4);
                suffix = ".vub.ub";
            } else {
                std::tie(mpys, rest) = find_mpy_ops(operands, UInt(8).with_lanes(lanes), Int(8), 4);
                suffix = ".vub.b";
            }
            if (mpys.size() == 4) {
                // TODO: It's possible that permuting the order of the
                // multiply operands can simplify the shuffle away.
                Expr a0123 = Shuffle::make_interleave({mpys[0].first, mpys[1].first, mpys[2].first, mpys[3].first});
                Expr b0123 = Shuffle::make_interleave({mpys[0].second, mpys[1].second, mpys[2].second, mpys[3].second});
                a0123 = simplify(a0123);
                b0123 = simplify(b0123);
                b0123 = reinterpret(Type(b0123.type().code(), 32, 1), b0123);
                expr = halide_hexagon_add_4mpy(suffix, a0123, b0123);
                if (rest.defined()) {
                    expr = Add::make(expr, rest);
                }
                expr = mutate(expr);
                return;
            }

            // Now try to match vector*vector vrmpy expressions.
            if (op->type.is_uint()) {
                std::tie(mpys, rest) = find_mpy_ops(operands, UInt(8).with_lanes(lanes), UInt(8).with_lanes(lanes), 4);
                suffix = ".vub.vub";
            } else {
                std::tie(mpys, rest) = find_mpy_ops(operands, UInt(8).with_lanes(lanes), Int(8).with_lanes(lanes), 4);
                suffix = ".vub.vb";
            }
            // TODO: suffix = ".vb.vb"
            if (mpys.size() == 4) {
                // TODO: It's possible that permuting the order of the
                // multiply operands can simplify the shuffle away.
                Expr a0123 = Shuffle::make_interleave({mpys[0].first, mpys[1].first, mpys[2].first, mpys[3].first});
                Expr b0123 = Shuffle::make_interleave({mpys[0].second, mpys[1].second, mpys[2].second, mpys[3].second});
                a0123 = simplify(a0123);
                b0123 = simplify(b0123);
                expr = halide_hexagon_add_4mpy(suffix, a0123, b0123);
                if (rest.defined()) {
                    expr = Add::make(expr, rest);
                }
                expr = mutate(expr);
                return;
            }
        }

        // Find opportunities vdmpy or vmpa.
        if (op->type.is_vector() && operands.size() >= 2) {
            int lanes = op->type.lanes();

            vector<MulExpr> mpys;
            Expr rest;
            string vmpa_suffix;
            string vdmpy_suffix;

{'fraction_non_alphanumeric': 0.13311451495258936, 'fraction_numerical': 0.021474235208306516, 'mean_word_length': 2.226621443898387, 'pattern_counts': {'":': 0, '<': 177, '<?xml version=': 0, '>': 129, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '203149', 'n_tokens_mistral': 16052, 'n_tokens_neox': 14951, 'n_words': 3568}