AMReX_Parser.H

Go to the documentation of this file.

#ifndef AMREX_PARSER_H_
#define AMREX_PARSER_H_
 
#include <AMReX_Arena.H>
#include <AMReX_Array.H>
#include <AMReX_GpuDevice.H>
#include <AMReX_Parser_Exe.H>
#include <AMReX_REAL.H>
#include <AMReX_Vector.H>
 
#include <concepts>
#include <map>
#include <memory>
#include <string>
#include <set>
#include <utility>
 
namespace amrex {
 
namespace parser_detail {
 
// Store the arguments in an array.  On the host, volatile keeps the stores
// scalar so that parser_exe_eval's loads can forward from them.
template <typename... Ts>
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void store_args (double* l_var, Ts... var) noexcept
{
    AMREX_IF_ON_DEVICE((
        int i = 0;
        ([&] { l_var[i++] = static_cast<double>(var); }(), ...);
    ))
    AMREX_IF_ON_HOST((
        volatile double* p = l_var;
        int i = 0;
        ([&] { p[i++] = static_cast<double>(var); }(), ...);
    ))
}
 
}
 
template <int N>
struct ParserExecutor
{
    [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
    double operator() () const noexcept requires (N==0)
    {
        AMREX_IF_ON_DEVICE((return parser_exe_eval(m_device_executor, nullptr);))
        AMREX_IF_ON_HOST((return eval_host(nullptr);))
    }
 
    template <typename... Ts>
    requires (N > 0 && sizeof...(Ts) == N && !amrex::Same<float,Ts...>::value)
    [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
    double operator() (Ts... var) const noexcept
    {
        double l_var[N];
        parser_detail::store_args(l_var, var...);
        AMREX_IF_ON_DEVICE((return parser_exe_eval(m_device_executor, l_var);))
        AMREX_IF_ON_HOST((return eval_host(l_var);))
    }
 
    template <std::same_as<float>... Ts>
    requires (N > 0 && sizeof...(Ts) == N)
    [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
    float operator() (Ts... var) const noexcept
    {
        double l_var[N];
        parser_detail::store_args(l_var, var...);
        AMREX_IF_ON_DEVICE((return static_cast<float>(parser_exe_eval(m_device_executor, l_var));))
        AMREX_IF_ON_HOST((return static_cast<float>(eval_host(l_var));))
    }
 
    [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
    double operator() (GpuArray<double,N> const& var) const noexcept
    {
        AMREX_IF_ON_DEVICE((return parser_exe_eval(m_device_executor, var.data());))
        AMREX_IF_ON_HOST((return eval_host(var.data());))
    }
 
    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
    explicit operator bool () const {
        AMREX_IF_ON_DEVICE((return m_device_executor != nullptr;))
        AMREX_IF_ON_HOST((return m_host_executor != nullptr;))
    }
 
    char* m_host_executor = nullptr;  
#ifdef AMREX_USE_GPU
    char* m_device_executor = nullptr; 
#endif
    bool m_threaded = false; 
 
private:
    AMREX_FORCE_INLINE
    double eval_host (double const* var) const noexcept
    {
#if AMREX_PARSER_THREADED_DISPATCH
        if (m_threaded) { return parser_exe_eval_threaded(m_host_executor, var); }
#endif
        return parser_exe_eval(m_host_executor, var);
    }
};
 
 
class Parser
{
public:
    Parser (std::string const& func_body);
    Parser () = default;
    void define (std::string const& func_body);
 
    explicit operator bool () const;
 
    void setConstant (std::string const& name, double c);
 
    void registerVariables (Vector<std::string> const& vars);
 
    void registerUserFn1 (std::string const& name, ParserUserFn1 fh, ParserUserFn1 fd);
    void registerUserFn2 (std::string const& name, ParserUserFn2 fh, ParserUserFn2 fd);
    void registerUserFn3 (std::string const& name, ParserUserFn3 fh, ParserUserFn3 fd);
    void registerUserFn4 (std::string const& name, ParserUserFn4 fh, ParserUserFn4 fd);
 
    void print () const;
    void printExe () const;
 
    [[nodiscard]] int depth () const;
    [[nodiscard]] int maxStackSize () const;
 
    [[nodiscard]] std::string expr () const;
 
    [[nodiscard]] std::set<std::string> symbols () const;
 
    [[nodiscard]] std::map<std::string,int> const& userFunctions () const;
 
    template <int N> [[nodiscard]] ParserExecutor<N> compile () const;
 
    template <int N> [[nodiscard]] ParserExecutor<N> compileHost () const;
 
private:
 
    template <int argc, typename F>
    void register_user_fn (std::string const& name, F fh, F fd);
 
    struct Data {
        std::string m_expression;
        struct amrex_parser* m_parser = nullptr;
        int m_nvars = 0;
        bool m_use_arena = true;
        char* m_host_executor = nullptr;
#ifdef AMREX_USE_GPU
        char* m_device_executor = nullptr;
#endif
        int m_max_stack_size = 0;
        int m_exe_size = 0;
        Vector<char const*> m_locals;
        Data () = default;
        ~Data ();
        void clear_host_executor ();
        Data (Data const&) = delete;
        Data (Data &&) = delete;
        Data& operator= (Data const&) = delete;
        Data& operator= (Data &&) = delete;
    };
 
    std::shared_ptr<Data> m_data;
    std::map<std::string,int> m_ufs;
    std::map<std::string,std::pair<void*,void*>> m_uf_ptrs;
    Vector<std::string> m_vars;
};
 
template <int N>
ParserExecutor<N>
Parser::compileHost () const
{
    if (m_data && m_data->m_parser) {
        AMREX_ALWAYS_ASSERT(N == m_data->m_nvars);
 
        // Make sure all user functions have been registered.
        for (auto const& [ufname, nargs] : m_ufs) {
            if (m_uf_ptrs.find(ufname) == m_uf_ptrs.cend()) {
                amrex::Abort("amrex::Parser: Unknown user function "+ufname
                             + " in " + m_data->m_expression);
            }
            amrex::ignore_unused(nargs);
        }
 
        if (!(m_data->m_host_executor)) {
            int stack_size;
            m_data->m_exe_size = static_cast<int>
                (parser_exe_size(m_data->m_parser, m_data->m_max_stack_size,
                                 stack_size));
 
            if (m_data->m_max_stack_size > AMREX_PARSER_STACK_SIZE) {
                amrex::Abort("amrex::Parser: AMREX_PARSER_STACK_SIZE, "
                             + std::to_string(AMREX_PARSER_STACK_SIZE) + ", is too small for "
                             + m_data->m_expression);
            }
            if (stack_size != 0) {
                amrex::Abort("amrex::Parser: something went wrong with parser stack! "
                             + std::to_string(stack_size));
            }
 
            m_data->m_host_executor = (char*)The_Pinned_Arena()->alloc(m_data->m_exe_size);
            if (m_data->m_host_executor == nullptr) { // Arena is not ready yet
                m_data->m_host_executor = (char*) std::malloc(m_data->m_exe_size);
                m_data->m_use_arena = false;
            }
 
            try {
                m_data->m_locals = parser_compile(m_data->m_parser, m_uf_ptrs,
                                                  m_data->m_host_executor);
            } catch (const std::runtime_error& e) {
                // Do not leave a partially compiled executor behind.
                m_data->clear_host_executor();
                throw std::runtime_error(std::string(e.what()) + " in Parser expression \""
                                         + m_data->m_expression + "\"");
            }
        }
 
        bool threaded = parser_exe_use_threaded(m_data->m_exe_size,
                                                m_data->m_max_stack_size);
#ifdef AMREX_USE_GPU
        return ParserExecutor<N>{m_data->m_host_executor, m_data->m_device_executor, threaded};
#else
        return ParserExecutor<N>{m_data->m_host_executor, threaded};
#endif
    } else {
        return ParserExecutor<N>{};
    }
}
 
template <int N>
ParserExecutor<N>
Parser::compile () const
{
    if (m_data && m_data->m_parser) {
        auto const unsupported = parser_get_device_unsupported_functions(m_data->m_parser);
        if (!unsupported.empty()) {
            std::string msg("amrex::Parser: these functions cannot be used in device code:");
            for (auto const& f : unsupported) { msg.append(" ").append(f); }
            amrex::Abort(msg + ". Use compileHost instead for " + m_data->m_expression);
        }
    }
 
    auto exe = compileHost<N>();
 
#ifdef AMREX_USE_GPU
    if (m_data && m_data->m_parser && !(m_data->m_device_executor)
        && m_data->m_use_arena)
    {
        m_data->m_device_executor = (char*)The_Arena()->alloc(m_data->m_exe_size);
        Gpu::htod_memcpy_async(m_data->m_device_executor, m_data->m_host_executor,
                               m_data->m_exe_size);
        Gpu::streamSynchronize();
        exe.m_device_executor = m_data->m_device_executor;
    }
#endif
 
    return exe;
}
 
template <int argc, typename F>
void
Parser::register_user_fn (std::string const& name, F fh, F fd)
{
#if !defined(AMREX_USE_GPU)
    // For pure CPU builds, we allow the user to pass the function pointer with either fh or fd
    if (fh == nullptr) { std::swap(fh, fd); }
#endif
 
    auto it = m_ufs.find(name);
    if (it != m_ufs.end()) {
        if (it->second == argc) {
            m_uf_ptrs[name] = std::make_pair((void*)fh, (void*)fd);
        } else {
            amrex::Abort("Parser::registerUserFn: incorrect no. of arguments in "+name);
        }
    } else {
        amrex::Abort("Parser::registerUserFn: unknown user define function "+name);
    }
}
 
}
 
#endif