NEML2ModelHandle.h

Go to the documentation of this file.

//* This file is part of the MOOSE framework
//* https://mooseframework.inl.gov
//*
//* 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 "NEML2Utils.h"
 
#ifdef NEML2_ENABLED
 
#include <memory>
#include <string>
#include <map>
#include <vector>
#include <utility>
#include <cstddef>
#include <filesystem>
 
#include "neml2/csrc/eager/Model.h"
#include "neml2/csrc/eager/load_model.h"
#include "neml2/csrc/dispatchers/factory.h"
#include "neml2/csrc/dispatchers/DispatchedModel.h"
#include "neml2/csrc/dispatchers/MPISimpleScheduler.h"
 
class NEML2ModelHandle
{
public:
  virtual ~NEML2ModelHandle() = default;
 
  virtual const std::vector<std::string> & input_names() const = 0;
  virtual const std::vector<std::string> & output_names() const = 0;
  virtual const std::vector<std::vector<int64_t>> & input_base_shapes() const = 0;
  virtual const std::vector<std::vector<int64_t>> & output_base_shapes() const = 0;
  virtual const std::map<std::string, std::vector<int64_t>> & parameter_base_shapes() const = 0;
 
  virtual std::map<std::string, at::Tensor>
  value(const std::map<std::string, at::Tensor> & inputs) const = 0;
  virtual std::pair<std::map<std::string, at::Tensor>, neml2::aoti::VariablePairJacobian>
  jacobian(const std::map<std::string, at::Tensor> & inputs) const = 0;
  virtual std::pair<std::map<std::string, at::Tensor>, neml2::aoti::VariablePairJacobian>
  param_jacobian(const std::map<std::string, at::Tensor> & inputs) const = 0;
  virtual void set_parameter(const std::string & name, const at::Tensor & value) = 0;
 
  virtual at::Device device() const = 0;
};
 
class EagerModelHandle : public NEML2ModelHandle
{
public:
  EagerModelHandle(const std::string & input_file,
                   const std::string & model_name,
                   const at::Device & device,
                   const std::vector<std::string> & load)
    : _m(input_file, model_name, device, load)
  {
  }
 
  const std::vector<std::string> & input_names() const override { return _m.input_names(); }
  const std::vector<std::string> & output_names() const override { return _m.output_names(); }
  const std::vector<std::vector<int64_t>> & input_base_shapes() const override
  {
    return _m.input_base_shapes();
  }
  const std::vector<std::vector<int64_t>> & output_base_shapes() const override
  {
    return _m.output_base_shapes();
  }
  const std::map<std::string, std::vector<int64_t>> & parameter_base_shapes() const override
  {
    return _m.parameter_base_shapes();
  }
 
  std::pair<std::map<std::string, at::Tensor>, neml2::aoti::VariablePairJacobian>
  jacobian(const std::map<std::string, at::Tensor> & inputs) const override
  {
    return _m.jacobian(inputs);
  }
  std::pair<std::map<std::string, at::Tensor>, neml2::aoti::VariablePairJacobian>
  param_jacobian(const std::map<std::string, at::Tensor> & inputs) const override
  {
    return _m.param_jacobian(inputs);
  }
  void set_parameter(const std::string & name, const at::Tensor & value) override
  {
    _m.set_parameter(name, value);
  }
  std::map<std::string, at::Tensor>
  value(const std::map<std::string, at::Tensor> & inputs) const override
  {
    return _m.forward(inputs);
  }
  at::Device device() const override { return _m.device(); }
 
private:
  neml2::eager::Model _m;
};
 
class AOTIModelHandle : public NEML2ModelHandle
{
public:
  AOTIModelHandle(const std::string & input,
                  const std::string & model_name,
                  const std::vector<std::string> & devices,
                  const std::vector<std::size_t> & batch_sizes,
                  const void * comm)
    : _m(load(input, model_name, devices, batch_sizes, comm))
  {
  }
 
  const std::vector<std::string> & input_names() const override { return _m.input_names(); }
  const std::vector<std::string> & output_names() const override { return _m.output_names(); }
  const std::vector<std::vector<int64_t>> & input_base_shapes() const override
  {
    return _m.input_base_shapes();
  }
  const std::vector<std::vector<int64_t>> & output_base_shapes() const override
  {
    return _m.output_base_shapes();
  }
  const std::map<std::string, std::vector<int64_t>> & parameter_base_shapes() const override
  {
    return _m.parameter_base_shapes();
  }
 
  std::pair<std::map<std::string, at::Tensor>, neml2::aoti::VariablePairJacobian>
  jacobian(const std::map<std::string, at::Tensor> & inputs) const override
  {
    return _m.jacobian(inputs);
  }
  std::pair<std::map<std::string, at::Tensor>, neml2::aoti::VariablePairJacobian>
  param_jacobian(const std::map<std::string, at::Tensor> & inputs) const override
  {
    return _m.param_jacobian(inputs);
  }
  void set_parameter(const std::string & name, const at::Tensor & value) override
  {
    _m.set_parameter(name, value);
  }
  std::map<std::string, at::Tensor>
  value(const std::map<std::string, at::Tensor> & inputs) const override
  {
    return _m.forward(inputs);
  }
  at::Device device() const override { return _m.device(); }
 
private:
  static neml2::aoti::DispatchedModel load(const std::string & input,
                                           const std::string & model_name,
                                           const std::vector<std::string> & devices,
                                           const std::vector<std::size_t> & batch_sizes,
                                           const void * comm)
  {
    neml2::aoti::MPISimpleScheduler::Config config;
    config.devices = devices;
    config.batch_sizes = batch_sizes;
    config.comm = comm;
    auto scheduler = std::make_shared<neml2::aoti::MPISimpleScheduler>(config);
 
    const std::filesystem::path path(input);
    if (std::filesystem::is_directory(path))
      return neml2::aoti::DispatchedModel(path, scheduler);
    if (path.filename() == "metadata.json")
      return neml2::aoti::DispatchedModel(path.parent_path(), scheduler);
    return neml2::aoti::load_model(input, model_name, scheduler);
  }
 
  neml2::aoti::DispatchedModel _m;
};
 
inline std::unique_ptr<NEML2ModelHandle>
makeNEML2ModelHandle(bool eager,
                     const std::string & input,
                     const std::string & model,
                     const std::vector<std::string> & devices,
                     const std::vector<std::size_t> & batch_sizes,
                     const void * comm,
                     const std::vector<std::string> & load)
{
  if (eager)
    return std::make_unique<EagerModelHandle>(input, model, at::Device(devices.at(0)), load);
  return std::make_unique<AOTIModelHandle>(input, model, devices, batch_sizes, comm);
}
 
#endif // NEML2_ENABLED