KokkosFunctionParser.h

Go to the documentation of this file.

//* This file is part of the MOOSE framework
//* https://www.mooseframework.org
//*
//* All rights reserved, see COPYRIGHT for full restrictions
//* https://github.com/idaholab/moose/blob/master/COPYRIGHT
//*
//* Licensed under LGPL 2.1, please see LICENSE for details
//* https://www.gnu.org/licenses/lgpl-2.1.html

#pragma once

#include "KokkosTypes.h"
#include "KokkosVariableValue.h"
#include "KokkosMaterialProperty.h"
#include "KokkosMaterialPropertyValue.h"
#include "KokkosFunction.h"

#include "ConsoleStreamInterface.h"

#include "peglib.h"

namespace Moose::Kokkos
{

class FunctionBase;

class PEGParser
{
public:
  PEGParser(const std::string & expression, const ConsoleStream * console = nullptr);

  auto ast() const { return _ast; }
  const std::string & expression() const { return _expression; }

private:
  peg::parser _parser;
  std::shared_ptr<peg::Ast> _ast;
  const std::string _expression;
};

class RPNBuilder
{
public:
  RPNBuilder(const std::string & expression, const ConsoleStream * console = nullptr)
    : _parser(expression, console)
  {
  }

  enum class Opcode
  {
    NUM,
    VAR,
    NEG,
    NOT,
    // Binary operators
    ADD,
    SUB,
    MUL,
    DIV,
    AND,
    OR,
    EQ,
    NEQ,
    LT,
    LEQ,
    GT,
    GEQ,
    // Functions
    ABS,
    ACOS,
    ACOSH,
    ASIN,
    ASINH,
    ATAN,
    ATAN2,
    ATANH,
    CBRT,
    CEIL,
    COS,
    COSH,
    COT,
    CSC,
    EXP,
    EXP2,
    FLOOR,
    HYPOT,
    IF,
    INT,
    LOG,
    LOG2,
    LOG10,
    MAX,
    MIN,
    POW,
    SEC,
    SIN,
    SINH,
    SQRT,
    TAN,
    TANH,
    TRUNC
  };

  struct Instruction
  {
    // Opcode
    Opcode op;
    // Original string
    std::string text;
    // Constant of variable index
    unsigned int arg = libMesh::invalid_uint;
  };

  struct Variable
  {
    // Variable index
    unsigned int idx;
    // Pointer to a scalar value
    const Real * scalar = nullptr;
    // Pointer to a field variable
    const VariableValue * field = nullptr;
    // Pointer to a material property
    const MaterialProperty<Real> * property = nullptr;
    // Pointer to a function
    const Function * function = nullptr;
    // Whether the variable was associated
    bool associated() const { return scalar || field || property || function; }
  };

  void build(const peg::Ast & ast);
  void build() { build(*_parser.ast()); }

  void printRPN(const ConsoleStream & console) const;
  const std::vector<Instruction> & getRPN() const { return _rpn; }

  void addDefaultVariables();
  bool hasDefaultVariables() const { return _has_default_variables; }

  unsigned int addNumber(Real number);
  unsigned int addVariable(const std::string & name);
  void associateScalar(const std::string & name, const Real * scalar);
  void associateField(const std::string & name, const VariableValue * field);
  void associateProperty(const std::string & name, const MaterialProperty<Real> * property);
  void associateFunction(const std::string & name, const Function * function);

  const std::vector<Real> & getNumbers() const { return _numbers; }
  const std::unordered_map<std::string, Variable> & getVariables() const { return _variables; }

  void finalize() { _finalized = true; }
  bool finalized() const { return _finalized; }

private:
  static inline const std::map<std::string, Opcode> _unary_opcode_map = {{"-", Opcode::NEG},
                                                                         {"!", Opcode::NOT}};

  static inline const std::map<std::string, Opcode> _binary_opcode_map = {{"+", Opcode::ADD},
                                                                          {"-", Opcode::SUB},
                                                                          {"*", Opcode::MUL},
                                                                          {"/", Opcode::DIV},
                                                                          {"^", Opcode::POW},
                                                                          {"&", Opcode::AND},
                                                                          {"|", Opcode::OR},
                                                                          {"=", Opcode::EQ},
                                                                          {"!=", Opcode::NEQ},
                                                                          {"<", Opcode::LT},
                                                                          {"<=", Opcode::LEQ},
                                                                          {">", Opcode::GT},
                                                                          {">=", Opcode::GEQ}};

