tensormesh.distributed

Mesh partitioning and parallel assembly across multiple devices, with integration into torch-sla’s distributed sparse solver. See the user guide chapter for a worked walkthrough.

DistributedMesh

class DistributedMesh(mesh: Mesh, num_partitions: int | None = None, method: str = 'coordinate', devices: List[device] | None = None)

Bases: object

Partitioned mesh for multi-GPU parallel assembly.

Wraps partition_mesh to split a global mesh into submeshes, each assigned to a separate device. Each submesh stores an orig_nid mapping (local node index → global node index) in point_data.

Parameters:

• mesh (Mesh) – Global mesh (typically on CPU).

• num_partitions (int, optional) – Number of partitions. Defaults to torch.cuda.device_count(), or 2 if no CUDA devices are available.

• method (str, optional) – Partitioning method: 'coordinate' (default, fast RCB), 'spectral', or 'metis'.

• devices (list of device, optional) – Devices to assign partitions to. Defaults to cuda:0, cuda:1, ... or cpu if CUDA is unavailable.

Examples

>>> mesh = tm.Mesh.gen_rectangle(chara_length=0.05)
>>> dmesh = DistributedMesh(mesh, num_partitions=4)
>>> print(dmesh.num_partitions, dmesh.n_global_points)

__init__(mesh: Mesh, num_partitions: int | None = None, method: str = 'coordinate', devices: List[device] | None = None)

DSparseMatrix

class DSparseMatrix(dsparse_tensor: DSparseTensor, partition_uuid: int | None = None)

Bases: _FEMSparsityMixin, DSparseTensor

Distributed FEM sparse matrix.

Subclass of torch_sla.distributed.DSparseTensor so isinstance(D, DSparseTensor) holds and torch-sla’s free functions (solve / io / arithmetic dispatch) accept it.

Constructed by wrapping an existing DSparseTensor (typically returned by DSparseTensor.partition() or distributed_element_assemble()).

Parameters:

• dsparse_tensor – The underlying distributed sparse tensor. DSparseMatrix copies its state and adds the FEM identity layer on top.

• partition_uuid – 63-bit UUID identifying the partition build. Pass when deriving a new DSparseMatrix from an existing one (so caches are shared); leave as None to freshly generate + broadcast.

__init__(dsparse_tensor: DSparseTensor, partition_uuid: int | None = None)

property partition_uuid: int

property partition

The Partition from the spec.

property row_indices: Tensor

Local row indices (in local coordinates).

property col_indices: Tensor

Local column indices (in local coordinates).

property layout_signature: Tuple

Sequence-identity signature scoped to the partition build.

Combines the mixin’s local sequence identity with the partition UUID so two matrices on the same partition share signatures (and therefore caches) while two matrices on independently-built partitions do not, even if local layouts coincidentally match.

to(*args, **kw) → DSparseMatrix

cuda(device=None) → DSparseMatrix

cpu() → DSparseMatrix

double() → DSparseMatrix

float() → DSparseMatrix

to_single() → SparseMatrix

Allgather the distributed matrix into a global single-device SparseMatrix materialised on every rank.

Used as a temporary bridge to FEM operators that have not yet learned to consume DSparseMatrix directly (notably Condenser). Costs an all-gather of the COO triples; avoid in hot paths.

The @distributed decorator

distributed(asm_cls: Type[T]) → Type[T]

Class decorator: turn an Assembler class into a distributed one.

The wrapped class shares the original’s weak form, quadrature setup and element kernel – only the entry points (from_mesh + __call__) are swapped to run a distributed assembly path.

Parameters:

asm_cls – An ElementAssembler or NodeAssembler subclass. The decision between matrix- and vector-flavoured assembly is taken at the wrapped __call__ based on the original class type.

Returns:

• A new subclass of asm_cls with overridden entry points; the

• name is prefixed with Distributed for repr / debugging clarity.

Distributed assembly

