Homogeneous batching

Example source: examples/15_homogeneous_batching.py

This example compares one homogeneous CKKS batch with an explicit loop over the same packed matrix-vector workload. The tutorial explains leading batch dimensions and keeps the latency-memory execution choice in application code.

Run the example

Start with the 8,192-slot, 40-bit-scale baseline:

python examples/15_homogeneous_batching.py \
  --preset slots8192-scale40-levels7-int64 \
  --level 0 \
  --batch-sizes 1,4,8 \
  --warmup 2 \
  --runs 10

Then compare the two ends of the 32,768-slot, 40-bit-scale chain:

python examples/15_homogeneous_batching.py \
  --preset slots32768-scale40-levels34-int64 --level 0 --batch-sizes 1,4,8

python examples/15_homogeneous_batching.py \
  --preset slots32768-scale40-levels34-int64 --level 30 --batch-sizes 1,4,8

The second command is not redundant. Level changes the active RNS row count, which changes both arithmetic work and the size of each NTT/key-switch working set.

1. A leading prefix is a value batch

The example constructs B independent vectors:

vectors.shape == (B, size)
messages = vectors.repeat(1, engine.num_slots // size)
source = engine.encrypt_message(messages, level=level)
assert tuple(source.batch_shape) == (B,)

For a ciphertext, the complete data layout is conceptually:

[component, *batch, limb, coefficient]

For an RNS plaintext it is:

[*batch, limb, coefficient_or_ntt_index]

The leading dimensions are semantic message dimensions. They are not RNS limbs, polynomial components, distributed ranks, or hybrid-decomposition digits. All members of one homogeneous value share its context, level, scale, polynomial domain, modulus basis, device, dtype, and component count.

They must also have the same effective encryption-key lineage. A context id describes parameters, not a particular secret key, so the engine cannot infer that independently produced ciphertexts are safe to stack. This example encrypts the complete message batch with one engine/key. When assembling existing ciphertexts, key-switch them when necessary before calling Ciphertext.stack_batch.

The matrix diagonals are stacked along a new term axis inside the evaluator. Any pre-existing ciphertext batch axes remain inner dimensions, so one public matrix is applied to every encrypted vector.

2. The evaluator is batch-polymorphic

matrix_vector contains no batch-specific branch:

def matrix_vector(source, *, engine, diagonals, rotation_keys):
    rotated_values = []
    for step in range(len(diagonals)):
        rotated = (
            source
            if step == 0
            else engine.rotate_with_key(source, rotation_keys[step])
        )
        rotated_values.append(rotated)
    rotated_ntt = engine.coefficient_domain_to_ntt_domain(
        Ciphertext.stack_batch(rotated_values)
    )
    weighted = engine.multiply_plaintext(
        rotated_ntt, Plaintext.stack_batch(diagonals)
    )
    return engine.rescale_to_next_level(
        engine.ntt_domain_to_coefficient_domain(
            engine.sum_ciphertext_batch(weighted)
        )
    )

The same function accepts either an unbatched ciphertext or one with a non-empty batch prefix. batch=1 therefore uses the same public program rather than a separate compatibility path. The outer term axis batches the forward NTT, plaintext products, and binary-tree reduction, avoiding one under-filled native launch per matrix diagonal.

3. Choose batch or loop execution

The example evaluates equivalent logical work in two ways.

Batched submission:

batched_result = matrix_vector(source, ...)

Explicit loop:

individual_sources = tuple(
    value.clone() for value in source.unbind_batch()
)
looped_result = [matrix_vector(value, ...) for value in individual_sources]

unbind_batch makes the logical members visible. The example clones the views so the loop represents independently owned request values. It then uses Ciphertext.stack_batch(looped_result) only to verify exact equivalence:

torch.testing.assert_close(
    batched_result.data,
    Ciphertext.stack_batch(looped_result).data,
    rtol=0,
    atol=0,
)

stack_batch allocates and copies; it is not a hidden performance shortcut.

Batch-versus-loop selection is a programmer or workload-scheduler decision based on deployment measurements, latency requirements, and memory budget.

4. Compare the complete workload

The example uses the cyclic-diagonal formulation of an 8-by-8 matrix-vector product:

$$y=\sum _s{p}_s\odot \mathrm{Rot}(x,s).$$

This workload composes:

• shared plaintext broadcasting;
• rotations and direct key switching;
• plaintext multiplication;
• rescale;
• ciphertext accumulation.

It is more informative than timing only one elementwise operator. Before timing, the example also decrypts the batched result and compares it with vectors @ matrix.T.

5. Read both latency and memory columns

For every requested B, the example reports:

• the aggregate size of one extended QP digit;
• synchronized batched and loop medians;
• loop / batch speedup;
• the faster path at that measured point;
• incremental peak allocated CUDA memory when running on CUDA (n/a on CPU);
• maximum absolute decryption error.

A ratio above one means the batch won. It does not imply that a larger batch will continue to scale. A batch can reduce launches and expose parallelism while simultaneously enlarging active NTT, automorphism, accumulator, and output tensors beyond the effective cache working set.

The faster column describes only the current command. The example does not write an automatic recommendation into engine configuration or cache a hidden device policy.

6. Select a policy from the deployed point

Use these rules as a measurement plan, not as hard-coded library behavior:

1. Verify the batched result exactly against an explicit loop from the same installed build.
2. Compare [1, slots] with [slots] to isolate B1 overhead.
3. Compare B4/B8 with an explicit loop over the same members at the same preset and level.
4. Repeat at the levels used by the real evaluator.
5. Measure the complete workload and peak memory, not only an NTT kernel.
6. Keep the resulting choice in application or scheduler code.

The worked RTX PRO 6000 measurements and the working-set explanation are in Choose a homogeneous batch size. The stable mechanism is summarized in the CKKS workload cost model.

Complete runnable source

#!/usr/bin/env python3

"""Compare one homogeneous CKKS batch with an explicit per-message loop.

The evaluator computes a packed ``y = A @ x`` with the cyclic-diagonal
method.  FHElium preserves every leading message dimension as a homogeneous
batch prefix; the application still chooses whether to submit that batch or
to evaluate its members one at a time.

Example:
    python examples/15_homogeneous_batching.py \
        --preset slots8192-scale40-levels7-int64 --level 0 --batch-sizes 1,4,8
"""

from __future__ import annotations

import argparse
import gc
import statistics as stats
import time
from collections.abc import Callable
from typing import Any

import torch
from common import (
    add_engine_args,
    error_stats,
    format_bytes,
    make_engine,
    print_table,
    sync_if_cuda,
)

from fhelium import Ciphertext, CkksEngine, Plaintext, RotationKey


def _parse_batch_sizes(text: str) -> list[int]:
    batch_sizes = [
        int(item.strip()) for item in text.split(',') if item.strip()
    ]
    if not batch_sizes or any(size <= 0 for size in batch_sizes):
        raise argparse.ArgumentTypeError(
            "batch sizes must be a comma-separated list of positive integers"
        )
    return batch_sizes


def matrix_and_vectors(
    size: int,
    batch_size: int,
    *,
    seed: 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)
    generator = torch.Generator().manual_seed(seed)
    vectors = (
        torch.randn(
            (batch_size, size),
            generator=generator,
            dtype=torch.float64,
        )
        * 0.025
    )
    return matrix, vectors


def periodic_slots(values: torch.Tensor, num_slots: int) -> torch.Tensor:
    return values.repeat(1, num_slots // values.size(-1))


def cyclic_diagonal_slots(
    matrix: torch.Tensor,
    rotation_step: int,
    num_slots: int,
) -> torch.Tensor:
    row = torch.arange(num_slots) % matrix.size(0)
    column = torch.remainder(row - rotation_step, matrix.size(0))
    return matrix[row, column]


def prepare_constants(
    engine: CkksEngine,
    matrix: torch.Tensor,
    *,
    level: int,
) -> tuple[list[Plaintext], dict[int, RotationKey]]:
    diagonals = [
        engine.prepare_plaintext_for_multiplication(
            engine.encode(
                cyclic_diagonal_slots(matrix, step, engine.num_slots),
                level=level,
            )
        )
        for step in range(matrix.size(0))
    ]
    rotation_keys = {
        step: engine.rotation_key(step) for step in range(1, matrix.size(0))
    }
    return diagonals, rotation_keys


def matrix_vector(
    source: Ciphertext,
    *,
    engine: CkksEngine,
    diagonals: list[Plaintext],
    rotation_keys: dict[int, RotationKey],
) -> Ciphertext:
    """Evaluate the same program for an unbatched or batched ciphertext."""

    rotated_values = []
    for step in range(len(diagonals)):
        rotated = (
            source
            if step == 0
            else engine.rotate_with_key(source, rotation_keys[step])
        )
        rotated_values.append(rotated)
    rotated_ntt = engine.coefficient_domain_to_ntt_domain(
        Ciphertext.stack_batch(rotated_values)
    )
    diagonal_batch = Plaintext.stack_batch(diagonals)
    if source.batch_shape:
        assert diagonal_batch.data is not None
        expanded = (
            diagonal_batch.data.reshape(
                diagonal_batch.batch_shape
                + (1,) * len(source.batch_shape)
                + diagonal_batch.data.shape[-2:]
            )
            .expand(
                diagonal_batch.batch_shape
                + source.batch_shape
                + diagonal_batch.data.shape[-2:]
            )
            .contiguous()
        )
        diagonal_batch = Plaintext(
            message=None,
            level=diagonal_batch.level,
            scale=diagonal_batch.scale,
            data=expanded,
            context_id=diagonal_batch.context_id,
            representation=diagonal_batch.representation,
            polynomial_domain=diagonal_batch.polynomial_domain,
            modulus_basis=diagonal_batch.modulus_basis,
            residue_representation=diagonal_batch.residue_representation,
            prime_ids=diagonal_batch.prime_ids,
        )
    weighted = engine.multiply_plaintext(rotated_ntt, diagonal_batch)
    return engine.rescale_to_next_level(
        engine.ntt_domain_to_coefficient_domain(
            engine.sum_ciphertext_batch(weighted)
        )
    )


def _time_call(fn: Callable[[], Any], device: torch.device) -> float:
    sync_if_cuda(device)
    start = time.perf_counter()
    result = fn()
    sync_if_cuda(device)
    elapsed_ms = (time.perf_counter() - start) * 1e3
    del result
    return elapsed_ms


def alternating_benchmark(
    batched: Callable[[], Any],
    looped: Callable[[], Any],
    *,
    warmup: int,
    runs: int,
    device: torch.device,
) -> tuple[float, float]:
    """Return paired medians while alternating first-run order."""

    for _ in range(warmup):
        _time_call(batched, device)
        _time_call(looped, device)

    batched_times = []
    looped_times = []
    for run_index in range(runs):
        if run_index % 2 == 0:
            batched_times.append(_time_call(batched, device))
            looped_times.append(_time_call(looped, device))
        else:
            looped_times.append(_time_call(looped, device))
            batched_times.append(_time_call(batched, device))
    return stats.median(batched_times), stats.median(looped_times)


def peak_allocated_mib(
    fn: Callable[[], Any],
    *,
    device: torch.device,
) -> float | None:
    if device.type != "cuda":
        return None
    gc.collect()
    torch.cuda.empty_cache()
    sync_if_cuda(device)
    baseline = torch.cuda.memory_allocated(device)
    torch.cuda.reset_peak_memory_stats(device)
    result = fn()
    sync_if_cuda(device)
    peak = torch.cuda.max_memory_allocated(device) - baseline
    del result
    gc.collect()
    torch.cuda.empty_cache()
    return peak / (1 << 20)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    add_engine_args(parser, default_preset="slots8192-scale40-levels7-int64")
    parser.add_argument("--level", type=int, default=0)
    parser.add_argument("--size", type=int, default=8)
    parser.add_argument(
        "--batch-sizes",
        type=_parse_batch_sizes,
        default=[1, 4, 8],
    )
    parser.add_argument("--warmup", type=int, default=2)
    parser.add_argument("--runs", type=int, default=10)
    args = parser.parse_args()

    engine = make_engine(args)
    if engine.device.type == "cuda":
        # The shared synchronization helper and allocator statistics follow
        # the process-current CUDA device.
        torch.cuda.set_device(engine.device)
    if args.size <= 0 or engine.num_slots % args.size != 0:
        parser.error(
            f"--size must be positive and divide num_slots={engine.num_slots}"
        )
    if not 0 <= args.level < engine.config.num_q_primes - 1:
        parser.error(
            "--level must leave at least one Q prime for the workload's "
            f"rescale; got {args.level} with "
            f"{engine.config.num_q_primes} Q primes"
        )
    if args.warmup < 0 or args.runs <= 0:
        parser.error("--warmup must be non-negative and --runs positive")

    matrix, _ = matrix_and_vectors(args.size, 1, seed=700)
    diagonals, rotation_keys = prepare_constants(
        engine,
        matrix,
        level=args.level,
    )

    active_q_rows = engine.config.num_q_primes - args.level
    active_qp_rows = active_q_rows + engine.config.num_p_primes
    qp_digit_bytes = active_qp_rows * engine.config.N * torch.int64.itemsize
    print(
        "Homogeneous batching benchmark\n"
        f"  preset={args.preset}, level={args.level}, "
        f"device={args.device}\n"
        f"  active rows: Q={active_q_rows}, QP={active_qp_rows}\n"
        "  one extended QP digit per message: "
        f"{format_bytes(qp_digit_bytes)}\n"
        "  policy: the application executes and compares and compares both "
        "paths"
    )

    rows = []
    for batch_size in args.batch_sizes:
        _, vectors = matrix_and_vectors(
            args.size,
            batch_size,
            seed=700 + batch_size,
        )
        source = engine.encrypt_message(
            periodic_slots(vectors, engine.num_slots),
            level=args.level,
        )
        # unbind_batch returns views. Clone here so the loop represents
        # independently owned, ordinary unbatched request values.
        individual_sources = tuple(
            value.clone() for value in source.unbind_batch()
        )

        def batched() -> Ciphertext:
            return matrix_vector(
                source,
                engine=engine,
                diagonals=diagonals,
                rotation_keys=rotation_keys,
            )

        def looped() -> list[Ciphertext]:
            return [
                matrix_vector(
                    value,
                    engine=engine,
                    diagonals=diagonals,
                    rotation_keys=rotation_keys,
                )
                for value in individual_sources
            ]

        batched_result = batched()
        looped_result = looped()
        stacked_loop = Ciphertext.stack_batch(looped_result)
        torch.testing.assert_close(
            batched_result.data,
            stacked_loop.data,
            rtol=0,
            atol=0,
        )
        actual = engine.decrypt_message(batched_result, is_real=True)[
            ..., : args.size
        ]
        error = error_stats(actual, vectors @ matrix.T)
        del batched_result, looped_result, stacked_loop, actual

        batched_ms, looped_ms = alternating_benchmark(
            batched,
            looped,
            warmup=args.warmup,
            runs=args.runs,
            device=engine.device,
        )
        speedup = looped_ms / batched_ms
        faster_path = "batch" if speedup >= 1.0 else "loop"
        batch_peak = peak_allocated_mib(batched, device=engine.device)
        loop_peak = peak_allocated_mib(looped, device=engine.device)
        rows.append(
            [
                batch_size,
                format_bytes(qp_digit_bytes * batch_size),
                f"{batched_ms:.4f}",
                f"{looped_ms:.4f}",
                f"{speedup:.3f}x",
                faster_path,
                "n/a" if batch_peak is None else f"{batch_peak:.1f}",
                "n/a" if loop_peak is None else f"{loop_peak:.1f}",
                f"{error['max_abs']:.3e}",
            ]
        )

    print_table(
        [
            "B",
            "B x QP digit",
            "batch ms",
            "loop ms",
            "loop/batch",
            "faster",
            "batch peak MiB",
            "loop peak MiB",
            "max abs error",
        ],
        rows,
    )
    print(
        "\nThe faster column describes only this measured point. FHElium "
        "does not select or cache an execution policy; keep the choice in "
        "application code and remeasure the deployed preset, level, batch "
        "size, device, and complete workload."
    )


if __name__ == "__main__":
    main()

Related concepts and guides

• Values and state
• CKKS workload cost model
• Choose a homogeneous batch size
• Benchmark a workload correctly
• CUDA Graph matrix-vector