reassoc.h

Go to the documentation of this file.

#pragma once
 
#include <optional>
 
#include <mim/phase.h>
 
namespace mim::plug::tensor::phase {
 
/// One node of a bracketing: `i … j` splits after `s`.
/// Spelling the fields out keeps `size_t` out of the type - it is *not* mim::u64 everywhere.
/// A transposed node @p t is computed as the product of its operands' transposes, `(Rᵀ · Lᵀ)ᵀ`.
struct Split {
    u64 i, s, j;
    bool t = false;
};
 
/// A bracketing of a matrix chain, innermost node first.
using Splits = fe::Vector<Split>;
 
/// A chain operand; the chain reads its transpose if @p t.
struct Leaf {
    const Def* def;
    bool t;
};
 
using Leaves = fe::Vector<Leaf>;
 
/// Reassociates chains of tensor.product_2d with the classic matrix-chain-order dynamic program,
/// so that a chain is evaluated with the least number of vector-lane slots.
///
/// Extents need not be literal: a cost is kept as a polynomial in the symbolic extents, ordered by
/// coefficient-wise `≤`.
/// Since extents are `Nat`s and hence non-negative, that order proves `≤` under *every* instantiation -
/// but it is only a *partial* order, so a chain can have several bracketings that each win for some
/// instantiation (a batch dimension favouring a different one when small than when large).
/// Up to Reassoc::max_dispatch_ matrices those survivors are emitted side by side behind a runtime
/// comparison of their costs; a longer chain is only reassociated where one bracketing provably wins.
class Reassoc : public RWPhase {
public:
    using RWPhase::RWPhase;
 
private:
    static constexpr u64 Default_max_dispatch = 4;
    static constexpr u64 Default_vec          = 8;
 
    /// One `tensor.product_2d` of a chain: `«m, k» · «k, l»`.
    struct Link {
        const App* app;
        const Def* m;
        const Def* k;
        const Def* l;
    };
 
    void start() override;
    const Def* rewrite_imm_App(const App*) final;
 
    const Def* reassoc(const App*);
    std::optional<Link> isa_link(const Def* def, const Def* ring) const;
 
    /// Appends the chain's leaves to @p mats, each leaf's *row* extent to @p dims, and the bracketing as
    /// written to @p orig; the caller appends the chain's trailing column extent and its own split.
    void flatten(const Def* def, const Def* ring, const Def* rows, bool t, Leaves& mats, DefVec& dims, Splits& orig);
 
    /// Rebuilds `mats[i … j]` in the new world, parenthesized according to @p split (indexed `i * n + j`).
    const Def* build(const Def* head, fe::View<Leaf> mats, Defs dims, fe::View<u64> split, u64 i, u64 j);
 
    /// Emits every bracketing in @p cands as a thunk and selects the cheapest one by comparing their costs
    /// at run time.
    const Def* dispatch(const Def* head, const Def* res_ty, fe::View<Leaf> mats, Defs dims, fe::View<Splits> cands);
 
    /// The number of lane slots @p splits costs, as a `Nat` expression in the new world.
    const Def* cost_expr(fe::View<Leaf> mats, Defs dims, const Splits& splits);
 
    /// Old-world consumer count per `product_2d` and 2-D transpose app, attributed through tuple wrappers.
    DefMap<u64> consumers_;
 
    /// Longest chain whose bracketings are enumerated - and, failing a unique winner, dispatched over.
    /// The number of bracketings is `Catalan(n − 1)`, so this cannot grow much.
    /// Set with `-X tensor:reassoc-max=<n>`; below `3` nothing is ever dispatched.
    u64 max_dispatch_ = Default_max_dispatch;
 
    /// Lanes of the vector loop, which `dot_schedule` runs over each product's trailing extent - or,
    /// for a transposed operand, `dot_schedule_kvec` over its contraction.
    /// A literal extent is charged rounded up to a whole number of these, so a bracketing whose
    /// intermediates are too narrow to fill a vector pays for the lanes it leaves idle.
    /// Taken from `tensor.vec_width`, i.e. `-X tensor:vec-width=<n>`.
    u64 vec_ = Default_vec;
};
 
} // namespace mim::plug::tensor::phase