  static inline const std::map<std::string, std::pair<Opcode, unsigned int>> _function_opcode_map =
      {{"abs", {Opcode::ABS, 1}},     {"acos", {Opcode::ACOS, 1}},   {"acosh", {Opcode::ACOSH, 1}},
       {"asin", {Opcode::ASIN, 1}},   {"asinh", {Opcode::ASINH, 1}}, {"atan", {Opcode::ATAN, 1}},
       {"atan2", {Opcode::ATAN2, 2}}, {"atanh", {Opcode::ATANH, 1}}, {"cbrt", {Opcode::CBRT, 1}},
       {"ceil", {Opcode::CEIL, 1}},   {"cos", {Opcode::COS, 1}},     {"cosh", {Opcode::COSH, 1}},
       {"cot", {Opcode::COT, 1}},     {"csc", {Opcode::CSC, 1}},     {"exp", {Opcode::EXP, 1}},
       {"exp2", {Opcode::EXP2, 1}},   {"floor", {Opcode::FLOOR, 1}}, {"hypot", {Opcode::HYPOT, 2}},
       {"if", {Opcode::IF, 3}},       {"int", {Opcode::INT, 1}},     {"log", {Opcode::LOG, 1}},
       {"log2", {Opcode::LOG2, 1}},   {"log10", {Opcode::LOG10, 1}}, {"max", {Opcode::MAX, 2}},
       {"min", {Opcode::MIN, 2}},     {"pow", {Opcode::POW, 2}},     {"sec", {Opcode::SEC, 1}},
       {"sin", {Opcode::SIN, 1}},     {"sinh", {Opcode::SINH, 1}},   {"sqrt", {Opcode::SQRT, 1}},
       {"tan", {Opcode::TAN, 1}},     {"tanh", {Opcode::TANH, 1}},   {"trunc", {Opcode::TRUNC, 1}}};

  PEGParser _parser;
  std::vector<Instruction> _rpn;
  std::vector<Real> _numbers;
  std::unordered_map<std::string, Variable> _variables;
  bool _has_default_variables = false;
  bool _finalized = false;

  [[noreturn]] void builderError(const peg::Ast & ast, const std::string & message) const;

  void checkFinalized();
};

class RPNEvaluator
{
  using Opcode = RPNBuilder::Opcode;

public:
  RPNEvaluator() = default;
  RPNEvaluator(const RPNEvaluator & evaluator);

  void init(const RPNBuilder & builder);

  KOKKOS_FUNCTION Real eval(const Real t,
                            const Real3 p,
                            const unsigned int qp = 0,
                            Datum * datum = nullptr) const;

private:
  struct Instruction
  {
    // Opcode
    Opcode op;
    // Constant of variable index
    unsigned int arg = libMesh::invalid_uint;
  };
  enum class VariableType
  {
    SCALAR,
    FIELD,
    MATERIAL,
    FUNCTION
  };

  Array<Instruction> _rpn;
  Array<Real> _numbers;
  Array<Real> _scalars;
  Array<VariableValue> _fields;
  Array<MaterialProperty<Real>> _properties;
  Array<Function> _functions;
  Array<::Kokkos::pair<VariableType, unsigned int>> _variables;
  std::vector<const void *> _pointers;

  unsigned int _num_scalars = 0;
  unsigned int _num_fields = 0;
  unsigned int _num_properties = 0;
  unsigned int _num_functions = 0;

