expression.hpp

// Copyright (c) Sleipnir contributors
 
#pragma once
 
#include <stdint.h>
 
#include <algorithm>
#include <array>
#include <concepts>
#include <cstddef>
#include <memory>
#include <string_view>
#include <tuple>
#include <type_traits>
#include <utility>
 
#include <gch/small_vector.hpp>
 
#include "sleipnir/autodiff/expression_type.hpp"
#include "sleipnir/util/force_inline.hpp"
#include "sleipnir/util/function_ref.hpp"
#include "sleipnir/util/intrusive_shared_ptr.hpp"
#include "sleipnir/util/pool.hpp"
 
namespace slp::detail {
 
#ifdef _WIN32
inline constexpr bool USE_POOL_ALLOCATOR = false;
#else
inline constexpr bool USE_POOL_ALLOCATOR = true;
#endif
 
template <typename Scalar>
struct Expression;
 
inline constexpr size_t MAX_ARGS = 4;
 
template <typename T>
using ArgArray = std::array<T, MAX_ARGS>;
 
template <typename Scalar>
using ExpressionPtr = IntrusiveSharedPtr<Expression<Scalar>>;
 
template <typename T, typename... Args>
static ExpressionPtr<typename T::Scalar> make_expression_ptr(Args&&... args) {
  if constexpr (USE_POOL_ALLOCATOR) {
    return allocate_intrusive_shared<T>(global_pool_allocator<T>(),
                                        std::forward<Args>(args)...);
  } else {
    return make_intrusive_shared<T>(std::forward<Args>(args)...);
  }
}
 
template <typename Scalar, ExpressionType T>
struct BinaryMinusExpression;
 
template <typename Scalar, ExpressionType T>
struct BinaryPlusExpression;
 
template <typename Scalar, ExpressionType T>
struct DivExpression;
 
template <typename Scalar, ExpressionType T>
struct MultExpression;
 
template <typename Scalar, ExpressionType T>
struct UnaryMinusExpression;
 
template <template <typename, ExpressionType> typename T, typename Scalar,
          typename... Args>
SLEIPNIR_FORCE_INLINE ExpressionPtr<Scalar> make_typed_expression_ptr(
    ExpressionType type, const Args&... args) {
  using enum ExpressionType;
 
  if (type == LINEAR) {
    return make_expression_ptr<T<Scalar, LINEAR>>(args...);
  } else if (type == QUADRATIC) {
    return make_expression_ptr<T<Scalar, QUADRATIC>>(args...);
  } else {
    return make_expression_ptr<T<Scalar, NONLINEAR>>(args...);
  }
}
 
template <typename Scalar>
ExpressionPtr<Scalar> constant_ptr(Scalar value);
 
template <typename Value, typename Scalar>
Value lit(Scalar value) {
  if constexpr (std::same_as<Value, Scalar>) {
    return value;
  } else {
    return constant_ptr(value);
  }
}
 
template <typename Scalar_>
struct Expression {
  using Scalar = Scalar_;
 
  Scalar val{0};
 
  Scalar adj{0};
 
  ExpressionPtr<Scalar> adj_expr;
 
  int32_t scratch = -1;
 
  uint32_t ref_count : 31 = 0;
 
  uint32_t is_leaf : 1 = 0;
 
  constexpr Expression() = default;
 
  explicit constexpr Expression(Scalar value) : val{value}, is_leaf{1} {}
 
  virtual ~Expression() = default;
 
  SLEIPNIR_FORCE_INLINE constexpr bool is_constant(Scalar constant) const {
    return type() == ExpressionType::CONSTANT && val == constant;
  }
 
  friend ExpressionPtr<Scalar> operator*(const ExpressionPtr<Scalar>& lhs,
                                         const ExpressionPtr<Scalar>& rhs) {
    using enum ExpressionType;
 
    // Prune expression
    if (lhs->is_constant(Scalar(0))) {
      // Return zero, which lhs currently is
      return lhs;
    } else if (rhs->is_constant(Scalar(0))) {
      // Return zero, which rhs currently is
      return rhs;
    } else if (lhs->is_constant(Scalar(1))) {
      // Return rhs unmodified
      return rhs;
    } else if (rhs->is_constant(Scalar(1))) {
      // Return lhs unmodified
      return lhs;
    }
 
    // Evaluate constant
    if (lhs->type() == CONSTANT && rhs->type() == CONSTANT) {
      return constant_ptr(lhs->val * rhs->val);
    }
 
    // Evaluate expression type. Scaling by a constant keeps the other
    // operand's type.
    auto type = NONLINEAR;
    if (lhs->type() == CONSTANT || rhs->type() == CONSTANT) {
      type = std::max(lhs->type(), rhs->type());
    } else if (lhs->type() == LINEAR && rhs->type() == LINEAR) {
      type = QUADRATIC;
    }
    return make_typed_expression_ptr<MultExpression, Scalar>(type, lhs, rhs);
  }
 
