NEML2ModelInterface< T > Class Template Reference

Interface class to provide common input parameters, members, and methods for MOOSEObjects that use NEML2 models. More...

#include <NEML2ModelInterface.h>

Inheritance diagram for NEML2ModelInterface< T >:

Public Member Functions
template<typename... P>
	NEML2ModelInterface (const InputParameters &params, P &&... args)

Static Public Member Functions
static InputParameters	validParams ()

Protected Member Functions
virtual void	validateModel () const
	Validate the NEML2 material model.
const NEML2ModelHandle &	model () const
	Get the NEML2 model handle (cpp-aoti by default, cpp-eager when eager=true)
NEML2ModelHandle &	model ()
	Non-const access to the NEML2 model handle, for runtime parameter writes (set_parameter)
const at::Device &	device () const
	Get the target compute device.
const at::Device &	output_device () const
	Get the target output device.
unsigned int	neml2NumThreads () const
	libtorch intra-op thread count to use while evaluating the NEML2 model.

Static Private Member Functions
static std::vector< std::string >	loadExtensionPaths (const InputParameters &params)
	Resolve the 'load' DataFileName list (whose relative paths InputParameters has already resolved against the input file + data search path) to plain path strings for the eager runtime's external-extension loader.
static std::unique_ptr< NEML2ModelHandle >	makeModelHandle (const InputParameters &params, NEML2ModelInterface< T > &self)
	Parse the runtime parameters and build the model handle.

Private Attributes
std::unique_ptr< NEML2ModelHandle >	_model
	The NEML2 material model handle.
