FHElium

Appearance

fhelium.distributed ​

PyTorch distributed operations and CKKS value collectives.

For ordinary torch.Tensor values, this namespace exposes selected torch.distributed functions with their native signatures, mutation rules, ProcessGroup behavior, and Work handles. FHElium value operations use typed names such as broadcast_ciphertext and all_reduce_ciphertext.

Backend ​

re-export

python
Backend = torch.distributed.Backend

P2POp ​

re-export

python
P2POp = torch.distributed.P2POp

ProcessGroup ​

re-export

python
ProcessGroup = torch.distributed.ProcessGroup

ReduceOp ​

re-export

python
ReduceOp = torch.distributed.ReduceOp

Work ​

re-export

python
Work = torch.distributed.Work

all_gather ​

re-export

python
all_gather = torch.distributed.all_gather

all_gather_ciphertexts ​

function View source

python
def all_gather_ciphertexts(value: Ciphertext, *, group: torch.distributed.ProcessGroup | None=None) -> list[Ciphertext]: ...

All-gather independent ciphertexts with one identical layout.

Equality covers context, level, scale, components, prime IDs, domain, basis, Montgomery state, degree, shape, and dtype. Rank-local CUDA device indices may differ because descriptors intentionally preserve device type rather than one global device index. This is transport of independent values, not algebraic reduction. The operation is synchronous, accepts no async_op argument, and returns no torch.distributed.Work.

Parameters

  • value: One independent rank-local ciphertext. All group ranks must provide ciphertexts with identical layout metadata. The input is read but not mutated.
  • group: Participating process group. None selects the default process group. No root rank is used.

Returns

A newly allocated list on every rank containing one ciphertext per group rank in process-group-rank order. Entries, including the caller's entry, are new ciphertext objects whose tensors use the same device as the caller's input. For world size one, a one-element copied result is returned.

Raises

  • TypeError: If value is not a ciphertext.
  • ValueError: If ciphertext shapes or arithmetic metadata differ across group ranks.
  • RuntimeError: If distributed communication is uninitialized or the caller is not a member of group.

all_gather_compressed_plaintexts ​

function View source

python
def all_gather_compressed_plaintexts(value: CompressedPlaintext, *, group: torch.distributed.ProcessGroup | None=None) -> list[CompressedPlaintext]: ...

All-gather independent compressed plaintexts of one exact layout.

all_gather_plaintexts ​

function View source

python
def all_gather_plaintexts(value: Plaintext, *, group: torch.distributed.ProcessGroup | None=None) -> list[Plaintext]: ...

All-gather independent plaintexts with one arithmetic state.

Every rank must provide the same message/encoded representation kind, shape, dtype, context, level, scale, layout, domain, basis, Montgomery state, and prime IDs. This is representation-preserving transport, not an arithmetic operation. It is synchronous, accepts no async_op argument, and returns no torch.distributed.Work.

Parameters

  • value: One independent rank-local plaintext. It is read but not mutated.
  • group: Participating process group. None selects the default process group. No root rank is used.

Returns

A newly allocated list on every rank containing one plaintext per group rank in process-group-rank order. Entries, including the caller's entry, are new plaintext objects whose tensors use the same device as the caller's input. For world size one, a one-element copied result is returned.

Raises

  • TypeError: If value is not a plaintext.
  • ValueError: If plaintext representations, shapes, or arithmetic metadata differ across group ranks.
  • RuntimeError: If distributed communication is uninitialized or the caller is not a member of group.

all_gather_single ​

data View source

python
all_gather_single = getattr(_torch_distributed, 'all_gather_single', _torch_distributed.all_gather_into_tensor)

all_reduce ​

re-export

python
all_reduce = torch.distributed.all_reduce

all_reduce_ciphertext ​

function View source

python
def all_reduce_ciphertext(value: Ciphertext, *, engine: CkksEngine, group: torch.distributed.ProcessGroup | None=None) -> None: ...

Synchronously sum ciphertext partials and update every rank in place.

