API: Connector

API: Connector#

Interface#

struct ModelProperties#

Properties of the model that can be queried by the consumer of the interface, e.g., LAMMPS.

Public Members

double cutoff#

The cutoff distance for the potential.

double comm_dist = 0.0#

Minimum distance to local atoms for which ghost atoms must be communicated.

const char *unit_style#

The unit style of the model.

struct jcn::ModelProperties::[anonymous] neighbor_list#

The neighbor list settings of the model.

struct ConnectorConfig#

Configurations for the evaluation of the potential

Public Members

std::string backend#

String identifying the backend on which to evaluate the model. The backend name is inferred from the name of the PjRt plugin file.

int device = 0#

Integer specifying the local device ID on which to exectute the model.

float memory_fraction = 0.75f#

Fraction of device memory to allocate to execute the model.

struct Statistics#

Statistics of the computation, e.g., the expected FLOPs

Public Members

double flops#

The estimated number of floating point operations.

bool recompiled#

Whether recompilation was necessary

struct Results#

Results of the evaluation of the potential.

Public Members

double potential#

The potential energy of the system.

Statistics stats#

The statistics of the computation.

class Connector#

Pointer to implementation class.

Public Functions

ModelProperties load_model(ModelConfig config)#

Loads the model and initializes the interface. @params config: The configuration of the model and interface.

Returns:

The properties of the model.

Results compute_force(int lnum, int gnum, double **x, double **f, int *type, int inum, int *ilist, int *numneigh, int **firstneigh, bool list_changed, bool allow_recompile)#

Computes the forces for a system given the atom positions and (if required) the neighbor list.

Parameters:

  • lnum – Number of local atoms in the domain.

  • gnum – Number of ghost atoms associated with the domain.

  • x – Pointer to the atom positions for local and ghost atoms.

  • f – Pointer to the force array for local and ghost atoms. Array is overwritten with computed forces by the interface.

  • type – Pointer to the atom types (species) for local and ghost atoms.

  • inum – Number of sender atoms in the domain, corresponding to the rows of the neighbor list.

  • ilist – Array holding the local indices of the senders.

  • numneigh – Array holding the number of neighbors for each sender atom.

  • firstneigh – Array holding the local indices of the neighbor atoms.

  • list_changed – Whether the neighbor list has changed since the the last call to the function.

  • allow_recompile – Allows recompilation of the executable. If set to false while a recompilation is required, will throw a RecompilationRequired exception.

Returns:

The potential energy of the system and statistics of the computation.

Public Static Functions

static void initialize()#

Initializes the PjRt backends.

class Impl#

Data Transfer#

Atoms#

class AtomBuilder#

Transforms atom data from the local domain into a padded XLA-copmpatible format

Public Functions

AtomBuilder(float atom_multiplier, bool newton)#

Constructor

Parameters:

atom_multiplier – Fraction of extra atoms to consider when re-allocate the arrays @params newton True if the forces should be computed according to LAMMPS newton setting

std::vector<xla::PjRtBuffer*> build_domain(xla::PjRtClient *client, int device_id, int inum, int gnum, double **x, int *type)#

Writes atom positions into device buffers. Padds the atom data to reduce the number of recompilations.

Parameters:

  • client – PjRt client to allocate buffers

  • device_id – Device ID to allocate buffers

  • inum – Number of local atoms

  • gnum – Number of ghost atoms

  • x – Pointer to atom positions

  • type – Pointer to atom types (one-based species)

Returns:

A vector holding references to the buffers

double evaluate_domain(bool success, int inum, int gnum, double **f, std::vector<std::vector<std::unique_ptr<xla::PjRtBuffer>>> &results)#

Writes back the force to the original array and returns the potential

Parameters:

  • success – True if the computation was successful, i.e., the neighbor list did not overflow.

  • inum – Number of local atoms

  • f – Pointer to target force array

  • results – Vector of vector of pointers to the result buffers.

Returns:

The potential energy of the system

struct AtomShapes#

Contains all per-atom data of a domnain

Public Members

int n_atoms#

Maximum possible number of atoms in the domain, including local, ghost, and padded atoms.

bool reallocate#

Tracks change of maxmimum number of atoms in the domain. If true, must re-compile the mlir module with new shapes.

Neighbor Lists#

class GraphBuilder#

Abstract baseclass for all graph / neighborlist builders.

Subclassed by jcn::DeviceSparseNeighborList, jcn::SimpleDenseNeighborList, jcn::SimpleSparseNeighborList

class SimpleSparseNeighborList : public jcn::GraphBuilder#

Class that interfaces neighborlists, e.g., from LAMMPS, to a sparse neighbor list in chemtrain.

Public Functions

virtual void initialize(std::vector<float> multipliers) override#

Initializes the interface.

Simple sparse neighbor list

Parameters:

multipliers – A vector of multipliers specifying the relative increase of the neighborlist buffers. Only the first element is used.

virtual NeighborListShapes get_neighbor_list_shapes(int max_atoms, int inum, int *numneigh, bool check_buffers) override#

