Customize a Compile pass and pipeline

Example source: examples/14_compile_custom_pass.py

Example 14 implements a caller-defined Compile pass that recognizes a captured torch.matmul with a compile-time square matrix and rewrites it as a baby-step/giant-step (BSGS) encrypted matrix-vector schedule. The source applies an illustrative elementwise square after the linear result, adding one ciphertext-ciphertext multiplication. Caller-selected passes then place relinearization and rescales, assign depths and actual scales, lower the schedule into RNS/NTT operations, link materials and evaluation keys, and execute it through the Backend.

The example demonstrates a transformation and execution architecture. Its CPU execution validates the selected matrix-multiplication schedule for the shown inputs.

Run the example

python examples/14_compile_custom_pass.py

The output presents three Programs:

1. the captured Program containing torch.call for torch.matmul;
2. the semantic BSGS Program produced by the custom pass;
3. the executable RNS/NTT Program produced by the complete selected pipeline.

Captured operation

The source callable is ordinary PyTorch:

def matrix_vector(x):
    linear = torch.matmul(_WEIGHT, x)
    return linear * linear

_WEIGHT is a module-level Tensor. PyTorch FX captures it as a graph-external fhelium.material.ref, and FHElium retains it in Compilation.material_bindings. PyTorch FX preserves encrypted matmul as a structural torch.call; the custom pass assigns its FHElium meaning.

A runtime matrix argument would not have a compile-time Tensor payload. This example therefore restricts the pattern to one captured square matrix material multiplied by one encrypted vector.

Cyclic-diagonal identity

For an n × n matrix A, define cyclic diagonal D_k by

$$D_k[r]=A[r,(r-k)modn].$$

Let R_s denote torch.roll(..., shifts=s). The direct diagonal form is

$$Ax=\sum _{k=0}^{n-1}{D}_k\odot R_k(x),$$

where ⊙ is pointwise multiplication.

Choose a baby-step width b and write k = jb + i. The BSGS form used by the pass is

$$Ax=\sum _j{R}_{jb}(\sum _i{R}_i(x)\odot R_{-jb}(D_{jb+i})),$$

for indices jb+i < n.

The adjustment R_{-jb}(D_{jb+i}) is required: the outer giant rotation also rotates every diagonal coefficient in the group.

For the example's n=8 and b=3, the encrypted schedule uses baby rotations 1 and 2, then giant rotations 3 and 6. It retains eight plaintext multiplications and seven ciphertext additions while reducing the independent input-rotation family.

Custom pass responsibilities

LowerConstantMatmulToBsgsPass performs the following work:

1. find torch.call with target torch.matmul;
2. locate the matrix's MaterialRefOp and Tensor in Compilation.material_bindings;
3. compute the eight cyclic diagonals, their group-specific adjustments, and periodic copies covering the configured CKKS slot count;
4. store those derived Tensors under readable material symbols;
5. emit reusable baby rotations;
6. emit semantic pointwise multiplication and addition for each giant group;
7. emit the giant rotations and final group reduction;
8. replace the original matmul result with the BSGS result;
9. publish the selected matrix size, baby width, group count, and material symbols in the Compile workspace.

The pass receives both the baby width and the CKKS slot count:

LowerConstantMatmulToBsgsPass(
    baby_step=3,
    slot_count=config.num_slots,
)

Dead-value elimination can then remove the original matrix reference after the custom rewrite has replaced its only use.

Lowering through existing operations

The custom pass emits only shared semantic operations and material references:

• fhelium_semantic.roll;
• fhelium_semantic.multiply;
• fhelium_semantic.add;
• fhelium.material.ref.

The subsequent passes classify encrypted/public operand roles, resolve logical rotation steps to caller-bound key operands, prepare each diagonal as a multiplication-ready plaintext, and lower the encrypted schedule to shared CKKS, RNS, and NTT operations.

The selected InsertRelinearizationPass materializes the square activation's three-component result as a two-component ciphertext. Independently, the selected LateRescalePass consolidates the BSGS plaintext-product rescales within each add tree. A rotation is a barrier in this conservative policy because moving a rescale across key switching changes the active-Q key-switch work and error. The three giant groups therefore produce three rescales rather than eight; the square activation contributes one additional rescale. AssignCkksDepthsPass then reads those concrete transitions, and AssignCkksScalesPass computes each actual scale using the complete dropped Q-group product.

This keeps BSGS as one caller-selected algebraic transformation while preserving other matrix-multiplication representations.

Clear and encrypted checks