  friend ExpressionPtr<Scalar> operator/(const ExpressionPtr<Scalar>& lhs,
                                         const ExpressionPtr<Scalar>& rhs) {
    using enum ExpressionType;
 
    // Prune expression
    if (lhs->is_constant(Scalar(0))) {
      // Return zero, which lhs currently is
      return lhs;
    } else if (rhs->is_constant(Scalar(1))) {
      // Return lhs unmodified
      return lhs;
    }
 
    // Evaluate constant
    if (lhs->type() == CONSTANT && rhs->type() == CONSTANT) {
      return constant_ptr(lhs->val / rhs->val);
    }
 
    // Evaluate expression type. Dividing by a constant keeps the numerator's
    // type.
    auto type = rhs->type() == CONSTANT ? lhs->type() : NONLINEAR;
    return make_typed_expression_ptr<DivExpression, Scalar>(type, lhs, rhs);
  }
 
  friend ExpressionPtr<Scalar> operator+(const ExpressionPtr<Scalar>& lhs,
                                         const ExpressionPtr<Scalar>& rhs) {
    using enum ExpressionType;
 
    // Prune expression. We check for nullptr because operator+ is used in
    // adjoint accumulation, and child nodes can be null.
    if (lhs == nullptr || lhs->is_constant(Scalar(0))) {
      // Return rhs unmodified
      return rhs;
    } else if (rhs == nullptr || rhs->is_constant(Scalar(0))) {
      // Return lhs unmodified
      return lhs;
    }
 
    // Evaluate constant
    if (lhs->type() == CONSTANT && rhs->type() == CONSTANT) {
      return constant_ptr(lhs->val + rhs->val);
    }
 
    auto type = std::max(lhs->type(), rhs->type());
    return make_typed_expression_ptr<BinaryPlusExpression, Scalar>(type, lhs,
                                                                   rhs);
  }
 
  friend ExpressionPtr<Scalar> operator+=(ExpressionPtr<Scalar>& lhs,
                                          const ExpressionPtr<Scalar>& rhs) {
    return lhs = lhs + rhs;
  }
 
  friend ExpressionPtr<Scalar> operator-(const ExpressionPtr<Scalar>& lhs,
                                         const ExpressionPtr<Scalar>& rhs) {
    using enum ExpressionType;
 
    // Prune expression
    if (lhs->is_constant(Scalar(0))) {
      if (rhs->is_constant(Scalar(0))) {
        // Return zero, which rhs currently is
        return rhs;
      } else {
        // Return rhs negated
        return -rhs;
      }
    } else if (rhs->is_constant(Scalar(0))) {
      // Return lhs unmodified
      return lhs;
    }
 
    // Evaluate constant
    if (lhs->type() == CONSTANT && rhs->type() == CONSTANT) {
      return constant_ptr(lhs->val - rhs->val);
    }
 
    auto type = std::max(lhs->type(), rhs->type());
    return make_typed_expression_ptr<BinaryMinusExpression, Scalar>(type, lhs,
                                                                    rhs);
  }
 
  friend ExpressionPtr<Scalar> operator-(const ExpressionPtr<Scalar>& lhs) {
    using enum ExpressionType;
 
    // Prune expression
    if (lhs->is_constant(Scalar(0))) {
      // Return zero, which lhs currently is
      return lhs;
    }
 
    // Evaluate constant
    if (lhs->type() == CONSTANT) {
      return constant_ptr(-lhs->val);
    }
 
    return make_typed_expression_ptr<UnaryMinusExpression, Scalar>(lhs->type(),
                                                                   lhs);
  }
 
  friend ExpressionPtr<Scalar> operator+(const ExpressionPtr<Scalar>& lhs) {
    return lhs;
  }
 
  virtual void visit_args(
      [[maybe_unused]] function_ref<void(Expression<Scalar>* arg)> func) const {
  }
 
  virtual Scalar value() const = 0;
 
  virtual ExpressionType type() const = 0;
 
  virtual std::string_view name() const = 0;
 
  virtual void accumulate_adjoints() const {}
 
  virtual void accumulate_adjoints_expr() const {}
};
 
template <typename Derived, typename Scalar, size_t NumArgs,
          ExpressionType Type = ExpressionType::NONLINEAR>
struct NonLeafExpression : Expression<Scalar> {
  static constexpr size_t NUM_ARGS = NumArgs;
 