The reduction phase uses the arbitrary-world-size binomial tree documented by reduce_ciphertext. The root then uses a representation-preserving tensor broadcast because broadcasting ciphertext payload bits requires no arithmetic specialization. Asynchronous composite Work semantics are intentionally deferred until a measured workload demonstrates useful communication/add overlap. The implicit root is the first global rank in process-group-rank order.

Parameters

  • value: Rank-local additive ciphertext partial. Every rank's object is overwritten in place with the component-wise CKKS sum modulo the active Q primes.
  • engine: Rank-local CKKS engine used during tree reduction. Every participating rank must provide an engine compatible with its value.
  • group: Participating process group. None selects the default process group. No caller-selected root is used.

Returns

None. Every rank's existing value object contains the global sum after return. For world size one, the value remains unchanged. The function accepts no async_op argument and returns no torch.distributed.Work.

Raises

  • ValueError: If rank-local values are not ciphertexts with identical shape and arithmetic metadata, or an engine rejects an incompatible value.
  • RuntimeError: If distributed communication is uninitialized or the caller is not a member of group.

all_to_all ​

re-export

python
all_to_all = torch.distributed.all_to_all

all_to_all_single ​

re-export

python
all_to_all_single = torch.distributed.all_to_all_single

barrier ​

re-export

python
barrier = torch.distributed.barrier

batch_isend_irecv ​

re-export

python
batch_isend_irecv = torch.distributed.batch_isend_irecv

broadcast ​

re-export

python
broadcast = torch.distributed.broadcast

broadcast_ciphertext ​

function View source

python
def broadcast_ciphertext(value_or_none: Ciphertext | None, *, src: int=0, group: torch.distributed.ProcessGroup | None=None) -> Ciphertext: ...

Broadcast one complete ciphertext and allocate non-source receivers.

A small internal object descriptor carries shape and CKKS metadata; the dense payload then uses the ordinary tensor broadcast. No placement or owner policy is inferred. The operation is synchronous: it accepts no async_op argument and returns no torch.distributed.Work.

Parameters

  • value_or_none: Ciphertext supplied only by src. Every non-source rank must pass None.
  • src: Global rank of the source process, which must belong to group. This is not a process-group-relative rank.
  • group: Participating process group. None selects the default process group.

Returns

The source's original ciphertext on src and a newly allocated, metadata-equivalent ciphertext on every other group rank. For a world-size-one group, the source object is returned unchanged and no payload communication or receiver allocation occurs.

Raises

  • ValueError: If src is outside group or ranks violate the source/non-source argument rules.
  • RuntimeError: If distributed communication is uninitialized, the caller is not a group member, or a valid transfer descriptor cannot be exchanged.

broadcast_compressed_plaintext ​

function View source

python
def broadcast_compressed_plaintext(value_or_none: CompressedPlaintext | None, *, src: int=0, group: torch.distributed.ProcessGroup | None=None) -> CompressedPlaintext: ...

Broadcast one exact compressed plaintext with typed allocation.

The compact tensor, ring dimension, encoded repetition layout, and all arithmetic metadata are preserved. The operation is synchronous and does not infer ownership, placement, or residency policy.

broadcast_key ​

function View source

python
def broadcast_key(key_or_none: _KeyT | None, *, src: int=0, group: torch.distributed.ProcessGroup | None=None) -> _KeyT: ...

Broadcast selected key material with typed allocation.

The function performs no key generation, owner inference, or automatic placement. Calling it is the program's visible decision to communicate one selected key. It is synchronous, accepts no async_op argument, and returns no torch.distributed.Work.

Parameters

  • key_or_none: Supported dense FHElium key supplied only by src. Every non-source rank must pass None.
  • src: Global rank of the source process, which must belong to group. This is not a process-group-relative rank.
  • group: Participating process group. None selects the default process group.

Returns

The source's original key object on src and a newly allocated key of the same exact type and metadata on every other group rank. For a world-size-one group, the source object is returned unchanged and no payload communication or receiver allocation occurs.