The example evaluates the same BSGS formula with clear float64 PyTorch tensors, applies the elementwise square, and compares it with the captured (_WEIGHT @ x) ** 2 computation. This check catches diagonal indexing, rotation-sign, and post-matvec dataflow mistakes in the example schedule.

The eight-element input and diagonals are repeated periodically across all configured slots. This makes a full-ring CKKS rotation agree with the intended eight-element cyclic rotation instead of padding the remaining slots with zeros. The example encrypts that tiled input, calls ProgramExecutable.run(ciphertext), decrypts the returned Ciphertext, and compares every decoded slot with the tiled clear result.

The example prepares numerical tables and key data separately from the Program:

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

The automatic preparation step adds parameters and key data to Compilation.material_bindings, leaving captured matrix materials in place. Backend linking reads this dictionary without generation or description-based correctness checks.

Backend implementations still consume and return Tensor payloads. The executable owns the public Program boundary: it unwraps the input ciphertext, runs the linked Tensor operations, and reconstructs the output ciphertext from the concrete result state assigned by Compile passes.

Transformation domain

The transformation rule in this example is defined for:

• a square matrix captured as a constant Tensor material;
• matrix-on-the-left vector multiplication;
• one-dimensional encrypted semantic input;
• a caller-selected baby-step width.

Dynamic matrices, rectangular matrices, batched matmul, transposed layouts, and other packing conventions require transformation rules that represent their own operand and layout semantics.

Source

#!/usr/bin/env python3

"""Define a custom Compile pass and pipeline for an executable CKKS plan.

A caller-defined Compile pass recognizes one constant square matrix multiplied
by an encrypted vector. It precomputes adjusted cyclic diagonals, replaces the
opaque PyTorch call with baby-step/giant-step (BSGS) rotations, pointwise
multiplications, and additions. The source then applies an illustrative
elementwise square activation, which contributes one ciphertext-ciphertext
multiplication. Caller-selected passes place relinearization and rescales,
lower the schedule into RNS/NTT operations, link materials and evaluation
keys, and execute the Program through the Backend.
"""

from __future__ import annotations


import math
from collections import Counter
from dataclasses import dataclass
from typing import cast

import torch
from common import print_table
from xdsl.dialects.builtin import IntegerAttr, StringAttr
from xdsl.ir import Operation
from xdsl.rewriter import Rewriter

from fhelium import Preset, 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.ir.dialects import core, semantic
from fhelium.ir.dialects import torch as torch_dialect
from fhelium.values import Ciphertext

_MATRIX_SIZE = 8
_BABY_STEP = 3
_INDEX = torch.arange(_MATRIX_SIZE, dtype=torch.float64)
_WEIGHT = (
    0.08 * torch.sin((_INDEX[:, None] + 1.0) * (_INDEX[None, :] + 2.0))
    + 0.04 * torch.cos(_INDEX[:, None] - 2.0 * _INDEX[None, :])
    + 0.35 * torch.eye(_MATRIX_SIZE, dtype=torch.float64)
)


def matrix_vector(x: torch.Tensor) -> torch.Tensor:
    """Apply a fixed matrix and an illustrative elementwise square."""

    linear = torch.matmul(_WEIGHT, x)
    return linear * linear


def _cyclic_diagonal(matrix: torch.Tensor, shift: int) -> torch.Tensor:
    """Return coefficients multiplying ``torch.roll(x, shift)``."""

    row = torch.arange(matrix.shape[0])
    column = (row - shift) % matrix.shape[0]
    return matrix[row, column]


def _bsgs_reference(
    matrix: torch.Tensor,
    vector: torch.Tensor,
    baby_step: int,
) -> torch.Tensor:
    """Evaluate the same BSGS schedule used by the custom pass in cleartext."""

    babies = [torch.roll(vector, shifts=step) for step in range(baby_step)]
    result = torch.zeros_like(vector)
    groups = math.ceil(matrix.shape[0] / baby_step)
    for group in range(groups):
        giant_shift = group * baby_step
        partial = torch.zeros_like(vector)
        for baby in range(baby_step):
            diagonal_index = giant_shift + baby
            if diagonal_index >= matrix.shape[0]:
                break
            diagonal = _cyclic_diagonal(matrix, diagonal_index)
            adjusted = torch.roll(diagonal, shifts=-giant_shift)
            partial = partial + babies[baby] * adjusted
        result = result + torch.roll(partial, shifts=giant_shift)
    return result


