VariableStore.h

// Copyright 2023, UChicago Argonne, LLC
// All Rights Reserved
// Software Name: NEML2 -- the New Engineering material Model Library, version 2
// By: Argonne National Laboratory
// OPEN SOURCE LICENSE (MIT)
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
 
#pragma once
 
#include "neml2/base/NEML2Object.h"
#include "neml2/base/Storage.h"
#include "neml2/tensors/Variable.h"
#include "neml2/tensors/LabeledVector.h"
#include "neml2/tensors/LabeledMatrix.h"
#include "neml2/tensors/LabeledTensor3D.h"
 
namespace neml2
{
class VariableStore
{
public:
  VariableStore(const OptionSet & options, NEML2Object * object);
 
  LabeledAxis & declare_axis(const std::string & name);
 
  virtual void setup_layout();
 
  template <typename T = BatchTensor>
  Variable<T> & get_input_variable(const VariableName & name)
  {
    auto var_base_ptr = _input_views.query_value(name);
    neml_assert(var_base_ptr, "Input variable ", name, " does not exist.");
    auto var_ptr = dynamic_cast<Variable<T> *>(var_base_ptr);
    neml_assert(
        var_ptr, "Input variable ", name, " exist but cannot be cast to the requested type.");
    return *var_ptr;
  }
  template <typename T = BatchTensor>
  const Variable<T> & get_input_variable(const VariableName & name) const
  {
    const auto var_base_ptr = _input_views.query_value(name);
    neml_assert(var_base_ptr, "Input variable ", name, " does not exist.");
    const auto var_ptr = dynamic_cast<const Variable<T> *>(var_base_ptr);
    neml_assert(
        var_ptr, "Input variable ", name, " exist but cannot be cast to the requested type.");
    return *var_ptr;
  }
 
  template <typename T = BatchTensor>
  const Variable<T> & get_output_variable(const VariableName & name)
  {
    return std::as_const(*this).get_output_variable<T>(name);
  }
  template <typename T = BatchTensor>
  const Variable<T> & get_output_variable(const VariableName & name) const
  {
    const auto var_base_ptr = _output_views.query_value(name);
    neml_assert(var_base_ptr, "Output variable ", name, " does not exist.");
    const auto var_ptr = dynamic_cast<const Variable<T> *>(var_base_ptr);
    neml_assert(
        var_ptr, "Output variable ", name, " exist but cannot be cast to the requested type.");
    return *var_ptr;
  }
 
  LabeledAxis & input_axis() { return _input_axis; }
  const LabeledAxis & input_axis() const { return _input_axis; }
 
  LabeledAxis & output_axis() { return _output_axis; }
  const LabeledAxis & output_axis() const { return _output_axis; }
 
  Storage<VariableName, VariableBase> & input_views() { return _input_views; }
  const Storage<VariableName, VariableBase> & input_views() const { return _input_views; }
 
  Storage<VariableName, VariableBase> & output_views() { return _output_views; }
  const Storage<VariableName, VariableBase> & output_views() const { return _output_views; }
 
  LabeledVector & input_storage() { return _in; }
  const LabeledVector & input_storage() const { return _in; }
 
  LabeledVector & output_storage() { return _out; }
  const LabeledVector & output_storage() const { return _out; }
 
  LabeledMatrix & derivative_storage() { return _dout_din; }
  const LabeledMatrix & derivative_storage() const { return _dout_din; }
 
  LabeledTensor3D & second_derivative_storage() { return _d2out_din2; }
  const LabeledTensor3D & second_derivative_storage() const { return _d2out_din2; }
 
  VariableBase * input_view(const VariableName &);
  VariableBase * output_view(const VariableName &);
 
protected:
  virtual void cache(TorchShapeRef batch_shape);
 
  virtual void allocate_variables(TorchShapeRef batch_shape,
                                  const torch::TensorOptions & options,
                                  bool in,
                                  bool out,
                                  bool dout_din,
                                  bool d2out_din2);
 
  virtual void setup_input_views();
 
  virtual void setup_output_views();
 
  virtual void reinit_input_views();
 
  virtual void reinit_output_views(bool out, bool dout_din = true, bool d2out_din2 = true);
 
  virtual void detach_and_zero(bool out, bool dout_din = true, bool d2out_din2 = true);
 
  template <typename T, typename... S>
  const Variable<T> & declare_input_variable(S &&... name)
  {
    const auto var_name = variable_name(std::forward<S>(name)...);
    declare_variable<T>(_input_axis, var_name);
    return *create_variable_view<T>(_input_views, var_name);
  }
 
