FHElium

Appearance

Late relinearization and NTT reuse ​

Example source: examples/05_eager_ntt_reuse.py

This example accumulates three-component products before relinearization and reuses fixed multiplication operands in NTT form. The tutorial explains the state preconditions that make both optimizations valid.

Run the example ​

bash
python examples/05_eager_ntt_reuse.py \
  --preset slots8192-scale40-depth7-int64 \
  --pair-count 3

1. Prepare multiplication operands ​

python
multiplicand_ntt = engine.coefficient_domain_to_ntt_domain(engine.encrypt_message(multiplicand))
multiplier_ntt = engine.coefficient_domain_to_ntt_domain(engine.encrypt_message(multiplier))

fhelium.eager.Engine.multiply has these fixed preconditions:

  • both inputs have two components;
  • both inputs are in the same NTT/Montgomery representation;
  • depth, scale, basis, and active prime IDs are compatible;
  • the caller established compatible CKKS parameter provenance;
  • the result is a three-component NTT ciphertext;
  • no implicit relinearization or rescale occurs.

2. Accumulate three-component products ​

python
product = engine.multiply(multiplicand_ntt, multiplier_ntt)
accumulator = (
    product
    if accumulator is None
    else engine.add(accumulator, product)
)

Compatible three-component products can be added before relinearization. This turns a sum of products from:

Expanded diagram
100%

into:

Expanded diagram
100%

The optimization is valid only while all terms share a matching layout and scale. An intervening operation that requires an ordinary two-component ciphertext creates a point at which relinearization becomes necessary.

3. Relinearize and rescale once ​

python
output = engine.rescale_to_next_depth(engine.relinearize(accumulator))

Relinearization key-switches the c2 contribution back into two ciphertext components. It is usually much more expensive than an elementwise modular addition, so reducing its count is useful for dot products, matrix methods, and polynomial schedules. The accumulated product still carries scale Δ2; the rescale consumes one depth after relinearization.

4. Keep reusable operands in NTT form ​

python
fixed = engine.coefficient_domain_to_ntt_domain(engine.encrypt_message(fixed_values))
source = engine.coefficient_domain_to_ntt_domain(engine.encrypt_message(source_values))

product = engine.multiply(source, fixed)

The example transforms the fixed operand once and reuses it in three products with independently encrypted sources. It accumulates those products in NTT form before one relinearization and rescale.

The application must still account for:

  • the memory cost of retaining the prepared operand;
  • its depth and scale;
  • whether consumers mutate it;
  • whether later operations require coefficient-domain form.

5. Inspect state instead of assuming it ​

The example prints:

text
component count, polynomial domain, residue representation, and scale

Use those fields when debugging a schedule. A tensor with the expected shape but the wrong domain or Montgomery representation is not a compatible operand.

For represented programs, the Compile stack can insert relinearization automatically. Compose and execute a pipeline from built-in Compile passes shows InsertRelinearizationPass placing the operation in a transformed Program.

Source
py
#!/usr/bin/env python3

"""Explicit late-relinearization and NTT-reuse workflows."""

from __future__ import annotations

import argparse
from typing import cast

import torch
from common import (
    add_engine_args,
    error_stats,
    make_engine,
    print_table,
    small_complex_vector,
)

import fhelium as fh
from fhelium.eager import Engine


def _state(ct: fh.Ciphertext) -> str:
    return (
        f"components={ct.component_count},polynomial_domain={ct.polynomial_domain},"
        f"residue_representation={ct.residue_representation},scale={ct.scale:.3e}"
    )


def late_relinearization(
    engine: Engine, pair_count: int
) -> tuple[fh.Ciphertext, torch.Tensor]:
    accumulator: fh.Ciphertext | None = None
    reference: torch.Tensor | None = None
    for index in range(pair_count):
        multiplicand = small_complex_vector(
            engine.num_slots, seed=100 + index, scale=0.005
        )
        multiplier = small_complex_vector(
            engine.num_slots, seed=200 + index, scale=0.005
        )
        multiplicand_ntt = cast(
            fh.Ciphertext,
            engine.coefficient_domain_to_ntt_domain(
                engine.encrypt_message(multiplicand)
            ),
        )
        multiplier_ntt = cast(
            fh.Ciphertext,
            engine.coefficient_domain_to_ntt_domain(
                engine.encrypt_message(multiplier)
            ),
        )
        product = engine.multiply(multiplicand_ntt, multiplier_ntt)
        accumulator = (
            product if accumulator is None else engine.add(accumulator, product)
        )
        reference = (
            multiplicand * multiplier
            if reference is None
            else reference + multiplicand * multiplier
        )

    if accumulator is None or reference is None:
        raise ValueError("pair_count must be positive")
    output = engine.rescale_to_next_depth(engine.relinearize(accumulator))
    return output, reference


def ntt_reuse(
    engine: Engine,
) -> tuple[fh.Ciphertext, torch.Tensor]:
    fixed_values = small_complex_vector(engine.num_slots, seed=300, scale=0.005)
    fixed = cast(
        fh.Ciphertext,
        engine.coefficient_domain_to_ntt_domain(
            engine.encrypt_message(fixed_values)
        ),
    )
    accumulator = None
    reference = torch.zeros_like(fixed_values)
    for index in range(3):
        source_values = small_complex_vector(
            engine.num_slots, seed=301 + index, scale=0.005
        )
        source = engine.coefficient_domain_to_ntt_domain(
            engine.encrypt_message(source_values)
        )
        # The same transformed fixed operand participates in every product.
        product = engine.multiply(source, fixed)
        accumulator = (
            product if accumulator is None else engine.add(accumulator, product)
        )
        reference += source_values * fixed_values
    assert accumulator is not None
    output = engine.rescale_to_next_depth(engine.relinearize(accumulator))
    return output, reference


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    add_engine_args(parser)
    parser.add_argument("--pair-count", type=int, default=3)
    args = parser.parse_args()
    engine = make_engine(args)

    late, late_reference = late_relinearization(engine, args.pair_count)
    reuse, reuse_reference = ntt_reuse(engine)
    rows = []
    for name, value, reference in [
        ("late relinearization", late, late_reference),
        ("NTT operand reuse", reuse, reuse_reference),
    ]:
        error = error_stats(
            engine.decrypt_message(value).cpu(),
            reference,
            engine.num_slots,
        )
        rows.append(
            [name, value.depth, _state(value), f"{error['max_abs']:.3e}"]
        )
    print_table(["demo", "depth", "state", "max abs error"], rows)


if __name__ == "__main__":
    main()