ll.h

Go to the documentation of this file.

#pragma once
 
#include <bit>
#include <deque>
#include <format>
#include <fstream>
#include <iterator>
#include <optional>
#include <ostream>
#include <set>
#include <string>
 
#include <fe/term.h>
 
#include <mim/driver.h>
#include <mim/phase.h>
#include <mim/schedule.h>
 
#include <mim/plug/clos/clos.h>
#include <mim/plug/math/math.h>
#include <mim/plug/mem/mem.h>
#include <mim/plug/vec/vec.h>
 
#include "mim/plug/ll/autogen.h"
 
// Lessons learned:
// * **Always** follow all ops - even if you actually want to ignore one.
//   Otherwise, you might end up with an incorrect schedule.
//   This was the case for an Extract of type Mem.
//   While we want to ignore the value obtained from that, since there is no Mem value in LLVM,
//   we still want to **first** recursively emit code for its operands and **then** ignore the Extract itself.
// * i1 has a different meaning in LLVM then in Mim:
//      * Mim: {0,  1} = i1
//      * LLVM:   {0, -1} = i1
//   This is a problem when, e.g., using an index of type i1 as LLVM thinks like this:
//   getelementptr ..., i1 1 == getelementptr .., i1 -1
namespace mim {
 
class World;
 
namespace plug::ll {
 
/// Prefix for this backend's diagnostics; concatenate it with the format literal.
/// @note Not usable with the Def::expect family: those take the *inner* "what was expected" description.
#define MIM_LL_BE "ll backend: "
 
namespace math = mim::plug::math;
namespace mem  = mim::plug::mem;
 
namespace detail {
inline const char* math_suffix(const Def* type) {
    if (auto w = math::isa_f(type)) {
        switch (*w) {
            case 32: return "f";
            case 64: return "";
        }
    }
    type->blame(MIM_LL_BE "unsupported floating-point type `{}`", type).bail();
}
 
inline const char* llvm_suffix(const Def* type) {
    if (auto w = math::isa_f(type)) {
        switch (*w) {
            case 16: return ".f16";
            case 32: return ".f32";
            case 64: return ".f64";
        }
    }
    type->blame(MIM_LL_BE "unsupported floating-point type `{}`", type).bail();
}
 
// [mem.M 0, T] => T
// TODO there may be more instances where we have to deal with this trickery
inline const Def* isa_mem_sigma_2(const Def* type) {
    if (auto sigma = type->isa<Sigma>())
        if (sigma->num_ops() == 2 && Axm::isa<mem::M>(sigma->op(0))) return sigma->op(1);
    return {};
}
} // namespace detail
 
struct BB {
    BB()                    = default;
    BB(const BB&)           = delete;
    BB(BB&& other) noexcept = default;
    BB& operator=(BB other) noexcept { return swap(*this, other), *this; }
 
    std::deque<std::ostringstream>& head() { return parts[0]; }
    std::deque<std::ostringstream>& body() { return parts[1]; }
    std::deque<std::ostringstream>& tail() { return parts[2]; }
 
    template<class... Args>
    inline std::string assign(std::string_view name, std::format_string<Args...> s, Args&&... args) {
        auto& os = body().emplace_back();
        std::print(os, "{} = ", name);
        std::print(os, s, std::forward<Args>(args)...);
        return std::string(name);
    }
 
    template<class... Args>
    inline void tail(std::format_string<Args...> s, Args&&... args) {
        std::print(tail().emplace_back(), s, std::forward<Args>(args)...);
    }
 
    friend inline void swap(BB& a, BB& b) noexcept {
        using std::swap;
        swap(a.phis, b.phis);
        swap(a.parts, b.parts);
    }
 