  static constexpr unsigned int _stack_size = 10;
  static constexpr double _epsilon = 1.0e-12;
};

#define KOKKOS_FPARSER_COMPARE(code, op, epsilon)                                                  \
  case Opcode::code:                                                                               \
    stack[head - 2] = stack[head - 2] - stack[head - 1] op epsilon;                                \
    --head;                                                                                        \
    break

#define KOKKOS_FPARSER_BINARY(code, op)                                                            \
  case Opcode::code:                                                                               \
    stack[head - 2] = stack[head - 2] op stack[head - 1];                                          \
    --head;                                                                                        \
    break

#define KOKKOS_FPARSER_FUNCTION_1(code, func)                                                      \
  case Opcode::code:                                                                               \
    stack[head - 1] = ::Kokkos::func(stack[head - 1]);                                             \
    break

#define KOKKOS_FPARSER_FUNCTION_2(code, func)                                                      \
  case Opcode::code:                                                                               \
    stack[head - 2] = ::Kokkos::func(stack[head - 2], stack[head - 1]);                            \
    --head;                                                                                        \
    break

#define KOKKOS_FPARSER_FUNCTION_INV_1(code, func)                                                  \
  case Opcode::code:                                                                               \
    stack[head - 1] = 1.0 / ::Kokkos::func(stack[head - 1]);                                       \
    break

KOKKOS_FUNCTION inline Real
RPNEvaluator::eval(const Real t, const Real3 p, const unsigned int qp, Datum * datum) const
{
  Real stack[_stack_size];

  // Stack head position
  unsigned int head = 0;

  for (unsigned int pos = 0; pos < _rpn.size(); ++pos)
  {
    KOKKOS_ASSERT(head < _stack_size);

    const auto inst = _rpn[pos];

    switch (_rpn[pos].op)
    {
      case Opcode::NUM:
        stack[head] = _numbers[inst.arg];
        ++head;
        break;
      case Opcode::VAR:
        if (inst.arg < 3)
          stack[head] = p(inst.arg);
        else if (inst.arg == 3)
          stack[head] = t;
        else
        {
          const auto type = _variables[inst.arg].first;
          const auto idx = _variables[inst.arg].second;

          switch (type)
          {
            case VariableType::SCALAR:
              stack[head] = _scalars[idx];
              break;
            case VariableType::FIELD:
            {
              KOKKOS_ASSERT(datum);

              stack[head] = _fields[idx](*datum, qp);
              break;
            }
            case VariableType::MATERIAL:
            {
              KOKKOS_ASSERT(datum);

              stack[head] = _properties[idx](*datum, qp);
              break;
            }
            case VariableType::FUNCTION:
              stack[head] = _functions[idx].value(t, p);
              break;
          }
        }
        ++head;
        break;
      case Opcode::NEG:
        stack[head - 1] = -stack[head - 1];
        break;
      case Opcode::NOT:
        stack[head - 1] = !::Kokkos::round(stack[head - 1]);
        break;
      case Opcode::EQ:
        stack[head - 2] = ::Kokkos::abs(stack[head - 2] - stack[head - 1]) <= _epsilon;
        --head;
        break;
      case Opcode::NEQ:
        stack[head - 2] = ::Kokkos::abs(stack[head - 2] - stack[head - 1]) > _epsilon;
        --head;
        break;
      case Opcode::AND:
        stack[head - 2] = ::Kokkos::round(stack[head - 2]) && ::Kokkos::round(stack[head - 1]);
        --head;
        break;
      case Opcode::OR:
        stack[head - 2] = ::Kokkos::round(stack[head - 2]) || ::Kokkos::round(stack[head - 1]);
        --head;
        break;
      case Opcode::IF:
        stack[head - 3] = ::Kokkos::round(stack[head - 3]) ? stack[head - 2] : stack[head - 1];
        head -= 2;
        break;
        KOKKOS_FPARSER_COMPARE(LT, <, -_epsilon);
        KOKKOS_FPARSER_COMPARE(LEQ, <=, _epsilon);
        KOKKOS_FPARSER_COMPARE(GT, >, _epsilon);
        KOKKOS_FPARSER_COMPARE(GEQ, >=, -_epsilon);
        KOKKOS_FPARSER_BINARY(ADD, +);
        KOKKOS_FPARSER_BINARY(SUB, -);
        KOKKOS_FPARSER_BINARY(MUL, *);
        KOKKOS_FPARSER_BINARY(DIV, /);
        KOKKOS_FPARSER_FUNCTION_1(ABS, abs);
        KOKKOS_FPARSER_FUNCTION_1(ACOS, acos);
        KOKKOS_FPARSER_FUNCTION_1(ACOSH, acosh);
        KOKKOS_FPARSER_FUNCTION_1(ASIN, asin);
        KOKKOS_FPARSER_FUNCTION_1(ASINH, asinh);
        KOKKOS_FPARSER_FUNCTION_1(ATAN, atan);
        KOKKOS_FPARSER_FUNCTION_2(ATAN2, atan2);
        KOKKOS_FPARSER_FUNCTION_1(ATANH, atanh);
        KOKKOS_FPARSER_FUNCTION_1(CBRT, cbrt);
        KOKKOS_FPARSER_FUNCTION_1(CEIL, ceil);
        KOKKOS_FPARSER_FUNCTION_1(COS, cos);
        KOKKOS_FPARSER_FUNCTION_1(COSH, cosh);
        KOKKOS_FPARSER_FUNCTION_INV_1(COT, tan);
        KOKKOS_FPARSER_FUNCTION_INV_1(CSC, sin);
        KOKKOS_FPARSER_FUNCTION_1(EXP, exp);
        KOKKOS_FPARSER_FUNCTION_1(EXP2, exp2);
        KOKKOS_FPARSER_FUNCTION_1(FLOOR, floor);
        KOKKOS_FPARSER_FUNCTION_2(HYPOT, hypot);
        KOKKOS_FPARSER_FUNCTION_1(INT, round);
        KOKKOS_FPARSER_FUNCTION_1(LOG, log);
        KOKKOS_FPARSER_FUNCTION_1(LOG2, log2);
        KOKKOS_FPARSER_FUNCTION_1(LOG10, log10);
        KOKKOS_FPARSER_FUNCTION_2(MAX, max);
        KOKKOS_FPARSER_FUNCTION_2(MIN, min);
        KOKKOS_FPARSER_FUNCTION_2(POW, pow);
        KOKKOS_FPARSER_FUNCTION_INV_1(SEC, cos);
        KOKKOS_FPARSER_FUNCTION_1(SIN, sin);
        KOKKOS_FPARSER_FUNCTION_1(SINH, sinh);
        KOKKOS_FPARSER_FUNCTION_1(SQRT, sqrt);
        KOKKOS_FPARSER_FUNCTION_1(TAN, tan);
        KOKKOS_FPARSER_FUNCTION_1(TANH, tanh);
        KOKKOS_FPARSER_FUNCTION_1(TRUNC, trunc);
    }
  }

  KOKKOS_ASSERT(head == 1);

  return stack[head - 1];
}

} // namespace Moose::Kokkos