  // Ensure ArgArray is big enough to hold one element per argument
  static_assert(NumArgs <= MAX_ARGS, "Increase MAX_ARGS");
 
  static constexpr bool CONSTANT_PARTIALS = false;
 
  std::array<ExpressionPtr<Scalar>, NumArgs> args;
 
  template <typename... Args>
    requires(sizeof...(Args) == NumArgs &&
             (std::same_as<std::remove_cvref_t<Args>, ExpressionPtr<Scalar>> &&
              ...))
  explicit constexpr NonLeafExpression(Args&&... args)
      : args{std::forward<Args>(args)...} {}
 
  void visit_args(
      function_ref<void(Expression<Scalar>* arg)> func) const final {
    [&]<size_t... Is>(std::index_sequence<Is...>) {
      (func(args[Is].get()), ...);
    }(std::make_index_sequence<NumArgs>{});
  }
 
  Scalar value() const final {
    return [&]<size_t... Is>(std::index_sequence<Is...>) {
      return static_cast<const Derived&>(*this).eval(args[Is]->val...);
    }(std::make_index_sequence<NumArgs>{});
  }
 
  ExpressionType type() const final { return Type; }
 
  std::string_view name() const final { return Derived::NAME; }
 
  void accumulate_adjoints() const final {
    [&]<size_t... Is>(std::index_sequence<Is...>) {
      // Evaluate gradients
      ArgArray<Scalar> g{};
      grad(g, this->val, args[Is]->val...);
 
      // Accumulate adjoints
      ((args[Is]->adj += this->adj * g[Is]), ...);
    }(std::make_index_sequence<NumArgs>{});
  }
 
  void accumulate_adjoints_expr() const final {
    [&]<size_t... Is>(std::index_sequence<Is...>) {
      // If the partial derivatives are constant, evaluate them as Scalars
      // instead of building them as expressions
      if constexpr (Derived::CONSTANT_PARTIALS) {
        // Evaluate gradients
        ArgArray<Scalar> g{};
        grad(g, this->val, args[Is]->val...);
 
        // Accumulate adjoints
        (
            [&] {
              // Fold the adjoint accumulation expression if the gradient is
              // 1, -1, or 0
              if (g[Is] == Scalar(1)) {
                args[Is]->adj_expr += this->adj_expr;
              } else if (g[Is] == Scalar(-1)) {
                args[Is]->adj_expr += -this->adj_expr;
              } else if (g[Is] != Scalar(0)) {
                args[Is]->adj_expr += this->adj_expr * constant_ptr(g[Is]);
              }
            }(),
            ...);
      } else {
        // Evaluate gradients
        ArgArray<ExpressionPtr<Scalar>> g;
        grad(g, ExpressionPtr<Scalar>{const_cast<NonLeafExpression*>(this)},
             args[Is]...);
 
        // Accumulate adjoints
        (
            [&] {
              // A null partial derivative is structurally zero
              if (g[Is] != nullptr) {
                args[Is]->adj_expr += this->adj_expr * g[Is];
              }
            }(),
            ...);
      }
    }(std::make_index_sequence<NumArgs>{});
  }
 
 private:
  template <typename Value, typename... Args>
  void grad(ArgArray<Value>& g, const Value& val, const Args&... x) const {
    if constexpr (requires(const Derived& expr) {
                    expr.eval_grad(g, val, x...);
                  }) {
      static_cast<const Derived&>(*this).eval_grad(g, val, x...);
    }
  }
};
 
template <typename T, typename... Args>
ExpressionPtr<typename T::Scalar> fold_or_make(const Args&... args) {
  using Scalar = typename T::Scalar;
 
  if (((args->type() == ExpressionType::CONSTANT) && ...)) {
    Scalar val = T::eval(args->val...);
    if constexpr (sizeof...(Args) == 1) {
      const auto& x = std::get<0>(std::forward_as_tuple(args...));
      if (x->val == val) {
        return x;
      }
    }
    return constant_ptr(val);
  }
 
  return make_expression_ptr<T>(args...);
}
 
template <typename Scalar, ExpressionType T>
struct BinaryMinusExpression final
    : NonLeafExpression<BinaryMinusExpression<Scalar, T>, Scalar, 2, T> {
  using BinaryMinusExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "binary minus";
 
  static constexpr bool CONSTANT_PARTIALS = true;
 
  static Scalar eval(Scalar lhs, Scalar rhs) { return lhs - rhs; }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        [[maybe_unused]] const Value& lhs,
                        [[maybe_unused]] const Value& rhs) {
    g[0] = lit<Value>(Scalar(1));
    g[1] = lit<Value>(Scalar(-1));
  }
};
 