    DefMap<std::deque<std::pair<std::string, std::string>>> phis;
    std::array<std::deque<std::ostringstream>, 3> parts;
};
 
class Emitter;
 
/// The heavy, target-independent emitter methods are compiled once into `libmim_ll` (see ll.cpp).
/// `libmim_nvptx` reaches them through these `extern "C"` shims instead of recompiling the bodies;
/// Emitter caches the pointers via `GET_FUN_PTR` in its constructor.
extern "C" {
MIM_EXPORT void mim_ll_convert(Emitter&, const Def*, bool simd, std::string& res);
MIM_EXPORT void mim_ll_finalize(Emitter&);
MIM_EXPORT void mim_ll_emit_epilogue(Emitter&, Lam*);
MIM_EXPORT void mim_ll_emit_bb(Emitter&, BB&, const Def*, std::string& res);
}
 
class Emitter : public NestPhase<Lam> {
public:
    using Super = NestPhase<Lam>;
 
    Emitter(World& world, std::string name, std::ostream& ostream)
        : Super(world, std::move(name), false)
        , ostream_(ostream) {
        auto& driver = world.driver();
        // Ensure libmim_ll is loaded so the shims below resolve (e.g. when a derived backend like
        // nvptx uses us). Loading merely registers ll.emit; it does not run it.
        if (!driver.is_loaded("ll")) driver.load("ll");
        convert_       = driver.GET_FUN_PTR("ll", mim_ll_convert);
        finalize_      = driver.GET_FUN_PTR("ll", mim_ll_finalize);
        emit_epilogue_ = driver.GET_FUN_PTR("ll", mim_ll_emit_epilogue);
        emit_bb_       = driver.GET_FUN_PTR("ll", mim_ll_emit_bb);
    }
 
    fe::Tab tab = fe::Tab::spaces();
 
    void start() override;
    void visit(const Nest&) override;
    void emit_imported(Lam*);
    virtual std::string prepare();
 
    // Thin forwarders into `libmim_ll`; the real bodies (`*_impl`) live in ll.cpp.
    virtual void emit_epilogue(Lam* lam) { emit_epilogue_(*this, lam); }
    void finalize() { finalize_(*this); }
    std::string emit_bb(BB& bb, const Def* def) {
        std::string res;
        emit_bb_(*this, bb, def, res);
        return res;
    }
 
    virtual inline std::optional<std::string> isa_targetspecific_intrinsic(BB&, const Def*) { return std::nullopt; }
 
    template<class... Args>
    void declare(std::format_string<Args...> s, Args&&... args) {
        decls_.emplace("declare " + std::format(s, std::forward<Args>(args)...));
    }
 
    /// How the C runtime wrappers (compiled to a `<name>.ll` via `add_mim_runtime`) reach the output.
    enum class Rt {
        embed, ///< Splice the wrapper IR into the emitted module so it is self-contained.
        ext,   ///< Only `declare` the wrappers; the runtime is linked in externally.
    };
 
    void rt_mode(Rt rt) { rt_ = rt; }
 
    /// Locates the runtime module `rt/<filename>` (produced by `add_mim_runtime`) in the driver's
    /// search paths, reads it, and stores it for `Rt::embed` splicing.
    /// Backends share this instead of duplicating the lookup: `ll` loads `ll_rt.ll`, `nvptx`
    /// loads `nvptx_rt.ll`, and so on — one merged module per plugin.
    /// @returns whether the module was found.
    bool load_rt_module(std::string_view filename);
 
    /// Declares a runtime wrapper @p sig (implemented in a C runtime, see `add_mim_runtime`) and
    /// records that the runtime is required by this module.
    /// In `Rt::ext` mode the declaration is emitted like any other `declare`.
    /// In `Rt::embed` mode the wrapper's *definition* is spliced into the output, so emitting a
    /// `declare` as well would be a redefinition — hence it is suppressed here.
    template<class... Args>
    void declare_rt(std::format_string<Args...> sig, Args&&... args) {
        rt_used_ = true;
        if (rt_ == Rt::ext) declare(sig, std::forward<Args>(args)...);
    }
 
protected:
    std::ostream& ostream() const { return ostream_; }
 