Raises

  • ValueError: If src is outside group, the source value is not a supported key, or ranks violate the source/non-source argument rules.
  • RuntimeError: If distributed communication is uninitialized, the caller is not a group member, or a valid transfer descriptor cannot be exchanged.

broadcast_plaintext ​

function View source

python
def broadcast_plaintext(value_or_none: Plaintext | None, *, src: int=0, group: torch.distributed.ProcessGroup | None=None) -> Plaintext: ...

Broadcast one complete plaintext and allocate non-source receivers.

The message or encoded data representation and its arithmetic metadata are preserved. The operation does not mutate the source plaintext, accepts no async_op argument, and returns no torch.distributed.Work.

Parameters

  • value_or_none: Plaintext supplied only by src. Every non-source rank must pass None.
  • src: Global rank of the source process, which must belong to group. This is not a process-group-relative rank.
  • group: Participating process group. None selects the default process group.

Returns

The source's original plaintext on src and a newly allocated, metadata-equivalent plaintext on every other group rank. For a world-size-one group, the source object is returned unchanged without payload communication or receiver allocation.

Raises

  • ValueError: If src is outside group or ranks violate the source/non-source argument rules.
  • RuntimeError: If distributed communication is uninitialized, the caller is not a group member, or a valid transfer descriptor cannot be exchanged.

destroy_process_group ​

re-export

python
destroy_process_group = torch.distributed.destroy_process_group

gather ​

re-export

python
gather = torch.distributed.gather

gather_ciphertext_limbs ​

function View source

python
def gather_ciphertext_limbs(local_shard: Ciphertext, *, dst: int=0, group: torch.distributed.ProcessGroup | None=None) -> Ciphertext | None: ...

Gather and reconstruct one ciphertext from disjoint RNS limb shards.

Group-rank order defines prime-interval order. On dst, compatible, nonempty, contiguous canonical intervals are concatenated along the RNS limb dimension. This is structural reconstruction, not ciphertext addition; use reduce_ciphertext for additive partials. The operation is synchronous, accepts no async_op argument, and returns no torch.distributed.Work.

Parameters

  • local_shard: Caller-owned rank-local ciphertext shard. It is read but not mutated.
  • dst: Global rank that reconstructs the ciphertext, which must belong to group. This is not a process-group-relative rank.
  • group: Participating process group. None selects the default process group.

Returns

The reconstructed ciphertext on dst and None on every other rank. The destination allocates receive buffers and the concatenated result. For world size one, the validated local_shard object is returned unchanged without concatenation.

Raises

  • ValueError: If dst is outside group or shards differ in logical ciphertext metadata, dtype, device, or contiguous prime layout.
  • RuntimeError: If distributed communication is uninitialized or the caller is not a member of group.

gather_ciphertexts ​

function View source

python
def gather_ciphertexts(value: Ciphertext, *, dst: int=0, group: torch.distributed.ProcessGroup | None=None) -> list[Ciphertext] | None: ...

Gather independent ciphertexts in process-group rank order.

This is representation-preserving transport, not an algebraic reduction. Use reduce_ciphertext when the inputs are additive partials of one logical encrypted value. The operation is synchronous, accepts no async_op argument, and returns no torch.distributed.Work.

Parameters

  • value: One independent rank-local ciphertext. It is read but not mutated.
  • dst: Global rank that receives the list, which must belong to group. This is not a process-group-relative rank.
  • group: Participating process group. None selects the default process group.

Returns

On dst, one ciphertext per group rank in process-group-rank order; the destination's list entry is its original object and other entries are newly allocated. Other ranks return None. For world size one, the sole rank returns [value] with the original object.

Raises

  • ValueError: If dst is outside group, a rank does not provide a ciphertext, or rank-local arguments disagree.
  • RuntimeError: If distributed communication is uninitialized or the caller is not a member of group.

get_backend ​

re-export

python
get_backend = torch.distributed.get_backend

get_global_rank ​

re-export

python
get_global_rank = torch.distributed.get_global_rank

