fhelium.backend.distributed

Process-group resources and registered collective execution.

PROCESS_GROUP_RESOURCE_KIND

constant View source

PROCESS_GROUP_RESOURCE_KIND = 'process-group'

TorchBroadcastImplementation

class View source

TorchBroadcastImplementation(name: str = 'torch-broadcast', supports_in_place: bool = False, operation_types: tuple[type[Operation], ...] = (distributed.BroadcastOp,))

Bases: _OperandResourceImplementation

Functionally broadcast one equal-layout Tensor from a group rank.

Attributes

Name	Type	Default/value
name	str	'torch-broadcast'
supports_in_place	bool	False
operation_types	tuple[type[Operation], ...]	(distributed.BroadcastOp,)

execute

method

def execute(invocation: OperationInvocation, inputs: tuple[torch.Tensor, ...], resources: tuple[BoundResource, ...], *, in_place: bool) -> tuple[torch.Tensor, ...]: ...

TorchCiphertextAddAllReduceImplementation

class View source

TorchCiphertextAddAllReduceImplementation(combine: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor], name: str = 'torch-ciphertext-add-all-reduce', supports_in_place: bool = False, operation_types: tuple[type[Operation], ...] = (distributed.AllReduceAddCiphertextOp,))

Bases: _OperandResourceImplementation

Reduce ciphertext payloads with an injected Backend Tensor addition.

Attributes

Name	Type	Default/value
combine	Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]	field(repr=False, compare=False)
name	str	'torch-ciphertext-add-all-reduce'
supports_in_place	bool	False
operation_types	tuple[type[Operation], ...]	(distributed.AllReduceAddCiphertextOp,)

execute

method

def execute(invocation: OperationInvocation, inputs: tuple[torch.Tensor, ...], resources: tuple[BoundResource, ...], *, in_place: bool) -> tuple[torch.Tensor, ...]: ...

TorchGenericAllReduceImplementation

class View source

TorchGenericAllReduceImplementation(name: str = 'torch-generic-all-reduce', supports_in_place: bool = False, operation_types: tuple[type[Operation], ...] = (distributed.AllReduceOp,))

Bases: _OperandResourceImplementation

Fold all rank-local payloads with the operation's combine region.

Attributes

Name	Type	Default/value
name	str	'torch-generic-all-reduce'
supports_in_place	bool	False
operation_types	tuple[type[Operation], ...]	(distributed.AllReduceOp,)

execute

method

def execute(invocation: OperationInvocation, inputs: tuple[torch.Tensor, ...], resources: tuple[BoundResource, ...], *, in_place: bool) -> tuple[torch.Tensor, ...]: ...

execute_regions

method

def execute_regions(invocation: OperationInvocation, inputs: tuple[torch.Tensor, ...], resources: tuple[BoundResource, ...], regions: tuple[TensorRegion, ...], *, in_place: bool) -> tuple[torch.Tensor, ...]: ...

ProcessGroupExecutionResource

class View source

ProcessGroupExecutionResource(group: torch.distributed.ProcessGroup | None = None)

Expose one initialized rank-local process group to Backend operations.

group=None names PyTorch's default process group. Ranks represented by distributed IR are process-group ranks; conversion to a global rank occurs only when a PyTorch collective requires it.

Attributes

Name	Type	Default/value
group	torch.distributed.ProcessGroup | None	None

rank

property

rank: int

Return this process's rank within the selected group.

size

property

size: int

Return the number of participating group ranks.

global_rank

method

def global_rank(group_rank: int) -> int: ...

Map one group-relative rank to its process-global rank.

TorchProcessGroupQueryImplementation

class View source

TorchProcessGroupQueryImplementation(name: str = 'torch-process-group-query', supports_in_place: bool = False, operation_types: tuple[type[Operation], ...] = (distributed.RankOp, distributed.GroupSizeOp))

Bases: _OperandResourceImplementation

Read group rank and size as scalar index Tensors.

Attributes

Name	Type	Default/value
name	str	'torch-process-group-query'
supports_in_place	bool	False
operation_types	tuple[type[Operation], ...]	(distributed.RankOp, distributed.GroupSizeOp)

execute

method

def execute(invocation: OperationInvocation, inputs: tuple[torch.Tensor, ...], resources: tuple[BoundResource, ...], *, in_place: bool) -> tuple[torch.Tensor, ...]: ...

distributed_operation_contributions

function View source

def distributed_operation_contributions(*, ciphertext_add: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor] | None=None) -> tuple[TorchProcessGroupQueryImplementation | TorchBroadcastImplementation | TorchGenericAllReduceImplementation | TorchCiphertextAddAllReduceImplementation, ...]: ...

Return distributed implementations for one execution owner.

The specialized ciphertext-add implementation is contributed only when an execution owner supplies addition through its existing Tensor Backend. A Compile pass may instead lower the specialized operation to generic all-reduce with a visible combine region.

prepare_ciphertext_add_combine

function View source

def prepare_ciphertext_add_combine(backend: OperationBackend, *, modulus_basis: str='Q', implementation: str | None=None) -> Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]: ...

Bind ciphertext addition through an existing operation backend.

The returned callable uses the selected RNS implementation and the arithmetic parameters supplied by each distributed reduction.