    /// Recursively emits code.
    /// `mem`-typed @p def%s yield an empty string, which this variant asserts against.
    std::string emit(const Def* def) {
        auto res = emit_unsafe(def);
        assert(!res.empty());
        return res;
    }
 
    /// As above but yielding an empty string is permitted.
    std::string emit_unsafe(const Def* def) {
        if (auto s = fe::lookup(locals_, def)) return *s;
 
        auto place          = scheduler_.smart(curr_lam_, def);
        auto& bb            = lam2bb_[place->mut()->as<Lam>()];
        auto val            = emit_bb(bb, def);
        return locals_[def] = val;
    }
 
    /// The Scheduler::schedule of the function currently being emitted; see Emitter::visit.
    const Scheduler::Schedule& schedule() const { return schedule_; }
 
    std::string id(const Def*, bool force_bb = false) const;
 
    /// Converts @p type to its LLVM representation as a value (@p simd) or in memory (`!simd`).
    /// Only values use LLVM vectors; memory keeps the plain array layout C expects.
    virtual std::string convert(const Def* type, bool simd = true) {
        std::string res;
        convert_(*this, type, simd, res);
        return res;
    }
    std::string convert_ret_pi(const Pi*);
 
    /// Lane count and element type if a value of @p type is an LLVM vector.
    std::optional<std::pair<nat_t, const Def*>> vectorize(const Def* type) const;
    /// Whether @p type's value representation contains an LLVM vector and hence differs from its memory one.
    bool has_vec(const Def* type) const;
 
    /// Registers @p arg as an incoming phi value for @p phi in @p callee, coming from predecessor @p pred.
    void emit_phi(Lam* callee, const Def* phi, std::string arg, Lam* pred) {
        // Inserting here would invalidate the `BB&` our caller holds.
        auto i = lam2bb_.find(callee);
        if (i == lam2bb_.end())
            callee->blame(MIM_LL_BE "`{}` jumps to `{}`, which is not a basic block of `{}`", pred, callee, root())
                .bail();
        i->second.phis[phi].emplace_back(std::move(arg), id(pred, true));
        locals_[phi] = id(phi);
    }
 
    /// Wires all non-`mem.M` arguments of @p app into @p callee's phis, coming from predecessor @p pred.
    void emit_phi_args(Lam* callee, const App* app, Lam* pred) {
        size_t n = callee->num_tvars();
        for (size_t i = 0; i != n; ++i)
            if (auto arg = emit_unsafe(app->arg(n, i)); !arg.empty()) {
                auto phi = callee->var(n, i);
                if (Axm::isa<mem::M>(phi->type())) continue;
                emit_phi(callee, phi, std::move(arg), pred);
            }
    }
 
    /// Emits the storage backing a `mem.slot` of type @p pointee and yields the pointer value.
    /// The generic backend allocates on the stack; targets may override (e.g. a global in a specific address space).
    /// Kept inline on purpose so `Emitter` retains no vtable key function (else its vtable would live in a single
    /// module and break derived backends loaded from a separate plugin).
    virtual std::string emit_slot(BB& bb, const App* app, const Def* pointee, const Def* /*addr_space*/) {
        auto v_ptr = "%" + app->unique_name() + ".slot";
        std::print(bb.body().emplace_back(), "{} = alloca {}", v_ptr, convert(pointee, false));
        return v_ptr;
    }
 
    Lam* curr_lam_ = nullptr;
    std::ostream& ostream_;
    Scheduler scheduler_;
    Scheduler::Schedule schedule_;
    DefMap<std::string> locals_;
    std::array<DefMap<std::string>, 2> types_; ///< Indexed by the `simd` flag of Emitter::convert.
    LamMap<BB> lam2bb_;
 