get_group_rank ​

re-export

python
get_group_rank = torch.distributed.get_group_rank

get_process_group_ranks ​

re-export

python
get_process_group_ranks = torch.distributed.get_process_group_ranks

get_rank ​

re-export

python
get_rank = torch.distributed.get_rank

get_world_size ​

re-export

python
get_world_size = torch.distributed.get_world_size

init ​

function View source

python
def init(*, backend: str | None=None, init_method: str | None=None, timeout: timedelta | None=None, world_size: int=-1, rank: int=-1, store: Any | None=None, pg_options: Any | None=None, device_id: torch.device | int | None=None) -> None: ...

Initialize the PyTorch default process group and local CUDA device.

The helper uses PyTorch's standard process group and the torchrun rank environment. A direct world-size-one program uses a local HashStore while retaining ordinary PyTorch collective and asynchronous Work behavior. rank is the global rank in the newly created default process group; initialization completes synchronously and returns None.

Parameters

  • backend: Process-group backend. Defaults to "nccl" when CUDA is available and "gloo" otherwise.
  • init_method: Optional PyTorch rendezvous URL. When omitted for a direct world-size-one launch, a local HashStore is used unless store is supplied.
  • timeout: Optional process-group operation timeout forwarded to torch.distributed.init_process_group.
  • world_size: Number of processes in the default group. -1 resolves WORLD_SIZE from the environment.
  • rank: Global rank in the default group. -1 resolves RANK from the environment.
  • store: Optional rendezvous key-value store forwarded to PyTorch.
  • pg_options: Optional backend-specific process-group options forwarded to PyTorch.
  • device_id: Rank-local device. When omitted with CUDA available, cuda:LOCAL_RANK is selected. A CUDA device is required when CUDA is available; only a CPU device is accepted otherwise.

Returns

None. If a process group is already initialized, the selected CUDA device is still applied and process-group initialization is skipped.

Raises

  • ValueError: If device_id is incompatible with CUDA availability.
  • RuntimeError: If torch.distributed is unavailable or PyTorch cannot initialize the requested process group.

init_process_group ​

re-export

python
init_process_group = torch.distributed.init_process_group

irecv ​

re-export

python
irecv = torch.distributed.irecv

is_available ​

re-export

python
is_available = torch.distributed.is_available

is_initialized ​

re-export

python
is_initialized = torch.distributed.is_initialized

isend ​

re-export

python
isend = torch.distributed.isend

local_device ​

function View source

python
def local_device() -> torch.device: ...

Return rank-local device metadata selected for the current process.

init selects the CUDA device identified by LOCAL_RANK. The helper reads the current CUDA device when CUDA is available and otherwise returns CPU; it is available before process-group initialization.

Returns

The current CUDA device when CUDA is available, otherwise the CPU device. World size does not affect the result.

new_group ​

re-export

python
new_group = torch.distributed.new_group

recv ​

re-export

python
recv = torch.distributed.recv

reduce ​

re-export

python
reduce = torch.distributed.reduce

reduce_ciphertext ​

function View source

python
def reduce_ciphertext(value: Ciphertext, *, dst: int=0, engine: CkksEngine, group: torch.distributed.ProcessGroup | None=None) -> None: ...

Synchronously sum ciphertext partials onto dst using add.

A binomial tree combines complete ciphertexts and works for every positive process-group size, not only powers of two. The tree uses P - 1 full ciphertext messages, O(log P) critical-path rounds, and one temporary receive buffer per active rank. Non-destination values are partial or unchanged after return and must not be interpreted as the global sum. This is component-wise CKKS ciphertext addition modulo the active Q primes, not gathering independent values or concatenating RNS limbs. The operation accepts no async_op argument and returns no torch.distributed.Work.

Parameters

  • value: Rank-local additive ciphertext partial. The object on dst is updated in place; intermediate receivers may also be mutated as the tree accumulates subtrees.
  • dst: Global rank that receives the complete sum, which must belong to group. This is not a process-group-relative rank.
  • engine: Rank-local CKKS engine used for in-place modular ciphertext addition. Every participating rank must provide an engine compatible with its value.
  • group: Participating process group. None selects the default process group.

