Generate editable Python from a Program

Example source: examples/15_compile_python_codegen.py

Example 15 uses two optional Compile passes to emit Python at two different Program stages:

• EmitEagerPythonPass translates CKKS operations into calls on the public Eager Engine;
• EmitBackendPythonPass resolves the operations visible at its position and emits direct calls to their Backend implementation classes.

Both passes scan the Program supplied at their own pipeline position, leave it unchanged, and publish source in the shared CompileWorkspace. Neither pass inserts other transformations or requires a particular predecessor. If an operation at the selected stage cannot be represented, that emitter reports the operation and stops.

Run the example

python examples/15_compile_python_codegen.py

The example emits both source forms, executes them on CPU, checks that their Tensor payloads are identical, and compares the decrypted Eager result with the clear rotated sum.

Eager Python emission

The example first transforms a captured torch.roll plus addition into CKKS operations and places the Eager emitter at that point:

from fhelium import compile as fh_compile

ckks_compilation = fh_compile.Pipeline(
    (
        fh_compile.EliminateDeadValuesPass(),
        fh_compile.LowerSemanticToLogicalPass(),
        fh_compile.LowerLogicalToCkksPass(),
        fh_compile.ResolveRotationKeyOperandsPass(),
        fh_compile.AssignCkksDepthsPass(entry_depth=0),
        fh_compile.AssignCkksScalesPass(
            entry_scale=config.default_scale,
        ),
        fh_compile.EmitEagerPythonPass(),
    )
).run(captured)

eager_source = ckks_compilation.workspace[
    fh_compile.EagerPythonSource
]

The emitted function contains ordinary Eager calls:

def generated_eager(engine, x, *, materials, resources):
    rotation_key = materials["rotation-key:3"]
    rotated = engine.rotate_with_key(x, rotation_key)
    result = engine.add(x, rotated)
    return result

The complete emitted source retains SSA aliases introduced by the current Program rather than reconstructing the original Python spelling. Materials and keys remain function inputs. The source does not contain Tensor payloads, secret-key data, devices, RNS contexts, or buffers.

The Eager emitter accepts flat CKKS operations that have a public Engine equivalent, one-to-one boundary casts, material/resource references, constants, and func.return. Semantic, logical, RNS, NTT, distributed, unknown, and region-owning operations are outside this emitter's current surface. A caller chooses a CKKS stage when Eager-style source is desired; the pass itself does not choose or enforce that stage.

Backend Python emission

The example continues from the CKKS Compilation, lowers arithmetic and rotation into RNS/NTT operations, and invokes the Backend emitter:

backend_compilation = fh_compile.Pipeline(
    (
        fh_compile.AssignNttImplementationPass("native-ntt", ntt_backend="radix2_indexed"),
        fh_compile.LowerCkksToRnsNttPass(),
        fh_compile.EmitBackendPythonPass(),
    )
).run(ckks_compilation)

backend_source = backend_compilation.workspace[
    fh_compile.BackendPythonSource
]

The generated module imports and instantiates the selected implementation classes, reconstructs literal OperationInvocation descriptors, and calls each implementation directly:

implementation_0 = NativeRnsLinearImplementation()
invocation_0 = OperationInvocation(
    operation_type=AddStandardOp,
    operand_count=3,
    result_count=1,
    ...
)

def generated_backend(x, y, *, materials, resources):
    parameters = materials["rns_parameters/Q/0"]
    result, = implementation_0.execute(
        invocation_0,
        (x, y, parameters),
        (),
        in_place=False,
    )
    return result

The generated function calls Backend implementations directly with Tensor inputs and outputs. The developer can edit implementation construction, resource use, call order, temporary values, or replace a call with experimental code.

EmitBackendPythonPass uses the built-in implementation registry by default. A developer using a custom registry supplies it directly:

fh_compile.EmitBackendPythonPass(registry=my_registry)

Resolution happens while that pass scans the current Program. It does not consume a ProgramDispatchTable produced by another pass and therefore does not impose a pipeline order. Unsupported or ambiguous operations are reported by the emitter instead of being lowered or silently assigned.

Implementation classes must be importable at module scope. A zero-state custom implementation can be constructed directly. The emitter refuses to copy constructor state from an external custom implementation into source; a developer can first replace it with an importable source-level implementation or emit at another stage. FHElium-owned implementation constructor fields contain implementation-selection data and may be reproduced when needed.

Materials and live resources

Both source forms use ordinary mappings. Eager source consumes public values and key objects; Backend source consumes numerical Tensor payloads:

generated_backend(
    input_tensor,
    materials={
        "rotation-key:3": rotation_key.data,
        "rns_parameters/Q/0": parameter_tensor,
    },
    resources={},
)

The source artifacts record referenced symbols. Backend resource requirements describe non-Tensor execution handles such as process groups or an encryption sampler. Numerical tables and key data arrive through Tensor operands, without Context/key resource wrappers. Eager emission rejects operations carrying numerical table operands; use Backend emission to preserve those supplied tables.

The source emitted by one pass describes the Program snapshot seen at that position. Later passes may continue transforming the Compilation; they do not rewrite or invalidate an earlier source artifact. The caller retains the source corresponding to the stage it intends to edit or execute.

