Compose and execute a pipeline from built-in Compile passes

Example source: examples/12_compile_pipeline.py

Example 12 is the ordinary end-to-end Compile workflow. It captures a supported PyTorch computation, composes existing FHElium passes, lowers the Program to operations covered by the Backend, links live keys and arithmetic resources, executes encrypted inputs, and checks the decrypted result.

Run the example

From the repository root:

python examples/12_compile_pipeline.py

The example runs on CPU and prints compact captured/backend-ready Program inventories, evaluation-key requirements, result state, numerical error, and pass reports.

Source computation

The source adds a vector to one cyclic rotation and squares the sum:

def rotated_quadratic(x, rotation):
    mixed = x + torch.roll(x, shifts=rotation, dims=-1)
    return mixed * mixed

Capture assigns x the encrypted role and treats rotation as a compile-time integer. The logical 16-slot input is tiled across the complete CKKS slot ring before encryption, so a full-ring rotation implements the same periodic layout used by the clear computation.

Caller-selected pipeline

Pipelines express workload-specific choices. This example composes the built-in passes directly rather than using the default_lower_and_fuse_pipeline recipe:

from fhelium import compile as fh_compile

pipeline = fh_compile.Pipeline(
    (
        fh_compile.EliminateDeadValuesPass(),
        fh_compile.LowerSemanticToLogicalPass(),
        fh_compile.InsertMultiplyNttTransitionsPass(),
        fh_compile.LowerLogicalToCkksPass(),
        fh_compile.ResolveRotationKeyOperandsPass(),
        fh_compile.InsertRelinearizationPass(),
        fh_compile.InsertRescalePass(),
        fh_compile.AssignCkksDepthsPass(entry_depth=0),
        fh_compile.AssignCkksScalesPass(
            entry_scale=config.default_scale,
        ),
        fh_compile.AssignNttImplementationPass("native-ntt", ntt_backend="radix2_indexed"),
        fh_compile.LowerCkksToRnsNttPass(),
    )
)
compiled = pipeline.run(captured)

This sequence records the following caller-selected pipeline choices:

• it selects the CPU indexed NTT schedule;
• it inserts immediate relinearization and rescale operations for the ciphertext product;
• it assigns concrete depths and per-value actual scales after transition placement;
• it lowers the remaining CKKS arithmetic into RNS and NTT operations.

Another caller may choose late transition placement, additional analyses, a different lowering route, or a pipeline that intentionally stops at an intermediate Program.

Keys, resources, and linking

analyze_evaluation_key_requirements(...) finds the required rotation and relinearization uses. The example creates those keys and uses prepare_material_bindings to prepare numerical data from saved descriptions before linking:

fh_compile.prepare_material_bindings(
    compiled,
    resources=resources,
    keys=(*rotation_keys, *((relinearization_key,) if relinearization_key is not None else ())),
)
backend = OperationBackend()
executable = backend.link(compiled)

The Program contains named placeholders. The separate dictionary supplies Tensor data without embedding an Engine or arithmetic Context in the Program. Linking does not generate data and leaves the source Compilation unchanged.

Encrypted execution

The example encrypts one tiled input, invokes the executable through the public Ciphertext boundary, decrypts the result, and compares it with the retained CapturedCallable reference:

result = executable.run(encrypted_x)
decoded = engine.decrypt_message(result, secret_key)
expected = captured_callable.reference(clear_x)

Registered Backend implementations remain Tensor-oriented internally. ProgramExecutable performs the public value adaptation at the outer execution interface.

Why this differs from Example 14

Example 12 uses only operations and passes already supplied by FHElium. Example 14 defines a new pass that rewrites captured torch.matmul into a baby-step/giant-step schedule before composing, linking, and executing its pipeline. Example 13 instead focuses on textual Program IR and analysis-pass composition and intentionally stops before Backend execution.

Source

#!/usr/bin/env python3

"""Compose built-in Compile passes, link the Program, and execute it."""

from __future__ import annotations

from collections import Counter
from typing import cast

import torch
from common import print_table

from fhelium import Preset
from fhelium import compile as fh_compile
from fhelium import ir
from fhelium.backend import OperationBackend
from fhelium.backend.ckks import CkksDeviceResources
from fhelium.config import CkksConfig
from fhelium.eager import Engine
from fhelium.values import Ciphertext

_ROTATION = 3
_LOGICAL_SLOTS = 16


def rotated_quadratic(
    x: torch.Tensor,
    rotation: int,
) -> torch.Tensor:
    """Square the sum of one vector and its cyclic rotation."""

    mixed = x + torch.roll(x, shifts=rotation, dims=-1)
    return mixed * mixed


def _dialect_counts(program: ir.Program) -> Counter[str]:
    """Count operation namespaces in one Program."""

    return Counter(
        operation.name.split(".", maxsplit=1)[0] for operation in program.walk()
    )


