Rotation-parallel matrix-vector multiplication

Example source: examples/09_spmd_rotation_parallel_mxv.py

This example partitions the cyclic-diagonal terms and exact rotation keys of one packed matrix-vector product across ranks, then reduces additive ciphertext partials. The tutorial explains term ownership, key movement, and the final reduction step.

Run on two GPUs

torchrun --standalone --nproc-per-node=2 \
  examples/09_spmd_rotation_parallel_mxv.py \
  --size 8

The matrix size must divide the CKKS slot count, and the world size must not exceed the matrix size.

1. Express matrix-vector multiplication with cyclic diagonals

For a packed vector x, the example evaluates:

$$\mathbf{y}=\sum _{s=0}^{n-1}{\mathbf{d}}_s\odot \mathrm{Rot}(\mathbf{x},s),$$

where ${\mathbf{d}}_s$ is the cyclic diagonal aligned with rotation step $s$, and $\odot$ denotes slot-wise multiplication.

The diagonal aligned with step s is built by:

row = torch.arange(num_slots) % size
column = torch.remainder(row - rotation_step, size)
diagonal = matrix[row, column]

Repeating the input vector periodically across all slots lets each slot block evaluate the same small matrix-vector problem.

2. Replicate one encrypted input

source = dist.broadcast_ciphertext(root_source, src=0)

Every rank needs the complete source because each rank rotates it by a different subset of steps. This is one logical ciphertext replicated for parallel evaluation, not independent data-parallel samples.

3. Partition terms by rotation step

local_rotation_steps = list(
    range(dist.get_rank(), size, dist.get_world_size())
)

For two ranks and size eight:

rank 0: steps 0, 2, 4, 6
rank 1: steps 1, 3, 5, 7

The cyclic assignment balances the number of diagonal terms. A production algorithm may choose another owner function based on measured cost or key locality.

4. Move only the exact keys each owner retains

Rank zero creates each key:

source_key = engine.create_rotation_key(rotation_step, secret_key)

All ranks participate in the typed broadcast when the owner is remote, but only that owner retains the transferred object:

transferred_key = dist.broadcast_key(source_key, src=0)
if dist.get_rank() == owner:
    local_keys[rotation_step] = transferred_key
else:
    del transferred_key

The secret key never leaves rank zero. FHElium does not infer step ownership or key placement.

5. Evaluate local diagonal terms

rotated = (
    source.clone()
    if rotation_step == 0
    else engine.rotate_with_key(source, local_keys[rotation_step])
)
diagonal = engine.prepare_plaintext_for_multiplication(
    engine.encode(diagonal_slots, level=rotated.level)
)
term = engine.rescale_to_next_level(
    engine.ntt_domain_to_coefficient_domain(
        engine.multiply_plaintext(
            engine.coefficient_domain_to_ntt_domain(rotated), diagonal
        )
    )
)

All local terms reach the same level and scale, so they can be summed with engine.sum_ciphertexts.

6. Reduce additive partials

local_partial = engine.sum_ciphertexts(local_terms)
dist.reduce_ciphertext(local_partial, dst=0, engine=engine)

Unlike the independent-ciphertext example, rank identity is no longer part of the result. Each local ciphertext represents a partial sum of the same mathematical output, so a typed ciphertext addition reduction is correct.

Expanded diagram

100%

Raw integer reduction of ciphertext.data is not a substitute for the typed collective because receiver construction and CKKS compatibility remain part of the operation.

Capacity implication

This partition can reduce per-rank rotation-key residency, but the input is replicated and every rank produces local terms. Measure:

• aggregate and maximum-rank key bytes;
• input replication;
• diagonal plaintext residency;
• reduction communication;
• load balance across step owners.

Complete runnable source

#!/usr/bin/env python3
"""Rotation-parallel packed matrix-vector multiplication.

Run on one process or one process per GPU:

    python examples/09_spmd_rotation_parallel_mxv.py

    torchrun --standalone --nproc-per-node=2 \
        examples/09_spmd_rotation_parallel_mxv.py --size 8

The packed input is replicated, while cyclic matrix diagonals and their
rotation keys are partitioned by rotation step. Each rank computes an
additive partial ciphertext. ``reduce_ciphertext`` combines those partials
with CKKS modular addition and leaves the final result only on rank 0.

Use this pattern for rotation/diagonal/head partitions whose rank-local
outputs are contributions to one logical encrypted result. The secret key
never leaves rank 0; only selected rotation keys are communicated.
"""

from __future__ import annotations

import argparse

import torch

import fhelium as fh
import fhelium.distributed as dist


def _matrix_and_vector(size: int) -> tuple[torch.Tensor, torch.Tensor]:
    row = torch.arange(size, dtype=torch.float64).view(-1, 1)
    column = torch.arange(size, dtype=torch.float64).view(1, -1)
    matrix = 0.018 * torch.sin((row + 1) * (column + 2) * 0.17)
    matrix += 0.007 * torch.cos((row + column + 1) * 0.23)
    vector = 0.025 * torch.cos(torch.arange(size, dtype=torch.float64) * 0.31)
    vector -= 0.009 * torch.sin(torch.arange(size, dtype=torch.float64) * 0.19)
    return matrix, vector


