tensormesh.distributed

Mesh partitioning and parallel assembly across multiple devices, with integration into torch-sla's distributed sparse solver. See 分布式 FEM for a worked walkthrough.

DistributedMesh

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

基类：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.

参数:

• 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.

示例

>>> 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)

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.

参数:

• 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().

返回:

Distributed sparse matrix ready for distributed solve.

返回类型:

DSparseTensor

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.