def main() -> None:
    config = CkksConfig.parse(Preset.slots8192_scale40_depth7_int64)
    device = "cpu"
    workspace = fh_compile.CompileWorkspace({CkksConfig: config})

    captured = fh_compile.capture(
        rotated_quadratic,
        inputs={
            "x": fh_compile.encrypted(
                slots=_LOGICAL_SLOTS,
                polynomial_domain="coefficient",
                residue_representation="standard",
            ),
            "rotation": fh_compile.static(_ROTATION),
        },
        workspace=workspace,
    )

    # This workload selects one concrete composition from the built-in pass
    # catalog. FHElium does not impose this sequence as a global default.
    pipeline = fh_compile.Pipeline(
        (
            fh_compile.EliminateDeadValuesPass(),
            fh_compile.LowerSemanticToLogicalPass(),
            fh_compile.InsertMultiplyNttTransitionsPass(),
            fh_compile.LowerLogicalToCkksPass(),
            fh_compile.ResolveRotationKeyOperandsPass(),
            fh_compile.InsertRelinearizationPass(),
            fh_compile.InsertRescalePass(),
            fh_compile.AssignCkksDepthsPass(entry_depth=0),
            fh_compile.AssignCkksScalesPass(
                entry_scale=config.default_scale,
            ),
            fh_compile.AssignNttImplementationPass(
                "native-ntt", ntt_backend="radix2_indexed"
            ),
            fh_compile.LowerCkksToRnsNttPass(),
            fh_compile.SelectNttImplementationsPass(
                OperationBackend().registry
            ),
            fh_compile.PrepareOperationOperandsPass(
                OperationBackend().registry
            ),
        )
    )
    compiled = pipeline.run(captured)

    requirements = ir.analyze_evaluation_key_requirements(compiled.program)
    engine = Engine(config, rng_seed=17)
    secret_key = engine.create_secret_key(device=device)
    public_key = engine.create_public_key(secret_key, device=device)
    rotation_keys = tuple(
        engine.create_rotation_key(step, secret_key, device=device)
        for step in sorted(requirements.rotation_steps)
    )
    relinearization_key = (
        engine.create_relinearization_key(secret_key, device=device)
        if requirements.requires_relinearization
        else None
    )
    resources = CkksDeviceResources(
        config=config,
        device=device,
        rng_seed=29,
    )
    fh_compile.prepare_material_bindings(
        compiled,
        resources=resources,
        keys=(
            *rotation_keys,
            *(
                (relinearization_key,)
                if relinearization_key is not None
                else ()
            ),
        ),
    )
    backend = OperationBackend()
    executable = backend.link(compiled)

    logical_x = torch.linspace(
        -0.04,
        0.04,
        _LOGICAL_SLOTS,
        dtype=torch.float64,
    )
    clear_x = logical_x.repeat(config.num_slots // _LOGICAL_SLOTS)
    captured_callable = cast(
        fh_compile.CapturedCallable[torch.Tensor],
        captured.workspace[fh_compile.CapturedCallable],
    )
    expected = captured_callable.reference(clear_x)

    encrypted_x = engine.encrypt_message(
        clear_x,
        public_key,
        depth=0,
        scale=config.default_scale,
        device=device,
    )
    result = cast(
        Ciphertext,
        executable.run(encrypted_x),
    )
    decoded = engine.decrypt_message(result, secret_key)
    torch.testing.assert_close(
        decoded.real,
        expected,
        rtol=5e-6,
        atol=5e-6,
    )
    torch.testing.assert_close(
        decoded.imag,
        torch.zeros_like(decoded.imag),
        rtol=0.0,
        atol=5e-6,
    )

    captured_inventory = ir.inventory_program(captured.program)
    compiled_inventory = ir.inventory_program(compiled.program)
    print_table(
        ["stage", "operations", "dialects"],
        [
            [
                "captured",
                sum(captured_inventory.operation_counts.values()),
                dict(sorted(_dialect_counts(captured.program).items())),
            ],
            [
                "backend-ready",
                sum(compiled_inventory.operation_counts.values()),
                dict(sorted(_dialect_counts(compiled.program).items())),
            ],
        ],
    )
    print()
    print_table(
        ["execution", "value"],
        [
            ["rotation steps", sorted(requirements.rotation_steps)],
            ["relinearization key", requirements.requires_relinearization],
            ["result depth", result.depth],
            ["result scale", f"{result.scale:.6e}"],
            [
                "maximum error",
                f"{float((decoded.real - expected).abs().max()):.6e}",
            ],
        ],
    )
    print()
    print_table(
        ["pass", "matched", "transformed", "inserted", "removed", "skipped"],
        [
            [
                report.name,
                report.stats.matched,
                report.stats.transformed,
                report.stats.inserted,
                report.stats.removed,
                report.stats.skipped,
            ]
            for report in compiled.reports
        ],
    )


if __name__ == "__main__":
    main()

Leaving NTT choices open

Logical NTT operations can leave their algorithm unassigned. Add SelectNttImplementationsPass(backend.registry) where the required facts become available, then supply any newly introduced table materials before linking. AssignNttImplementationPass can constrain an algorithm, group width or fixed radix without assigning every execution detail. See NTT semantics and selection for the selection order and table-data responsibilities.