mem_checks.cpp

Go to the documentation of this file.

#include "mim/plug/gpu/phase/mem_checks.h"
 
#include <mim/axm.h>
 
#include <mim/plug/mem/mem.h>
 
namespace mim::plug::gpu::phase {
 
namespace {
 
class MemFinder {
public:
    MemFinder(const Def* to_search)
        : to_search({to_search}) {}
    MemFinder(DefVec&& to_search)
        : to_search(std::move(to_search)) {}
 
    using IsaMem = Axm::IsA<mem::M, mim::App>;
 
    IsaMem next_mem();
 
private:
    DefVec to_search;
};
 
MemFinder::IsaMem MemFinder::next_mem() {
    while (!to_search.empty()) {
        auto cur = to_search.back();
        to_search.pop_back();
        if (cur->isa<Sigma>())
            for (auto op : cur->ops())
                to_search.push_back(op);
        else if (auto mem = Axm::isa<mem::M>(cur))
            return mem;
    }
    return IsaMem();
}
 
} // namespace
 
const Def* MemChecks::rewrite_imm_App(const App* app) {
    if (auto launch = Axm::isa<gpu::launch>(app)) {
        auto kernel        = launch->decurry()->decurry()->arg();
        auto kernel_args   = launch->decurry()->arg();
        auto kernel_args_t = kernel_args->type();
 
        MemFinder mem_finder(kernel_args_t);
        if (mem_finder.next_mem()) {
            launch
                ->blame("`mem.M` must not be passed across device boundaries: `{}` of type `{}` is passed from host "
                        "to kernel `{}`",
                        kernel_args, kernel_args_t, kernel)
                .bail();
        }
    }
    return Super::rewrite_imm_App(app);
}
 
void MemChecks::rewrite_external(Def* def) {
    auto lam = def->isa<Lam>();
    if (lam && lam->sym().str() == "main") {
        MemFinder intype_mem_finder(lam->type()->dom());
        while (auto mem = intype_mem_finder.next_mem()) {
            auto addr_space = mem->arg();
            if (Lit::as(addr_space) != 0)
                lam->blame("the `main` function must not take a `mem.M n` with a non-zero `n` as an argument").bail();
        }
 
        MemFinder outtype_mem_finder(lam->type()->ret_dom());
        while (auto mem = outtype_mem_finder.next_mem()) {
            auto addr_space = mem->arg();
            if (Lit::as(addr_space) != 0)
                lam->blame("the `main` function must not return a `mem.M n` with a non-zero `n`").bail();
        }
    }
    Super::rewrite_external(def);
}
 
} // namespace mim::plug::gpu::phase