NEML2ModelHandle.h File Reference

Go to the source code of this file.

Classes
class	NEML2ModelHandle
	Runtime-agnostic handle over a NEML2 v3 model. More...
class	EagerModelHandle
	Handle wrapping the cpp-eager runtime (embeds CPython; loads from a source .i). More...
class	AOTIModelHandle
	Handle wrapping the cpp-aoti runtime: a neml2::aoti::DispatchedModel backed by the self-describing artifact folder that neml2-compile writes (a shared metadata.json plus per-<device>/<dtype>/ .pt2 binaries). More...

Functions
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)
	Construct a NEML2 model handle for the requested runtime.

Function Documentation

makeNEML2ModelHandle()

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  )

inline

Construct a NEML2 model handle for the requested runtime.

Parameters

eager	When true, build a cpp-eager model from the source .i; when false, build a cpp-aoti dispatched model from the compiled artifact.
input	For cpp-eager, the path to the source NEML2 .i. For cpp-aoti, any of: the neml2-compile artifact folder, the metadata.json inside it, or the legacy stub .i (see AOTIModelHandle::load).
model	Name of the model in the .i. Required for cpp-eager, and for cpp-aoti only when input is a stub .i; the folder / metadata.json forms ignore it.
devices	Compute device(s). cpp-eager pins to the first; cpp-aoti dispatches over the list (a CUDA list round-robins over the MPI ranks; a single CPU runs locally).
batch_sizes	Per-device chunk size (0 = whole batch). Length 1 broadcasts to all devices.
comm	Pointer to the host MPI communicator (nullptr => MPI_COMM_WORLD); used by the CUDA MPI scheduler only.
load	External Python extension paths to import before building (cpp-eager only).

Definition at line 248 of file NEML2ModelHandle.h.

{
  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);
}

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