Returns the required shapes and types of the sparse neighborlist based on the reference neighborlist.

Parameters:

  • max_atoms – The maximum number of atoms in the system.

  • inum – The number of local atoms.

  • numneigh – Array holding the number of neighbors for each atom.

Returns:

Returns NeighborListShapes struct with necessary dimensions.

virtual std::vector<xla::PjRtBuffer*> build_graph(xla::PjRtClient *client, int device_id, int inum, int *ilist, int *numneigh, int **firstneigh, bool update) override#

Builds the sparse neighborlist from the reference neighborlist.

Parameters:

  • client – The PjRt client to allocate buffers.

  • device_id – The device ID on which to allocate buffers.

  • inum – The number of local atoms.

  • ilist – Array holding the indices of the senders of each neighborlist entry.

  • numneigh – Array holding the number of neighbors for each atom.

  • firstneigh – Array holding the index of the neighbors (receivers) for each sender.

  • update – Whether the neighbor list data must be updated.

Returns:

A vector holding references to the buffers.

class SimpleDenseNeighborList : public jcn::GraphBuilder#

Class that interfaces neighborlists, e.g., from LAMMPS, to a dense neighbor list in chemtrain.

Public Functions

virtual void initialize(std::vector<float> multipliers) override#

Initializes the interface.

Simple Dense Neighbor List

Parameters:

multipliers – A vector of multipliers specifying the relative increase of the neighborlist buffers. Only the first element is used.

virtual NeighborListShapes get_neighbor_list_shapes(int max_atoms, int inum, int *numneigh, bool check_buffers) override#

Returns the required shapes and types of the sparse neighborlist based on the reference neighborlist.

Parameters:

  • max_atoms – The maximum number of atoms in the system.

  • inum – The number of local atoms.

  • numneigh – Array holding the number of neighbors for each atom.

Returns:

Returns NeighborListShapes struct with necessary dimensions.

virtual std::vector<xla::PjRtBuffer*> build_graph(xla::PjRtClient *client, int device_id, int inum, int *ilist, int *numneigh, int **firstneigh, bool update) override#

Builds the sparse neighborlist from the reference neighborlist.

Parameters:

  • client – The PjRt client to allocate buffers.

  • device_id – The device ID on which to allocate buffers.

  • inum – The number of local atoms.

  • ilist – Array holding the indices of the senders of each neighborlist entry.

  • numneigh – Array holding the number of neighbors for each atom.

  • firstneigh – Array holding the index of the neighbors (receivers) for each sender.

  • update – Whether the neighbor list data must be updated.

Returns:

A vector holding references to the buffers.

class DeviceSparseNeighborList : public jcn::GraphBuilder#

Public Functions

virtual void initialize(std::vector<float> multipliers) override#

Initializes the interface.

@params multipliers: A vector of multipliers specifying the relative increase of the neighborlist buffers. TODO: Document the required multipliers.

Device Sparse Neighbor List

virtual NeighborListShapes get_neighbor_list_shapes(int max_atoms, int inum, int *numneigh, bool check_buffers) override#

Returns the required shapes and types of the sparse neighborlist based on the reference neighborlist.

Parameters:

  • max_atoms – The maximum number of atoms in the system.

  • inum – The number of local atoms.

  • numneigh – Array holding the number of neighbors for each atom.

Returns:

Returns NeighborListShapes struct with necessary dimensions.

virtual std::vector<xla::PjRtBuffer*> build_graph(xla::PjRtClient *client, int device_id, int inum, int *ilist, int *numneigh, int **firstneigh, bool update) override#

Builds the sparse neighborlist from the reference neighborlist.

Parameters:

  • client – The PjRt client to allocate buffers.

  • device_id – The device ID on which to allocate buffers.

  • inum – The number of local atoms.

  • ilist – Array holding the indices of the senders of each neighborlist entry.

  • numneigh – Array holding the number of neighbors for each atom.

  • firstneigh – Array holding the index of the neighbors (receivers) for each sender.

  • update – Whether the neighbor list should be regenerated, e.g., due to newly communicated atoms.

Returns:

A vector holding references to the buffers.

virtual bool evaluate_statistics(std::map<std::string, std::unique_ptr<xla::PjRtBuffer>> statistics, bool check_buffers) override#

Evaluates the statistics of the neighborlist computation and increase buffer dimensions if necessary.

This class generates the neighbor list completely on the device using a cell list. Therefore, it can be necessary to increase the buffers of the cell list and neighborlist or the dimension of the cell grid. Cells must be at least as large as the cutoff radius, but their size might change in NPT simulations.

Parameters:

results – A vector of vectors holding references to all result buffers.

Returns:

True if the neighborlist generation was sucessful and no overflow occured.

struct NeighborListShapes#

Contains shapes and types of the graph input buffers.

Public Members

std::vector<std::vector<int64_t>> graph_shapes#

Shapes of the neighborlist inputs.

std::vector<xla::PrimitiveType> graph_types#

Types of the neighborlist inputs.

bool reallocate#

Indicates whether a shape change requires recompilation.

Contents