    std::set<std::string> decls_;
    std::ostringstream type_decls_;
    std::ostringstream vars_decls_;
    std::ostringstream func_decls_;
    std::ostringstream func_impls_;
    LamMap<const Def*> simd_phi_;
 
    /// Loop-metadata node id per `ll.vec`/`ll.novec`-annotated loop header, and whether it may vectorize;
    /// `LoopMdBase` and `LoopMdBase + 1` hold the shared enable/disable hints, the per-loop nodes follow,
    /// clear of the ids the embedded runtime module brings along.
    static constexpr u64 LoopMdBase = 1000;
    LamMap<std::pair<u64, bool>> loop_md_;
 
    bool simd_    = true;
    Rt rt_        = Rt::embed;
    bool rt_used_ = false;
    std::string rt_module_;
 
private:
    // Real implementations; defined in ll.cpp and exported via the `mim_ll_*` shims above.
    std::string convert_impl(const Def*, bool simd);
    void finalize_impl();
    void emit_epilogue_impl(Lam*);
    std::string emit_bb_impl(BB&, const Def*);
 
    // Case groups of emit_bb_impl, split so one recursion level only pays the frame of the group it hits.
    MIM_NOINLINE std::string emit_lit(const Def*);
    MIM_NOINLINE std::string emit_tuple(BB&, const std::string& name, const Def* tuple);
    MIM_NOINLINE std::pair<std::string, std::string> emit_gep_index(BB&, const std::string& name, const Def* index);
    MIM_NOINLINE std::optional<std::string> emit_builtin(BB&, const std::string& name, const Def*);
    MIM_NOINLINE std::optional<std::string> emit_core(BB&, const std::string& name, const Def*);
    MIM_NOINLINE std::optional<std::string> emit_mem(BB&, const std::string& name, const Def*);
    MIM_NOINLINE std::optional<std::string> emit_math(BB&, const std::string& name, const Def*);
    MIM_NOINLINE std::optional<std::string> emit_vec(BB&, const std::string& name, const Def*);
 
    /// An `extractelement`/`insertelement` lane index; LLVM wants a plain `i32`, not the `Idx` carrier.
    std::string emit_simd_index(BB&, const std::string& name, const Def* index);
    /// Converts the value @p v of @p type into its value (@p simd) or memory (`!simd`) representation.
    std::string repack(BB&, const std::string& name, const Def* type, const std::string& v, bool simd);
    const Def* find_common_simd_src(const App*) const;
 
    decltype(&mim_ll_convert) convert_             = nullptr;
    decltype(&mim_ll_finalize) finalize_           = nullptr;
    decltype(&mim_ll_emit_epilogue) emit_epilogue_ = nullptr;
    decltype(&mim_ll_emit_bb) emit_bb_             = nullptr;
 