@dataclass(frozen=True)
class LowerConstantMatmulToBsgsPass:
    """Replace constant-matrix ``torch.matmul`` with a semantic BSGS schedule."""

    baby_step: int
    slot_count: int
    name: str = "lower-constant-matmul-to-bsgs"

    def run(self, compilation: fh_compile.Compilation) -> fh_compile.PassResult:
        program = compilation.program
        shared_data = compilation.workspace
        materials = compilation.material_bindings

        matched = transformed = inserted = 0
        plans: list[dict[str, object]] = []
        for operation in tuple(program.walk()):
            if not isinstance(operation, torch_dialect.CallOp):
                continue
            target = operation.attributes.get("fhelium.call.target")
            if (
                not isinstance(target, StringAttr)
                or target.data != "torch.matmul"
            ):
                continue
            matched += 1
            if len(operation.operands) != 2:
                raise ValueError("captured torch.matmul must have two operands")

            matrix_value, encrypted_vector = operation.operands
            matrix_owner = matrix_value.owner
            if not isinstance(matrix_owner, core.MaterialRefOp):
                raise ValueError(
                    "BSGS lowering requires a captured matrix material"
                )
            symbol = matrix_owner.attributes.get("symbol")
            if not isinstance(symbol, StringAttr):
                raise ValueError("captured matrix material has no symbol")
            matrix = materials[symbol.data]
            if not isinstance(matrix, torch.Tensor):
                raise TypeError("captured matrix material must be a Tensor")
            if matrix.ndim != 2 or matrix.shape[0] != matrix.shape[1]:
                raise ValueError("BSGS lowering requires a square matrix")

            matrix_size = matrix.shape[0]
            if self.slot_count < matrix_size or self.slot_count % matrix_size:
                raise ValueError(
                    "BSGS slot_count must be a positive multiple of the "
                    f"matrix size {matrix_size}, got {self.slot_count}"
                )
            groups = math.ceil(matrix_size / self.baby_step)
            created: list[Operation] = []
            babies = [encrypted_vector]
            for baby in range(1, self.baby_step):
                rotation = semantic.RollOp(
                    encrypted_vector,
                    result_type=operation.result.type,
                    attributes={
                        "shift": IntegerAttr(baby, 64),
                        "dimension": IntegerAttr(-1, 64),
                    },
                )
                rotation.result.name_hint = f"baby_{baby}"
                created.append(rotation)
                babies.append(rotation.result)

            result = None
            diagonal_symbols: list[str] = []
            for group in range(groups):
                giant_shift = group * self.baby_step
                partial = None
                for baby in range(self.baby_step):
                    diagonal_index = giant_shift + baby
                    if diagonal_index >= matrix_size:
                        break
                    diagonal = _cyclic_diagonal(matrix, diagonal_index)
                    adjusted = torch.roll(diagonal, shifts=-giant_shift).clone()
                    adjusted = adjusted.repeat(self.slot_count // matrix_size)
                    diagonal_symbol = f"compile/bsgs/diagonal/{group}/{baby}"
                    materials[diagonal_symbol] = adjusted
                    diagonal_symbols.append(diagonal_symbol)
                    diagonal_type = semantic.PublicType().with_state(
                        {"role": StringAttr("message")}
                    )
                    diagonal_ref = core.MaterialRefOp(
                        diagonal_type,
                        symbol=diagonal_symbol,
                        kind="tensor",
                    )
                    program.set_material_description(
                        diagonal_symbol,
                        {
                            "kind": "Tensor",
                            "label": f"matrix_diagonal_{diagonal_index}",
                        },
                    )
                    diagonal_ref.value.name_hint = f"diagonal_{diagonal_index}"
                    product = semantic.MultiplyOp(
                        babies[baby],
                        diagonal_ref.value,
                        result_type=operation.result.type,
                    )
                    product.result.name_hint = f"group_{group}_product_{baby}"
                    created.extend((diagonal_ref, product))
                    if partial is None:
                        partial = product.result
                    else:
                        addition = semantic.AddOp(
                            partial,
                            product.result,
                            result_type=operation.result.type,
                        )
                        addition.result.name_hint = f"group_{group}_sum"
                        created.append(addition)
                        partial = addition.result

                if partial is None:
                    raise RuntimeError("BSGS group contained no diagonal")
                if giant_shift:
                    giant = semantic.RollOp(
                        partial,
                        result_type=operation.result.type,
                        attributes={
                            "shift": IntegerAttr(giant_shift, 64),
                            "dimension": IntegerAttr(-1, 64),
                        },
                    )
                    giant.result.name_hint = f"giant_{giant_shift}"
                    created.append(giant)
                    partial = giant.result
                if result is None:
                    result = partial
                else:
                    addition = semantic.AddOp(
                        result,
                        partial,
                        result_type=operation.result.type,
                    )
                    addition.result.name_hint = "bsgs_result"
                    created.append(addition)
                    result = addition.result

            if result is None:
                raise RuntimeError("BSGS lowering produced no result")
            result.name_hint = operation.result.name_hint
            Rewriter.replace_op(
                operation,
                tuple(created),
                new_results=(result,),
            )
            transformed += 1
            inserted += len(created)
            plans.append(
                {
                    "matrix_size": matrix_size,
                    "slot_count": self.slot_count,
                    "baby_step": self.baby_step,
                    "giant_groups": groups,
                    "diagonal_symbols": tuple(diagonal_symbols),
                }
            )

        shared_data["compile/bsgs/plans"] = tuple(plans)
        if transformed:
            shared_data["compile/bsgs/semantic_program"] = program.clone()
        if transformed == 0:
            return fh_compile.PassResult.unchanged(program, matched=matched)
        return fh_compile.PassResult(
            program,
            fh_compile.PassStats(
                matched=matched,
                transformed=transformed,
                inserted=inserted,
                removed=transformed,
            ),
            diagnostics=("lowered torch.matmul to semantic BSGS",),
        )


def _torch_call_targets(program: ir.Program) -> list[str]:
    """Return encoded PyTorch call targets in structural order."""

    targets: list[str] = []
    for operation in program.walk():
        if not isinstance(operation, torch_dialect.CallOp):
            continue
        target = operation.attributes.get("fhelium.call.target")
        if isinstance(target, StringAttr):
            targets.append(target.data)
    return targets


def _program_summary(program: ir.Program) -> tuple[int, int, int]:
    """Return total operation, material-reference, and resource-reference counts."""

    operations = tuple(program.walk())
    return (
        len(operations),
        sum(
            isinstance(operation, core.MaterialRefOp)
            for operation in operations
        ),
        sum(
            isinstance(operation, core.ResourceRefOp)
            for operation in operations
        ),
    )


def _fhelium_operation_counts(program: ir.Program) -> list[tuple[str, int]]:
    """Count non-structural FHElium operations for compact terminal output."""

    counts = Counter(
        operation.name
        for operation in program.walk()
        if operation.name.startswith("fhelium_")
    )
    return sorted(counts.items())


def main() -> None:
    config = CkksConfig.parse(Preset.slots8192_scale40_depth7_int64)
    device = "cpu"
    capture_inputs = {
        "x": fh_compile.encrypted(
            slots=_MATRIX_SIZE,
            polynomial_domain="coefficient",
            residue_representation="standard",
        )
    }
    workspace = fh_compile.CompileWorkspace({CkksConfig: config})
    # Capture the source call while retaining the matrix Tensor in the
    # Compilation's material_bindings rather than embedding it in textual IR.
    captured = fh_compile.capture(
        matrix_vector,
        inputs=capture_inputs,
        workspace=workspace,
    )

    compiled = fh_compile.Pipeline(
        (
            # Replace torch.matmul with the inspectable BSGS rotation and
            # diagonal-product schedule, then remove the old matrix reference.
            LowerConstantMatmulToBsgsPass(
                _BABY_STEP,
                config.num_slots,
            ),
            fh_compile.EliminateDeadValuesPass(),
            # Classify encrypted/public operands and introduce concrete CKKS
            # plaintext preparation and representation transitions.
            fh_compile.LowerSemanticToLogicalPass(),
            fh_compile.InsertPlaintextPreparationPass(),
            fh_compile.InsertMultiplyNttTransitionsPass(),
            fh_compile.LowerLogicalToCkksPass(),
            # Give each rotation a named Tensor placeholder for its key data.
            fh_compile.ResolveRotationKeyOperandsPass(),
            # The square activation creates one CT×CT product. Select its
            # immediate relinearization independently from rescale placement.
            fh_compile.InsertRelinearizationPass(),
            # Consolidate product rescaling at legal add-tree frontiers before
            # assigning the resulting depths and per-value actual scales.
            fh_compile.LateRescalePass(),
            fh_compile.AssignCkksDepthsPass(entry_depth=0),
            fh_compile.AssignCkksScalesPass(
                entry_scale=config.default_scale,
            ),
            # Expose message encoding and lower the numerical computation.
            fh_compile.LowerMessagePlaintextPreparationPass(),
            fh_compile.AssignNttImplementationPass(
                "native-ntt", ntt_backend="radix2_indexed"
            ),
            fh_compile.LowerCkksToRnsNttPass(),
            fh_compile.SelectNttImplementationsPass(
                OperationBackend().registry
            ),
            fh_compile.PrepareOperationOperandsPass(
                OperationBackend().registry
            ),
        )
    ).run(captured)

    # Generate only the evaluation keys required by the transformed Program.
    requirements = ir.analyze_evaluation_key_requirements(compiled.program)
    rotation_steps = tuple(sorted(requirements.rotation_steps))
    engine = Engine(config, rng_seed=19)
    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 rotation_steps
    )
    relinearization_key = (
        engine.create_relinearization_key(secret_key, device=device)
        if requirements.requires_relinearization
        else None
    )
    # Prepare numerical data separately from the portable Program.
    device_resources = CkksDeviceResources(
        config=config,
        device=device,
        rng_seed=29,
    )
    fh_compile.prepare_material_bindings(
        compiled,
        resources=device_resources,
        keys=(
            *rotation_keys,
            *(
                (relinearization_key,)
                if relinearization_key is not None
                else ()
            ),
        ),
    )
    backend = OperationBackend()
    executable = backend.link(compiled)

    # Check the custom BSGS algebra independently with ordinary float64 data.
    clear_x = 0.03 * torch.sin(0.7 * _INDEX) + 0.01 * torch.cos(1.3 * _INDEX)
    captured_callable = cast(
        fh_compile.CapturedCallable[torch.Tensor],
        captured.workspace[fh_compile.CapturedCallable],
    )
    direct = captured_callable.reference(clear_x)
    bsgs_linear = _bsgs_reference(_WEIGHT, clear_x, _BABY_STEP)
    bsgs = bsgs_linear * bsgs_linear
    torch.testing.assert_close(
        bsgs,
        direct,
        rtol=8 * torch.finfo(torch.float64).eps,
        atol=8 * torch.finfo(torch.float64).eps,
    )
    # Repeat the logical eight-slot vector across the complete CKKS slot ring
    # so full-ring rotations implement the intended small cyclic layout.
    tiled_x = clear_x.repeat(config.num_slots // _MATRIX_SIZE)
    encrypted_x = engine.encrypt_message(
        tiled_x,
        public_key,
        depth=0,
        scale=config.default_scale,
        device=device,
    )
    result = executable.run(encrypted_x)
    encrypted_result = cast(Ciphertext, result)
    decoded = engine.decrypt_message(encrypted_result, secret_key)
    expected = direct.repeat(config.num_slots // _MATRIX_SIZE)
    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_targets = _torch_call_targets(captured.program)
    semantic_bsgs = cast(
        ir.Program,
        compiled.workspace["compile/bsgs/semantic_program"],
    )
    semantic_targets = _torch_call_targets(semantic_bsgs)
    compiled_targets = _torch_call_targets(compiled.program)
    captured_summary = _program_summary(captured.program)
    semantic_summary = _program_summary(semantic_bsgs)
    compiled_summary = _program_summary(compiled.program)

    print_table(
        [
            "stage",
            "IR ops",
            "materials",
            "resources",
            "torch.call targets",
            "rotation steps",
        ],
        [
            [
                "captured",
                *captured_summary,
                ", ".join(captured_targets) or "—",
                "—",
            ],
            [
                "semantic BSGS",
                *semantic_summary,
                ", ".join(semantic_targets) or "—",
                "—",
            ],
            [
                "compiled",
                *compiled_summary,
                ", ".join(compiled_targets) or "—",
                ", ".join(str(step) for step in rotation_steps) or "—",
            ],
        ],
    )
    print()
    print_table(
        ["BSGS property", "value"],
        [
            ["matrix size", _MATRIX_SIZE],
            ["CKKS slots", config.num_slots],
            ["baby step", _BABY_STEP],
            ["giant groups", math.ceil(_MATRIX_SIZE / _BABY_STEP)],
            ["derived diagonals", _MATRIX_SIZE],
            [
                "clear maximum error",
                f"{float((bsgs - direct).abs().max()):.3e}",
            ],
            [
                "Backend maximum error",
                f"{float((decoded.real - expected).abs().max()):.3e}",
            ],
            ["output depth", encrypted_result.depth],
            ["output scale", f"{encrypted_result.scale:.6e}"],
        ],
    )
    print()
    print_table(
        ["pass", "matched", "transformed", "inserted", "removed"],
        [
            [
                report.name,
                report.stats.matched,
                report.stats.transformed,
                report.stats.inserted,
                report.stats.removed,
            ]
            for report in compiled.reports
        ],
    )
    print()
    print_table(
        ["compiled FHElium operation", "count"],
        _fhelium_operation_counts(compiled.program),
    )


if __name__ == "__main__":
    main()

Next steps

• Compose and execute built-in Compile passes shows the end-to-end path without defining a new pass.
• Rank-local collective IR shows another caller-selected representation change.
• IR operation and implementation index lists the operations used by the resulting Program.