template <typename Scalar, ExpressionType T>
struct BinaryPlusExpression final
    : NonLeafExpression<BinaryPlusExpression<Scalar, T>, Scalar, 2, T> {
  using BinaryPlusExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "binary plus";
 
  static constexpr bool CONSTANT_PARTIALS = true;
 
  static Scalar eval(Scalar lhs, Scalar rhs) { return lhs + rhs; }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        [[maybe_unused]] const Value& lhs,
                        [[maybe_unused]] const Value& rhs) {
    g[0] = lit<Value>(Scalar(1));
    g[1] = lit<Value>(Scalar(1));
  }
};
 
template <typename Scalar, ExpressionType T>
struct DivExpression final
    : NonLeafExpression<DivExpression<Scalar, T>, Scalar, 2, T> {
  using DivExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "division";
 
  static Scalar eval(Scalar lhs, Scalar rhs) { return lhs / rhs; }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        const Value& lhs, const Value& rhs) {
    g[0] = lit<Value>(Scalar(1)) / rhs;
    g[1] = -lhs / (rhs * rhs);
  }
};
 
template <typename Scalar, ExpressionType T>
struct MultExpression final
    : NonLeafExpression<MultExpression<Scalar, T>, Scalar, 2, T> {
  using MultExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "multiplication";
 
  static Scalar eval(Scalar lhs, Scalar rhs) { return lhs * rhs; }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        const Value& lhs, const Value& rhs) {
    g[0] = rhs;
    g[1] = lhs;
  }
};
 
template <typename Scalar, ExpressionType T>
struct UnaryMinusExpression final
    : NonLeafExpression<UnaryMinusExpression<Scalar, T>, Scalar, 1, T> {
  using UnaryMinusExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "unary minus";
 
  static constexpr bool CONSTANT_PARTIALS = true;
 
  static Scalar eval(Scalar lhs) { return -lhs; }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        [[maybe_unused]] const Value& lhs) {
    g[0] = lit<Value>(Scalar(-1));
  }
};
 
template <typename Scalar>
constexpr void inc_ref_count(Expression<Scalar>* expr) {
  ++expr->ref_count;
}
 
template <typename Scalar>
void deallocate_expression(Expression<Scalar>* expr) {
  // If a deeply nested tree is being deallocated all at once, calling the
  // Expression destructor when expr's refcount reaches zero can cause a stack
  // overflow. Instead, we iterate over its children to decrement their
  // refcounts and deallocate them.
  gch::small_vector<Expression<Scalar>*> stack;
  stack.emplace_back(expr);
 
  // Decrements a child's refcount and enqueues it for deallocation if it
  // reaches zero
  auto release = [&stack](Expression<Scalar>* child) {
    if (--child->ref_count == 0) {
      stack.emplace_back(child);
    }
  };
 
  while (!stack.empty()) {
    auto elem = stack.back();
    stack.pop_back();
 
    if (elem->adj_expr != nullptr) {
      release(elem->adj_expr.get());
    }
    elem->visit_args(release);
 
    // Not calling the destructor here is safe because it only decrements
    // refcounts, which was already done above
    if constexpr (USE_POOL_ALLOCATOR) {
      auto alloc = global_pool_allocator<Expression<Scalar>>();
      std::allocator_traits<decltype(alloc)>::deallocate(alloc, elem, 1);
    } else {
      operator delete(elem);
    }
  }
}
 
template <typename Scalar>
constexpr void dec_ref_count(Expression<Scalar>* expr) {
  if (--expr->ref_count == 0) {
    deallocate_expression(expr);
  }
}
 
template <typename Scalar>
struct ConstantExpression final : Expression<Scalar> {
  explicit constexpr ConstantExpression(Scalar value)
      : Expression<Scalar>{value} {}
 
  Scalar value() const override { return this->val; }
 
  ExpressionType type() const override { return ExpressionType::CONSTANT; }
 
  std::string_view name() const override { return "constant"; }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> constant_ptr(Scalar value) {
  return make_expression_ptr<ConstantExpression<Scalar>>(value);
}
 
template <typename Scalar>
struct DecisionVariableExpression final : Expression<Scalar> {
  constexpr DecisionVariableExpression() : Expression<Scalar>{Scalar(0)} {}
 
  explicit constexpr DecisionVariableExpression(Scalar value)
      : Expression<Scalar>{value} {}
 
  Scalar value() const override { return this->val; }
 
  ExpressionType type() const override { return ExpressionType::LINEAR; }
 
  std::string_view name() const override { return "decision variable"; }
};
 
}  // namespace slp::detail