Returns

None. Only value on dst is guaranteed to contain the global sum. For world size one, the sole value is already the sum and remains unchanged.

Raises

  • ValueError: If dst is outside group; rank-local values are not ciphertexts with identical shape and arithmetic metadata; or the engine rejects an incompatible value.
  • RuntimeError: If distributed communication is uninitialized or the caller is not a member of group.

reduce_scatter ​

re-export

python
reduce_scatter = torch.distributed.reduce_scatter

reduce_scatter_single ​

data View source

python
reduce_scatter_single = getattr(_torch_distributed, 'reduce_scatter_single', _torch_distributed.reduce_scatter_tensor)

scatter ​

re-export

python
scatter = torch.distributed.scatter

scatter_ciphertext_limbs ​

function View source

python
def scatter_ciphertext_limbs(shards_or_none: Sequence[Ciphertext] | None, *, src: int=0, group: torch.distributed.ProcessGroup | None=None) -> Ciphertext: ...

Scatter RNS limb shards of one logical ciphertext from one source.

The source sequence must partition one ciphertext into nonempty, contiguous prime intervals in process-group-rank order. This function distributes those existing shards; it neither performs ciphertext arithmetic nor reconstructs the full ciphertext. It is synchronous, accepts no async_op argument, and returns no torch.distributed.Work.

Parameters

  • shards_or_none: On src, exactly one caller-prepared ciphertext limb shard per group rank, ordered by process-group rank and increasing canonical prime interval. Every non-source rank must pass None.
  • src: Global rank of the source process, which must belong to group.
  • group: Participating process group. None selects the default process group.

Returns

The shard assigned to the caller's process-group rank. The source receives its existing sequence element; other ranks receive newly allocated shards. For world size one, the sole validated shard is returned unchanged without communication or allocation.

Raises

  • ValueError: If src is outside group; the source/non-source calling rules or sequence-length requirements are violated; or shards do not describe compatible, contiguous prime intervals of one ciphertext.
  • RuntimeError: If distributed communication is uninitialized or the caller is not a member of group.

scatter_ciphertexts ​

function View source

python
def scatter_ciphertexts(values_or_none: Sequence[Ciphertext] | None, *, src: int=0, group: torch.distributed.ProcessGroup | None=None) -> Ciphertext: ...

Scatter independent ciphertext workload items from one source.

This is representation-preserving transport, not an arithmetic operation. Sequence position is process-group-rank order even though src is a global rank. The operation is synchronous, accepts no async_op argument, and returns no torch.distributed.Work.

Parameters

  • values_or_none: On src, one independent ciphertext per group rank in process-group-rank order. Every non-source rank must pass None.
  • src: Global rank of the source process, which must belong to group.
  • group: Participating process group. None selects the default process group.

Returns

The ciphertext assigned to the caller's process-group rank. The source receives its existing sequence element; other ranks receive newly allocated ciphertexts. For world size one, values_or_none[0] is returned unchanged without communication or allocation.

Raises

  • ValueError: If src is outside group; the source does not provide exactly one ciphertext per group rank; a non-source rank supplies values; or rank-local arguments disagree.
  • RuntimeError: If distributed communication is uninitialized or the caller is not a member of group.

send ​

re-export

python
send = torch.distributed.send

shutdown ​

function View source

python
def shutdown(group: torch.distributed.ProcessGroup | None=None) -> None: ...

Destroy an initialized process group.

Parameters

  • group: Process group to destroy. None selects the PyTorch default process group; this is a group object, not a rank identifier.

Returns

None. An unavailable or uninitialized distributed runtime, including a direct process before init, is a no-op. A world-size-one group is destroyed in the same way as any other initialized group. No asynchronous torch.distributed.Work handle is returned.

Towards full-stack encrypted execution infrastructures

Released under the MIT License. Copyright VisualDust and FHElium contributors