const at::Device	_device
	The device on which to evaluate the NEML2 model (this rank's scheduler-assigned device)
const at::Device	_output_device
	The device on which to store the outputs.
const unsigned int	_num_threads
	libtorch intra-op thread count applied (scoped) during model evaluation; defaults to the MOOSE/libMesh thread count.

Detailed Description

template<class T>
class NEML2ModelInterface< T >

Interface class to provide common input parameters, members, and methods for MOOSEObjects that use NEML2 models.

Definition at line 28 of file NEML2ModelInterface.h.

Constructor & Destructor Documentation

NEML2ModelInterface()

template<class T >

template<typename... P>

NEML2ModelInterface< T >::NEML2ModelInterface	(	const InputParameters &	params,
		P &&...	args
	)

Definition at line 219 of file NEML2ModelInterface.h.

  : T(params, args...)
{
}

Member Function Documentation

device()

template<class T >

const at::Device & NEML2ModelInterface< T >::device ( ) const

inlineprotected

Get the target compute device.

Definition at line 52 of file NEML2ModelInterface.h.

{ return _device; }

loadExtensionPaths()

template<class T >

static std::vector< std::string > NEML2ModelInterface< T >::loadExtensionPaths ( const InputParameters &  params )

inlinestaticprivate

Resolve the 'load' DataFileName list (whose relative paths InputParameters has already resolved against the input file + data search path) to plain path strings for the eager runtime's external-extension loader.

Definition at line 66 of file NEML2ModelInterface.h.

  {
    const auto files = params.get<std::vector<DataFileName>>("load");
    return std::vector<std::string>(files.begin(), files.end());
  }

Referenced by NEML2ModelInterface< T >::makeModelHandle().

makeModelHandle()

template<class T >

static std::unique_ptr< NEML2ModelHandle > NEML2ModelInterface< T >::makeModelHandle ( const InputParameters &  params, NEML2ModelInterface< T > &  self  )

inlinestaticprivate

Parse the runtime parameters and build the model handle.

Static so it can run in the member initializer list (before _device, which is read back from the handle). self supplies the MooseObject services (paramError, comm, app) once the T base is constructed.

Definition at line 75 of file NEML2ModelInterface.h.

  {
    // Route NEML2's log output to MOOSE's console before the model is built/evaluated.
    NEML2Utils::redirectLogsToConsole(self._console);
 
    const bool eager = params.get<bool>("eager");
    const auto load = loadExtensionPaths(params);
 
    // 'load' imports Python extensions into the embedded interpreter -- cpp-eager only.
    if (!eager && !load.empty())
      self.paramError("load",
                      "The 'load' parameter imports Python extensions into the embedded "
                      "interpreter and is only valid with the cpp-eager runtime. Set eager=true to "
                      "use it, or remove it for the cpp-aoti runtime.");
    // 'input' is the source .i (cpp-eager) or the neml2-compile artifact folder (cpp-aoti).
    if (!params.isParamValid("input"))
      self.paramError(
          "input",
          "'input' (the NEML2 source '.i' for cpp-eager, or the neml2-compile artifact folder for "
          "cpp-aoti) is required.");
 
    // Device list (default: the app's libtorch compute device).
    auto devices = params.get<std::vector<std::string>>("device");
    if (devices.empty())
      devices = {at::Device(self.getMooseApp().getLibtorchDevice()).str()};
 
    // Per-device chunk sizes (default 0 = whole batch); length 1 broadcasts, else matches 'device'.
    auto db = params.get<std::vector<unsigned int>>("device_batch");
    if (db.empty())
      db = {0};
    if (db.size() != 1 && db.size() != devices.size())
      self.paramError("device_batch",
                      "'device_batch' must have length 1 or the same length as 'device'.");
    const std::vector<std::size_t> batch_sizes(db.begin(), db.end());
 
    // Host MPI communicator for the cpp-aoti CUDA scheduler (read at construction only; unused for
    // cpp-eager and for CPU dispatch).
    const auto mpi_comm = self.comm().get();
 
    // MOOSE drives NEML2 exclusively in Float64. Force the libtorch default dtype to Float64 while
    // the model is built so a parameter-less eager model (whose dtype would otherwise fall back to
    // libtorch's Float32 default) resolves to Float64 and accepts MOOSE's double inputs.
    const NEML2Utils::ScopedDefaultDtype scoped_dtype;
 
    return makeNEML2ModelHandle(eager,
                                std::string(params.get<DataFileName>("input")),
                                params.get<std::string>("model"),
                                devices,
                                batch_sizes,
                                &mpi_comm,
                                load);
  }

model() [1/2]

template<class T >

NEML2ModelHandle & NEML2ModelInterface< T >::model ( )

inlineprotected

Non-const access to the NEML2 model handle, for runtime parameter writes (set_parameter)

Definition at line 49 of file NEML2ModelInterface.h.

{ return *_model; }

model() [2/2]

template<class T >

const NEML2ModelHandle & NEML2ModelInterface< T >::model ( ) const

inlineprotected

Get the NEML2 model handle (cpp-aoti by default, cpp-eager when eager=true)

Definition at line 46 of file NEML2ModelInterface.h.

{ return *_model; }

neml2NumThreads()

template<class T >

unsigned int NEML2ModelInterface< T >::neml2NumThreads ( ) const

inlineprotected

libtorch intra-op thread count to use while evaluating the NEML2 model.

Wrap the evaluation region in NEML2Utils::ScopedNumThreads guard(neml2NumThreads()); so the count is applied only during evaluation and the previous value is restored afterward.

Definition at line 60 of file NEML2ModelInterface.h.

{ return _num_threads; }

output_device()

template<class T >

const at::Device & NEML2ModelInterface< T >::output_device ( ) const

inlineprotected

Get the target output device.

Definition at line 55 of file NEML2ModelInterface.h.

{ return _output_device; }

validateModel()

template<class T >

void NEML2ModelInterface< T >::validateModel ( ) const

protectedvirtual

Validate the NEML2 material model.

Note that the developer is responsible for calling this method at the appropriate times, for example, at initialSetup().

Definition at line 250 of file NEML2ModelInterface.h.

{
  // NEML2 v3 provides no C++ diagnose() equivalent; constructing the eager model above already
  // parses and validates it. Kept as a hook for derived classes.
}

validParams()

template<class T >

InputParameters NEML2ModelInterface< T >::validParams ( )

static

Definition at line 146 of file NEML2ModelInterface.h.

{
  InputParameters params = T::validParams();
  params.addParam<bool>(
      "eager",
      false,
      "Use the cpp-eager NEML2 runtime, which embeds a Python interpreter and loads the model "
      "from the source 'input' file at runtime. Default false -- use the cpp-aoti runtime, which "
      "loads an ahead-of-time-compiled artifact (no Python at runtime) from the artifact folder "
      "given by 'input'. The 'load' parameter is only valid with eager=true.");
  params.addParam<DataFileName>(
      "input",
      "For the cpp-eager runtime, the path to the source NEML2 '.i' model file. For the cpp-aoti "
      "runtime, any of: the artifact folder produced by neml2-compile (a self-describing directory "
      "holding metadata.json plus per-<device>/<dtype> '.pt2' binaries), the metadata.json inside "
      "it, or the neml2-compile stub '.i'. Prefer the artifact folder -- it is self-contained and "
      "relocatable as a unit, and needs no 'model' name. The stub '.i' is accepted for "
      "compatibility but points at its artifact folder by path, so the two must stay together.");
  params.addParam<std::vector<std::string>>(
      "cli_args",
      {},
      "Additional command line arguments to use when parsing the NEML2 input file.");
  params.addParam<std::vector<DataFileName>>(
      "load",
      {},
      "External Python extension files (paths to .py files or package directories, resolved "
      "relative to the input file and the application's data search path) imported into the "
      "embedded interpreter before the NEML2 model is built. Importing them registers any "
      "@register_neml2_object types they define -- e.g. NEML2 models hosted inside a MOOSE app -- "
      "so the NEML2 input file can reference them. Mirrors the neml2 CLI --load flag. Only valid "
      "with the cpp-eager runtime (eager=true).");
  params.addParam<std::string>(
      "model",
      "",
      "Name of the NEML2 model, i.e., the string inside the brackets [] in the NEML2 input file "
      "that corresponds to the model you want to use. Required for the cpp-eager runtime, and for "
      "cpp-aoti only when 'input' is a stub '.i'; the cpp-aoti folder and metadata.json forms use "
      "the single model recorded in the artifact metadata and ignore it.");
  params.addParam<std::vector<std::string>>(
      "device",
      {},
      "Compute device(s) for the NEML2 model, each following the schema "
      "(cpu|cuda)[:<device-index>] "
      "(e.g. 'cpu', 'cuda:1'). cpp-eager pins the model to the first device. cpp-aoti dispatches "
      "over the list: a CUDA list is round-robined over the MPI ranks on each node (one device per "
      "rank when counts match, sharing when there are more ranks than devices, an error when there "
      "are fewer); a single 'cpu' runs each rank's local batch on CPU. If not specified, defaults "
      "to the compute device from the --compute-device command line argument.");
  params.addParam<std::vector<unsigned int>>(
      "device_batch",
      {},
      "Per-device chunk size along the leading batch axis (0 = run the whole local batch at once). "
      "Must have length 1 (broadcast to all devices) or match the length of 'device'. Defaults to "
      "0.");
  params.addParam<std::string>(
      "output_device",
      "Device on which to store the model outputs, following the same schema as a single 'device' "
      "entry. Defaults to cpu, where MOOSE consumes them. Set this only if a downstream object "
      "needs "
      "the outputs on a specific device.");
  params.addParam<unsigned int>(
      "num_threads",
      "Number of threads to use when evaluating the NEML2 model. Defaults to the number of threads "
      "MOOSE is run with (the --n-threads command-line value). Some models evaluate fastest with a "
      "single thread regardless of --n-threads; set this to 1 in that case.");
 
  return params;
}

Referenced by NEML2ActionCommon::commonParams(), and NEML2ModelExecutor::validParams().

Member Data Documentation

_device

template<class T >

const at::Device NEML2ModelInterface< T >::_device

private

The device on which to evaluate the NEML2 model (this rank's scheduler-assigned device)

Definition at line 134 of file NEML2ModelInterface.h.

Referenced by NEML2ModelInterface< T >::device().

_model

template<class T >

std::unique_ptr<NEML2ModelHandle> NEML2ModelInterface< T >::_model

private

The NEML2 material model handle.

Defaults to the cpp-aoti runtime (loads an ahead-of-time- compiled artifact from the 'input' folder); the cpp-eager runtime (embeds a CPython interpreter and loads the model from the source 'input' file) is opt-in via the 'eager' parameter.

Definition at line 132 of file NEML2ModelInterface.h.

Referenced by NEML2ModelInterface< T >::model(), and NEML2ModelInterface< T >::model().

_num_threads

template<class T >

const unsigned int NEML2ModelInterface< T >::_num_threads

private

libtorch intra-op thread count applied (scoped) during model evaluation; defaults to the MOOSE/libMesh thread count.

Definition at line 139 of file NEML2ModelInterface.h.

Referenced by NEML2ModelInterface< T >::neml2NumThreads().

_output_device

template<class T >

const at::Device NEML2ModelInterface< T >::_output_device

private

The device on which to store the outputs.

Definition at line 136 of file NEML2ModelInterface.h.

Referenced by NEML2ModelInterface< T >::output_device().

---

The documentation for this class was generated from the following file:

• include/neml2/interfaces/NEML2ModelInterface.h