distributed_element_assemble(assembler_cls: Type[ElementAssembler], dmesh: DistributedMesh, quadrature_order: int = 2, project: str = 'reduce', call_kwargs: dict | None = None, **assembler_kwargs) → DSparseTensor

Assemble element matrix in parallel across multiple devices.

Assemblers are created sequentially (for CUDA thread-safety), then assembly computation runs in parallel threads on separate GPUs.

Parameters:

• assembler_cls (Type[ElementAssembler]) – The assembler class (e.g., LaplaceElementAssembler).

• dmesh (DistributedMesh) – Partitioned mesh with device assignments.

• quadrature_order (int, optional) – Quadrature order for integration. Default: 2.

• project (str, optional) – Projection method: 'reduce' or 'sparse'. Default: 'reduce'.

• call_kwargs (dict, optional) – Extra keyword arguments passed to assembler.__call__() (e.g., point_data, scalar_data).

• **assembler_kwargs – Extra keyword arguments passed to assembler_cls.from_mesh().

Returns:

Distributed sparse matrix ready for distributed solve.

Return type:

DSparseTensor

distributed_element_assemble_per_rank(assembler_cls: Type[ElementAssembler], dmesh: DistributedMesh, rank: int, quadrature_order: int = 2, project: str = 'reduce', call_kwargs: dict | None = None, **assembler_kwargs) → DSparseTensor

True distributed assembly: each rank does ONLY its submesh.

Pipeline (every rank in lockstep, torch.distributed must be init’d):

1. Compute partition_ids deterministically from dmesh (lowest rank wins for shared boundary nodes). No communication.

2. This rank assembles submesh[rank] -> local triples in global coords. No threading, no work for other ranks’ submeshes.

3. For each local triple, route it to the rank that owns its row. all_to_all_single exchange (NCCL on CUDA, gloo on CPU).

4. Coalesce incoming triples (sum duplicate (row,col) pairs from elements assembled by different ranks that touch the same row).

5. Discover halo columns: distinct columns in received triples that aren’t owned by this rank.

6. Build Partition + remap to local coords + DSparseTensor.from_sparse_local.

Total compute is roughly the same as single-process assembly divided across ranks, instead of duplicated.

Caller is responsible for initialising torch.distributed before calling. The (rank, world) is taken from dist.

Returns a row-shard DSparseTensor with each rank holding only its owned-row contributions.

distributed_element_assemble_to_sparse(assembler_cls: Type[ElementAssembler], dmesh: DistributedMesh, quadrature_order: int = 2, project: str = 'reduce', call_kwargs: dict | None = None, **assembler_kwargs) → SparseMatrix

Assemble element matrix in parallel, returning a global SparseMatrix.

Same as distributed_element_assemble() but returns a standard SparseMatrix instead of torch-sla’s DSparseTensor.

distributed_node_assemble(assembler_cls: Type[NodeAssembler], dmesh: DistributedMesh, quadrature_order: int = 2, project: str = 'reduce', point_data: Dict[str, Tensor] | None = None, call_kwargs: dict | None = None, **assembler_kwargs) → Tensor

Assemble node vector (RHS) in parallel across multiple devices.

Utilities

broadcast_from_rank0(factory: Callable[[], Tensor], *, dst_device: device | None = None) → Tensor

Sample on rank 0, broadcast to every rank.

Parameters:

• factory – Zero-arg callable returning a fresh torch.Tensor. Called only on rank 0 when a process group is active; called on every rank in single-process / no-process-group mode.

• dst_device – Optional override for the device the broadcast lands on. Defaults to CUDA when NCCL is the active backend, CPU otherwise; the returned tensor is moved back to CPU so it composes with caller code that does its own .to(device) later. Pass an explicit device if you want a different placement.

Returns:

A tensor with the same content on every rank. Single-process mode returns factory() directly.

Return type:

Tensor

Notes

Uses torch.distributed.broadcast_object_list() under the hood, which pickles the tensor. This is fine for the typical use case (small coefficient matrices, scalar parameters); for large buffers prefer manual dist.broadcast with pre-allocated receive buffers on every rank.