def _periodic_slots(values: torch.Tensor, num_slots: int) -> torch.Tensor:
    if num_slots % values.numel() != 0:
        raise ValueError(
            f"matrix size {values.numel()} must divide num_slots={num_slots}"
        )
    return values.repeat(num_slots // values.numel())


def _cyclic_diagonal_slots(
    matrix: torch.Tensor,
    rotation_step: int,
    num_slots: int,
) -> torch.Tensor:
    """Return weights aligned with ``torch.roll(x, shifts=rotation_step)``."""

    size = matrix.size(0)
    row = torch.arange(num_slots) % size
    column = torch.remainder(row - rotation_step, size)
    return matrix[row, column]


def _provision_owned_rotation_keys(
    engine: fh.CkksEngine,
    secret_key: fh.SecretKey | None,
    size: int,
) -> dict[int, fh.RotationKey]:
    """Create each exact key on rank 0 and retain it only on its owner."""

    local_keys: dict[int, fh.RotationKey] = {}
    for rotation_step in range(1, size):
        owner = rotation_step % dist.get_world_size()
        source_key = None
        if dist.get_rank() == 0:
            assert secret_key is not None
            source_key = engine.create_rotation_key(rotation_step, secret_key)

        if owner == 0:
            if dist.get_rank() == 0:
                assert source_key is not None
                local_keys[rotation_step] = source_key
            continue

        # All ranks participate in the typed key broadcast, but only the
        # designated owner retains the transferred object.
        transferred_key = dist.broadcast_key(source_key, src=0)
        if dist.get_rank() == owner:
            local_keys[rotation_step] = transferred_key
        else:
            del transferred_key
    return local_keys


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--size", type=int, default=8)
    args = parser.parse_args()

    dist.init()
    engine = fh.CkksEngine(
        fh.Preset.slots32768_scale40_levels34_int64,
        device=dist.local_device(),
        allow_sk_gen=False,
    )
    if args.size <= 0 or args.size > engine.num_slots:
        raise ValueError(f"size must be in [1, {engine.num_slots}]")
    if engine.num_slots % args.size != 0:
        raise ValueError("size must divide the CKKS slot count")
    if dist.get_world_size() > args.size:
        raise ValueError(
            f"world_size={dist.get_world_size()} exceeds size={args.size}"
        )

    matrix, vector = _matrix_and_vector(args.size)
    if dist.get_rank() == 0:
        secret_key = engine.create_secret_key()
        public_key = engine.create_public_key(secret_key)
        root_source = engine.encrypt_message(
            _periodic_slots(vector, engine.num_slots),
            public_key,
        )
    else:
        secret_key = None
        root_source = None

    # One logical encrypted input is intentionally replicated because every
    # rank evaluates a different subset of its rotations.
    source = dist.broadcast_ciphertext(root_source, src=0)
    local_rotation_steps = list(
        range(dist.get_rank(), args.size, dist.get_world_size())
    )
    local_keys = _provision_owned_rotation_keys(
        engine,
        secret_key,
        args.size,
    )

    local_terms = []
    for rotation_step in local_rotation_steps:
        rotated = (
            source.clone()
            if rotation_step == 0
            else engine.rotate_with_key(source, local_keys[rotation_step])
        )
        diagonal = engine.prepare_plaintext_for_multiplication(
            engine.encode(
                _cyclic_diagonal_slots(matrix, rotation_step, engine.num_slots),
                level=rotated.level,
            )
        )
        local_terms.append(
            engine.rescale_to_next_level(
                engine.ntt_domain_to_coefficient_domain(
                    engine.multiply_plaintext(
                        engine.coefficient_domain_to_ntt_domain(rotated),
                        diagonal,
                    )
                )
            )
        )

    local_partial = engine.sum_ciphertexts(local_terms)

    # These are additive contributions to one result, so reduction is the
    # correct operation. Gathering would retain unnecessary per-rank objects.
    dist.reduce_ciphertext(local_partial, dst=0, engine=engine)
    print(
        f"rank={dist.get_rank()} rotation_steps={local_rotation_steps} "
        f"retained_rotation_keys={sorted(local_keys)}"
    )

    if dist.get_rank() == 0:
        assert secret_key is not None
        decoded = engine.decrypt_message(
            local_partial,
            secret_key=secret_key,
            is_real=True,
        )[: args.size]
        expected = matrix @ vector
        max_error = float(torch.max(torch.abs(decoded - expected)))
        torch.testing.assert_close(decoded, expected, atol=3e-5, rtol=0)
        print(
            "rotation_parallel_mxv_ok "
            f"size={args.size} world_size={dist.get_world_size()} "
            f"max_abs_error={max_error:.3e}"
        )

    dist.shutdown()


if __name__ == "__main__":
    main()

Related concepts and guides

• Communication semantics
• Key lifecycle
• Choose a multi-GPU partition