Generated Python is executable code. Execute only source emitted from trusted Programs and implementation registries, and review edited source before running it with live keys or resources.

Source

#!/usr/bin/env python3

"""Emit editable Eager and low-level Backend Python from two Program stages."""

from __future__ import annotations

from collections.abc import Callable
from typing import cast

import torch
from common import print_table

from fhelium import Preset
from fhelium import compile as fh_compile
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_sum(x: torch.Tensor, rotation: int) -> torch.Tensor:
    """Add one vector to its cyclic rotation."""

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


def _load_function(
    source: fh_compile.GeneratedPythonSource,
) -> Callable[..., object]:
    namespace: dict[str, object] = {}
    exec(source.source, namespace)
    return cast(Callable[..., object], namespace[source.entry_point])


def main() -> None:
    config = CkksConfig.parse(Preset.slots8192_scale40_depth7_int64)
    workspace = fh_compile.CompileWorkspace({CkksConfig: config})
    captured = fh_compile.capture(
        rotated_sum,
        inputs={
            "x": fh_compile.encrypted(
                slots=_LOGICAL_SLOTS,
                polynomial_domain="coefficient",
                residue_representation="standard",
            ),
            "rotation": fh_compile.static(_ROTATION),
        },
        workspace=workspace,
    )

    # Emit Eager Python from a CKKS-depth Program. The emitter is a normal
    # optional pass and does not prescribe the passes before or after it.
    ckks_compilation = fh_compile.Pipeline(
        (
            fh_compile.EliminateDeadValuesPass(),
            fh_compile.LowerSemanticToLogicalPass(),
            fh_compile.LowerLogicalToCkksPass(),
            fh_compile.ResolveRotationKeyOperandsPass(),
            fh_compile.AssignCkksDepthsPass(entry_depth=0),
            fh_compile.AssignCkksScalesPass(
                entry_scale=config.default_scale,
            ),
            fh_compile.EmitEagerPythonPass(),
        )
    ).run(captured)
    eager_source = cast(
        fh_compile.EagerPythonSource,
        ckks_compilation.workspace[fh_compile.EagerPythonSource],
    )

    # Continue lowering, then independently emit direct Backend implementation
    # calls. Moving either emitter to an unsupported stage makes that emitter
    # report the first operation it cannot convert.
    backend_compilation = fh_compile.Pipeline(
        (
            fh_compile.AssignNttImplementationPass(
                "native-ntt", ntt_backend="radix2_indexed"
            ),
            fh_compile.LowerCkksToRnsNttPass(),
            fh_compile.SelectNttImplementationsPass(
                OperationBackend().registry
            ),
            fh_compile.PrepareOperationOperandsPass(
                OperationBackend().registry
            ),
            fh_compile.EmitBackendPythonPass(),
        )
    ).run(ckks_compilation)
    backend_source = cast(
        fh_compile.BackendPythonSource,
        backend_compilation.workspace[fh_compile.BackendPythonSource],
    )

    engine = Engine(config, rng_seed=23)
    secret_key = engine.create_secret_key(device="cpu")
    public_key = engine.create_public_key(secret_key, device="cpu")
    rotation_key = engine.create_rotation_key(
        _ROTATION,
        secret_key,
        device="cpu",
    )
    logical_x = torch.linspace(
        -0.04,
        0.04,
        _LOGICAL_SLOTS,
        dtype=torch.float64,
    )
    clear_x = logical_x.repeat(config.num_slots // _LOGICAL_SLOTS)
    encrypted_x = engine.encrypt_message(
        clear_x,
        public_key,
        depth=0,
        scale=config.default_scale,
        device="cpu",
    )

    eager_function = _load_function(eager_source)
    eager_result = eager_function(
        engine,
        encrypted_x,
        materials={f"rotation-key:{_ROTATION}": rotation_key},
        resources={},
    )
    assert isinstance(eager_result, Ciphertext)

    materializer = CkksDeviceResources(
        config=config,
        device="cpu",
        rng_seed=29,
    )
    fh_compile.prepare_material_bindings(
        backend_compilation, resources=materializer, keys=(rotation_key,)
    )
    backend_function = _load_function(backend_source)
    backend_result = backend_function(
        encrypted_x.data,
        materials=backend_compilation.material_bindings,
        resources={},
    )
    assert isinstance(backend_result, torch.Tensor)
    assert torch.equal(backend_result, eager_result.data)

    decoded = engine.decrypt_message(eager_result, secret_key)
    expected = clear_x + torch.roll(clear_x, shifts=_ROTATION, dims=-1)
    torch.testing.assert_close(
        decoded.real,
        expected,
        rtol=5e-6,
        atol=5e-6,
    )

    print("--- Eager Python ---")
    print(eager_source.source)
    print("--- Backend Python ---")
    print(backend_source.source)
    print_table(
        ["generated target", "operations", "resources"],
        [
            [
                "Eager",
                eager_source.operation_count,
                eager_source.resource_symbols,
            ],
            [
                "Backend",
                backend_source.operation_count,
                backend_source.resource_symbols,
            ],
        ],
    )


if __name__ == "__main__":
    main()