  template <typename... S>
  const Variable<BatchTensor> & declare_input_variable(TorchSize sz, S &&... name)
  {
    const auto var_name = variable_name(std::forward<S>(name)...);
    declare_variable(_input_axis, var_name, sz);
    return *create_variable_view<BatchTensor>(_input_views, var_name, sz);
  }
 
  template <typename T, typename... S>
  const Variable<BatchTensor> & declare_input_variable_list(TorchSize list_size, S &&... name)
  {
    return declare_input_variable(list_size * T::const_base_storage, std::forward<S>(name)...);
  }
 
  template <typename T, typename... S>
  Variable<T> & declare_output_variable(S &&... name)
  {
    const auto var_name = variable_name(std::forward<S>(name)...);
    declare_variable<T>(_output_axis, var_name);
    return *create_variable_view<T>(_output_views, var_name);
  }
 
  template <typename... S>
  Variable<BatchTensor> & declare_output_variable(TorchSize sz, S &&... name)
  {
    const auto var_name = variable_name(std::forward<S>(name)...);
    declare_variable(_output_axis, var_name, sz);
    return *create_variable_view<BatchTensor>(_output_views, var_name, sz);
  }
 
  template <typename T, typename... S>
  Variable<BatchTensor> & declare_output_variable_list(TorchSize list_size, S &&... name)
  {
    return declare_output_variable(list_size * T::const_base_storage, std::forward<S>(name)...);
  }
 
  template <typename T>
  VariableName declare_variable(LabeledAxis & axis, const VariableName & var) const
  {
    return declare_variable(axis, var, T::const_base_storage);
  }
 
  VariableName declare_variable(LabeledAxis & axis, const VariableName & var, TorchSize sz) const
  {
    axis.add(var, sz);
    return var;
  }
 
  VariableName declare_subaxis(LabeledAxis & axis, const VariableName & subaxis) const
  {
    axis.add<LabeledAxis>(subaxis);
    return subaxis;
  }
 
private:
  // Helper method to construct variable name in place
  template <typename... S>
  VariableName variable_name(S &&... name) const
  {
    using FirstType = std::tuple_element_t<0, std::tuple<S...>>;
 
    if constexpr (sizeof...(name) == 1 && std::is_convertible_v<FirstType, std::string>)
    {
      if (_options.contains<VariableName>(name...))
        return _options.get<VariableName>(name...);
      return VariableName(std::forward<S>(name)...);
    }
    else
      return VariableName(std::forward<S>(name)...);
  }
 
  // Create a variable view (doesn't setup the view)
  template <typename T>
  Variable<T> * create_variable_view(Storage<VariableName, VariableBase> & views,
                                     const VariableName & name,
                                     TorchSize sz = -1)
  {
    if constexpr (std::is_same_v<T, BatchTensor>)
      neml_assert(sz > 0, "Allocating a BatchTensor requires a known storage size.");
 
    // Make sure we don't duplicate variable allocation
    VariableBase * var_base_ptr = views.query_value(name);
    neml_assert(!var_base_ptr,
                "Trying to allocate variable ",
                name,
                ", but a variable with the same name already exists.");
 
    // Allocate
    if constexpr (std::is_same_v<T, BatchTensor>)
    {
      auto var = std::make_unique<Variable<BatchTensor>>(name, sz);
      var_base_ptr = views.set_pointer(name, std::move(var));
    }
    else
    {
      auto var = std::make_unique<Variable<T>>(name);
      var_base_ptr = views.set_pointer(name, std::move(var));
    }
 
    // Cast it to the concrete type
    auto var_ptr = dynamic_cast<Variable<T> *>(var_base_ptr);
    neml_assert(
        var_ptr, "Internal error: Failed to cast variable ", name, " to its concrete type.");
 
    return var_ptr;
  }
 
  NEML2Object * _object;
 
  const OptionSet _options;
 
  Storage<std::string, LabeledAxis> _axes;
 
  Storage<VariableName, VariableBase> _input_views;
 
  Storage<VariableName, VariableBase> _output_views;
 
  LabeledAxis & _input_axis;
 
  LabeledAxis & _output_axis;
 
  LabeledVector _in;
 
  LabeledVector _out;
 
  LabeledMatrix _dout_din;
 
  LabeledTensor3D _d2out_din2;
};
} // namespace neml2