    friend void mim_ll_convert(Emitter&, const Def*, bool, std::string&);
    friend void mim_ll_finalize(Emitter&);
    friend void mim_ll_emit_epilogue(Emitter&, Lam*);
    friend void mim_ll_emit_bb(Emitter&, BB&, const Def*, std::string&);
};
 
/// Parses `-X <backend>:rt=extern|embed`; Emitter::Rt::embed is the default.
inline Emitter::Rt arg_rt(fe::View<std::string> args) {
    return arg_value(args, "rt") == "extern" ? Emitter::Rt::ext : Emitter::Rt::embed;
}
 
/*
 * convert
 */
 
inline static std::optional<std::pair<nat_t, const Def*>> is_simd(const Def* type) {
    if (auto arr = type->isa<Arr>(); arr && !arr->shape().is_fused()) {
        // LLVM rounds a vector's store size up to a power of two, so only then does `<n x T>` occupy the
        // `n * sizeof(T)` bytes that `mem`'s size arithmetic assumes.
        if (auto l = Lit::isa(arr->arity()); l && std::has_single_bit(*l)) {
            if (arr->body()->isa<Nat>() || Idx::isa(arr->body()) || Axm::isa<math::F>(arr->body()))
                return std::pair{*l, arr->body()};
        }
    }
    return {};
}
 
inline std::optional<std::pair<nat_t, const Def*>> Emitter::vectorize(const Def* type) const {
    return simd_ ? is_simd(type) : std::nullopt;
}
 
inline bool Emitter::has_vec(const Def* type) const {
    if (!simd_) return false;
    if (is_simd(type)) return true;
    if (auto arr = type->isa<Arr>()) return has_vec(arr->elem());
    if (auto sigma = type->isa<Sigma>())
        return std::ranges::any_of(sigma->ops(), [this](auto t) { return has_vec(t); });
    return false;
}
 
inline const Def* Emitter::find_common_simd_src(const App* app) const {
    const Def* common_src = nullptr;
    size_t lane           = 0;
    for (auto arg : app->args()) {
        if (Axm::isa<mem::M>(arg->type())) continue;
        auto extract = arg->isa<Extract>();
        if (!extract || !vectorize(extract->tuple()->type())) return nullptr;
        // Only devectorized args - lane i in position i - may forward the whole vector;
        // anything else (e.g. a Select with a non-literal index) must stay scalar.
        if (auto index = Lit::isa(extract->index()); !index || *index != lane++) return nullptr;
        if (!common_src)
            common_src = extract->tuple();
        else if (common_src != extract->tuple())
            return nullptr;
    }
    if (common_src && vectorize(common_src->type())->first != lane) return nullptr;
    return common_src;
}
 
inline std::string Emitter::id(const Def* def, bool force_bb /*= false*/) const {
    if (auto lam = def->isa_mut<Lam>(); lam && !force_bb) {
        if (lam->type()->ret_pi()) {
            if (lam->is_external() || !lam->is_set())
                return std::string("@") + lam->sym().str(); // TODO or use is_internal or sth like that?
            return std::string("@") + lam->unique_name();
        }
    }
 
    return std::string("%") + def->unique_name();
}
 
inline std::string Emitter::convert_ret_pi(const Pi* pi) {
    auto dom = mem::strip_mem_ty(pi->dom());
    if (dom == world().sigma()) return "void";
    return convert(dom);
}
 
/*
 * emit
 */
 
inline void Emitter::start() {
    Super::start();
 
    auto sep     = "";
    auto section = [this, &sep](std::string_view s) {
        while (s.ends_with('\n'))
            s.remove_suffix(1);
        if (s.empty()) return;
        std::println(ostream(), "{}{}", sep, s);
        sep = "\n";
    };
 
    // Splice the runtime wrapper module first (it carries the module's target triple/datalayout).
    if (rt_used_ && rt_ == Rt::embed) {
        if (rt_module_.empty())
            fe::throwf(MIM_LL_BE
                       "`-X ll:rt=embed` needs the runtime module `mim_rt.ll`, but it "
                       "was not found (build with clang / `MIM_BUILD_LL_RUNTIME=ON`, or use `-X ll:rt=extern`)");
        section(rt_module_);
    }
 
    section(type_decls_.str());
 
    std::ostringstream decls;
    for (auto&& decl : decls_)
        std::println(decls, "{}", decl);
    section(decls.str());
 
    section(func_decls_.str());
    section(vars_decls_.str());
    section(func_impls_.str());
 
    // One distinct `!llvm.loop` node per hinted loop header, sharing the vectorize-enable or -disable hint.
    if (!loop_md_.empty()) {
        std::ostringstream md_decls;
        std::println(md_decls, "!{} = !{{!\"llvm.loop.vectorize.enable\", i1 true}}", LoopMdBase);
        std::println(md_decls, "!{} = !{{!\"llvm.loop.vectorize.enable\", i1 false}}", LoopMdBase + 1);
        for (const auto& [_, md] : loop_md_)
            std::println(md_decls, "!{} = distinct !{{!{}, !{}}}", md.first, md.first,
                         md.second ? LoopMdBase : LoopMdBase + 1);
        section(md_decls.str());
    }
}
 
inline void Emitter::visit(const Nest& nest) {
    if (!root()->is_set()) return emit_imported(root());
 
    schedule_        = Scheduler::schedule(nest); // cached; finalize needs the very same schedule
    const auto& muts = schedule_;
 
    // make sure that we don't need to rehash later on
    for (auto mut : muts)
        if (auto lam = mut->isa<Lam>()) lam2bb_.try_emplace(lam, BB());
    auto old_size = lam2bb_.size();
 
    if (!root()->ret_var())
        root()->blame(MIM_LL_BE "top-level function is not a continuation with a return continuation").bail();
 
    prepare();
 
    Scheduler new_scheduler(nest);
    swap(scheduler_, new_scheduler);
 
    for (auto mut : muts) {
        if (auto lam = mut->isa<Lam>()) {
            curr_lam_ = lam;
            if (lam != root() && !Lam::isa_basicblock(lam))
                lam->blame(MIM_LL_BE "neither the entry nor a basic block of `{}`; it needs a phase that "
                                     "removes higher-order functions first",
                           root())
                    .bail();
            emit_epilogue(lam);
        }
    }
 
    finalize();
    locals_.clear();
    assert_unused(lam2bb_.size() == old_size && "really make sure we didn't trigger a rehash");
    // A BB never crosses a function boundary: Nest::contains is `def->has_free_vars_in(vars())`,
    // so a *closed* Lam is never a member of another Lam's Nest - and it cannot belong to two Nests either,
    // since a Lam free in the vars of two closed Lams would make the outer one open.
    // Every `BB&` handed out by emit_unsafe died with the calls above, so clearing here is safe.
    // Without it, finalize re-walks the BBs of all previously emitted functions - O(n²) in program size.
    lam2bb_.clear();
}
 
inline bool Emitter::load_rt_module(std::string_view filename) {
    for (const auto& dir : world().driver().rt_paths()) {
        auto path = dir / std::string(filename);
        std::error_code ec;
        if (!std::filesystem::is_regular_file(path, ec) || ec) continue;
        if (auto ifs = std::ifstream(path)) {
            world().log().d("ll backend: load runtime module `{}`", path.string());
            rt_module_.assign(std::istreambuf_iterator<char>(ifs), std::istreambuf_iterator<char>());
            return true;
        }
    }
    return false;
}
 
inline void Emitter::emit_imported(Lam* lam) {
    // TODO merge with declare method
    std::print(func_decls_, "declare {} {}(", convert_ret_pi(lam->type()->ret_pi()), id(lam));
 
    auto doms = lam->doms();
    for (auto sep = ""; auto dom : doms.view().rsubspan(1)) {
        if (Axm::isa<mem::M>(dom)) continue;
        std::print(func_decls_, "{}{}", sep, convert(dom));
        sep = ", ";
    }
 
    std::println(func_decls_, ")");
}
 
inline std::string Emitter::prepare() {
    auto internal = root()->is_external() ? "" : "internal ";
    auto ret_t    = convert_ret_pi(root()->type()->ret_pi());
    std::print(func_impls_, "define {}{} {}(", internal, ret_t, id(root()));
 
    auto vars = root()->vars();
    for (auto sep = ""; auto var : vars.view().rsubspan(1)) {
        if (Axm::isa<mem::M>(var->type())) continue;
        if (auto sigma = var->type()->isa<Sigma>(); sigma && sigma->num_ops() == 0) continue;
        auto name    = id(var);
        locals_[var] = name;
        std::print(func_impls_, "{}{} {}", sep, convert(var->type()), name);
        sep = ", ";
    }
 
    std::println(func_impls_, ") {{");
    return root()->unique_name();
}
 
} // namespace plug::ll
} // namespace mim