Trace, transform, and run a JIT program

Example source: examples/19_unified_jit.py

Example 19 traces a typed PyTorch square matrix-vector computation into one mixed-dialect xDSL Program, applies an selected lowering pipeline, provisions the evaluation keys required by the resulting operations, and runs the program with a retained workspace.

The example evaluates an 8 × 8 affine map in repeated packed slot blocks, then applies

p (z) = (z + 0.25) (z - 0.5)

and a caller-owned plaintext output gain. Public matrix preparation remains as preserved Torch operations. Arithmetic that mixes encrypted values is represented by FHElium semantic operations and then lowered to explicit CKKS operations.

Run the complete example

From the repository root:

bash

python examples/19_unified_jit.py \
  --preset slots8192-scale40-levels7-int64

The script runs on the selected CPU or CUDA engine and prints the final textual program, a report for every selected pass, the output level and actual scale, and the measured maximum absolute error against its semantic reference.

1. Define one semantic function

The input function uses ordinary PyTorch syntax:

python

def square_matvec_quadratic(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor,
    output_gain: torch.Tensor,
    matrix_size: int,
    repeats: int,
) -> torch.Tensor:
    main_diagonal = torch.diagonal(weight).repeat(repeats)
    affine = x * main_diagonal

    for shift in range(1, matrix_size):
        wrapped_head = torch.diagonal(
            weight,
            offset=matrix_size - shift,
        )
        ordinary_tail = torch.diagonal(weight, offset=-shift)
        cyclic_diagonal = torch.cat((wrapped_head, ordinary_tail))
        packed_diagonal = cyclic_diagonal.repeat(repeats)
        rotated = torch.roll(x, shifts=shift, dims=-1)
        affine = affine + rotated * packed_diagonal

    affine = affine + bias.repeat(repeats)
    quadratic = (affine + 0.25) * (affine - 0.5)
    return quadratic * output_gain

For rotation $r$ and matrix row $j$ , the packed diagonal is

d_{r} [j] = W [j, (j - r) mod 8] .

The packed affine computation is therefore

a = \sum_{r = 0}^{7} d_{r} ⊙ roll (x, r) + b .

Each public torch.diagonal, torch.cat, and repeat call remains in the mixed-dialect program as torch.call. Each supported operation whose result depends on the encrypted input is captured as an fhelium.semantic.* operation. The static loop bounds let FX unroll the seven nonzero rotations.

2. Trace with typed input roles

Import the unified package directly:

python

from fhelium.experimental import jit

Every function parameter receives one input declaration:

python

captured = jit.trace(
    square_matvec_quadratic,
    inputs={
        "x": jit.encrypted(),
        "weight": jit.message(),
        "bias": jit.message(),
        "output_gain": jit.plaintext(),
        "matrix_size": jit.static(8),
        "repeats": jit.static(engine.num_slots // 8),
    },
)

Declaration	Runtime meaning	Semantic reference meaning
`encrypted()`	A compatible core `Ciphertext`, or a Tensor encrypted online with `workspace["public_key"]`; the declaration carries exact level, scale, slot-extent, and batch policy	Tensor
`message()`	Public Python/PyTorch data; public-only computation remains a preserved Torch call until an encrypted consumer requires preparation	The same public value
`plaintext()`	A caller-owned core `Plaintext` whose representation and CKKS state reach the exact consumer	Public Tensor/scalar shadow
`static(value)`	A finite immutable scalar specialized during capture and omitted from the runtime signature	The specialized value restored by `CaptureResult.reference()`

encrypted() defaults to level zero and the selected engine's default scale at execution. slots="full" requires the complete engine slot axis; a positive integer declares a fixed logical final-axis extent. batch_mode="none" requires one slot axis, while "any" permits leading batch axes.

message() and plaintext() describe distinct public roles. A message is semantic Python/PyTorch data that can pass through public preprocessing. A plaintext is already a FHElium encoded value with exact representation, level, scale, basis, domain, and residue state.

3. Inspect the capture result

jit.trace() returns a CaptureResult, not a second program type:

python

source_program = captured.program
workspace = captured.workspace

print(source_program.to_text())
print(captured.fx_code)
print(captured.runtime_signature)

The result provides:

program: the canonical mixed-dialect xDSL Program;
workspace: retained graph-external materials and caller state;
signature, specs, and fx_code: frontend evidence;
reference(...): execution of the original callable with static values restored.

Tensor constants captured from Python are cloned into workspace["materials"]. Their SSA positions contain only symbolic fhelium.material.ref operations, so Program.save() does not serialize live Tensor values.

The source program is structurally valid immediately after capture. It can contain semantic operations and therefore need not yet be ready for the current interpreter:

python

source_report = source_program.readiness(workspace)
print(source_report.runnable)
print(source_report.diagnostics)

4. Apply a selected pass pipeline

Example 19 selects the general lowering policy:

python

lowered = jit.default_pipeline().run(
    captured.program,
    captured.workspace,
)
program = lowered.program
workspace = lowered.workspace

A pipeline clones the source Program once, runs the ordered pass tuple over the clone, and returns the same workspace object. The source program remains available for comparison.

The default pipeline currently runs:

unreachable pure-value elimination;
semantic-to-logical role classification;
operation-specific plaintext preparation insertion;
ciphertext-multiply NTT transition insertion;
logical-to-explicit-CKKS lowering;
relinearization insertion;
rescale insertion;
conservative late-rescale and late-relinearization passes.

Each pass handles only the local operations and roles it recognizes. A pass with no applicable pattern returns a legal unchanged result. Inspect the exact behavior through lowered.reports:

python

for report in lowered.reports:
    print(report.name, report.stats, report.diagnostics)

The pass counters distinguish matches, transformations, insertions, removals, and skips. The pipeline structurally verifies every pass result before invoking the next pass. Completion records the selected transformations; the independent readiness check decides whether the result can run.

Visualize the lowered Program

SvgGraphVisualizationPass renders the selected entry's lowered SSA/dataflow graph without changing the Program:

python

from pathlib import Path

from fhelium.experimental.jit.passes.visualize_svg import (
    SvgGraphPresentation,
)

output_path = Path("graph_exports/example19-lowered.svg")
output_path.parent.mkdir(parents=True, exist_ok=True)

visualization = jit.SvgGraphVisualizationPass(
    output_path,
    overwrite=True,
    presentation=SvgGraphPresentation(
        fields={
            "name",
            "opcode",
            "role",
            "operands",
            "attributes",
            "scheduling_obligations",
            "num_users",
        },
        attribute_names={
            "condition",
            "fhelium.call.target",
            "operation",
            "scale_mode",
            "shift",
        },
    ),
).run(program, workspace)
print(visualization.stats)

The pass draws entry arguments, explicit operations, SSA dependencies, result roles, and the function output. Arguments, constants, and outputs have dedicated colors. Operations of the same kind share a stable color selected from a diverse light palette derived from FHElium's cobalt, spectral-blue, helium-amber, warm-bridge, and neutral colors. Every fill uses the same high-contrast deep-neutral text color. Preserved torch.call nodes also distinguish their function or method target. fields selects complete node-record sections, while attribute_names optionally limits the exact xDSL attributes shown inside the attributes section.

Rendering requires the Python pydot package and the system Graphviz dot executable. The default overwrite=False refuses to replace an existing file; the example opts into replacement so the command is repeatable.

Figure 1. The Example 19 main entry after the default lowering pipeline. Open the SVG to inspect the complete 87-operation graph at full resolution.

Use Visualize and inspect a JIT Program to select state/type fields, control long attributes, compare pass snapshots, or customize node-record rows.

5. Analyze the transformed requirements

Analyze evaluation-key requirements after lowering, because the analysis scans the current explicit CKKS operations:

python

key_requirements = jit.analyze_evaluation_key_requirements(program)

evaluation_keys = fh.EvaluationKeySet(
    rotations=fh.RotationKeySet(
        {
            step: engine.rotation_key(step)
            for step in key_requirements.rotation_steps
        }
    ),
    relinearization=(
        engine.relinearization_key
        if key_requirements.requires_relinearization
        else None
    ),
)

Example 19's packed matrix-vector program requires rotation steps 1 through 7. Its quadratic contains a ciphertext-ciphertext multiplication and therefore requires relinearization.

program.requirements() provides the broader pure scan, including current operation names, symbolic materials/resources, preserved Torch targets, unknown operation names, engine requirement, and return count. This analysis does not inspect the workspace or decide whether execution should proceed.

6. Populate the retained workspace

Add live evaluator services to the same workspace:

python

workspace.update(
    {
        "engine": engine,
        "evaluation_keys": evaluation_keys,
    }
)

The program contains the computation and symbolic identities. The workspace contains the live engine, keys, captured materials, handlers, resources, policies, and caches for this request. Keeping these objects graph-external allows the same textual program to be paired with another compatible runtime.

When an encrypted argument is supplied as a Tensor, online encryption also requires:

python

workspace["public_key"] = engine.public_key

Example 19 supplies a caller-created Ciphertext, so online encryption is not part of its execution requirements.

7. Check readiness before execution

Readiness compares the selected entry with the exact workspace without running or materializing anything:

python

report = program.readiness(workspace)
if not report.runnable:
    for diagnostic in report.diagnostics:
        print(diagnostic.code, diagnostic.subject, diagnostic.message)
    raise RuntimeError("program is not ready")

The report covers:

the versioned program schema and dialect;
one structurally executable selected entry;
exact built-in operation schemas and cleared scheduling obligations;
trusted handlers for extension operations or FHE-touching Torch targets;
symbolic material and resource bindings;
an engine and the required evaluation keys.

program.run() performs the same gate and raises ProgramNotReadyError with its complete report when a requirement is missing.

8. Build semantic and encrypted inputs

The gain has separate semantic and encoded forms:

python

output_gain_value = 0.75
output_gain_plaintext = engine.plaintext(
    output_gain_value,
    level=2,
    scale=engine.config.default_scale,
)

The capture result's reference receives the semantic scalar:

python

reference = captured.reference(
    packed_x,
    weight,
    bias,
    output_gain_value,
)

The JIT program receives the core plaintext and an encrypted packed input:

python

encrypted_x = engine.encrypt_message(packed_x, engine.public_key)

result = program.run(
    encrypted_x,
    weight,
    bias,
    output_gain_plaintext,
    workspace=workspace,
)

Input names recorded by capture allow either positional or keyword binding. A caller-owned Ciphertext must match the workspace engine's context, device, dtype, ring dimension, declared level and actual scale, and batch policy. A plaintext() argument must be a core Plaintext; its exact representation state is handled at the consuming preparation operation.

The interpreter executes preserved public Torch calls through its audited public target table. A preserved Torch call that touches encrypted or plaintext roles requires an exact binding under workspace["torch_handlers"]. Unknown extension operations require an exact handler under workspace["handlers"].

9. Decrypt and apply the example's acceptance criterion

Example 19 decrypts and compares every packed slot:

python

decoded = engine.decrypt_message(
    result,
    engine.secret_key,
    is_real=True,
)
error = torch.abs(decoded - reference)
max_abs_error = float(error.max())
if max_abs_error > 3e-5:
    raise RuntimeError(
        "JIT execution exceeded its fixed CKKS validation threshold: "
        f"max_abs_error={max_abs_error:.3e}, atol={3e-5:.3e}"
    )

The fixed absolute tolerance belongs to this preset, circuit, and deterministic input construction. It is evidence for this maintained example rather than a general numerical guarantee for every JIT program.

Alternative entry paths

The same Program model supports workflows that start outside PyTorch capture:

Example 20: Import and execute textual JIT IR parses versioned mixed-dialect text, retains an application operation, binds its handler, and executes the selected entry.
Example 21: Customize and audit a JIT pass pipeline demonstrates custom analysis/pass composition and the retained workspace.

For direct textual input:

python

program = jit.parse(text, source_name="application.mlir")
# or: program = jit.load("application.mlir")

transformed = jit.default_pipeline().run(program, jit.Workspace())
print(transformed.program.to_text())

Parsing establishes structural validity. Select passes, handlers, materials, and runtime services separately for the intended execution request.

Complete source

#!/usr/bin/env python3

"""Trace, lower, check, and run a dense encrypted matrix-vector program.

This example follows the JIT's trace-first model: tracing produces one mixed-
dialect Program, the default pipeline lowers its recognized encrypted
operations, readiness compares the lowered Program with retained runtime
bindings, and only ``Program.run(...)`` executes it. Public matrix preparation stays
as Torch operations while encrypted arithmetic becomes explicit CKKS operations.
"""

from __future__ import annotations

import argparse

import torch
from common import add_engine_args, make_engine, print_table

import fhelium as fh
from fhelium.experimental import jit

_MATRIX_SIZE = 8
_VALIDATION_ATOL = 3e-5


def square_matvec_quadratic(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor,
    output_gain: torch.Tensor,
    matrix_size: int,
    repeats: int,
) -> torch.Tensor:
    """Evaluate a repeated packed affine map followed by a quadratic."""

    main_diagonal = torch.diagonal(weight).repeat(repeats)
    affine = x * main_diagonal

    for shift in range(1, matrix_size):
        wrapped_head = torch.diagonal(
            weight,
            offset=matrix_size - shift,
        )
        ordinary_tail = torch.diagonal(weight, offset=-shift)
        cyclic_diagonal = torch.cat((wrapped_head, ordinary_tail))
        packed_diagonal = cyclic_diagonal.repeat(repeats)
        rotated = torch.roll(x, shifts=shift, dims=-1)
        affine = affine + rotated * packed_diagonal

    affine = affine + bias.repeat(repeats)
    quadratic = (affine + 0.25) * (affine - 0.5)
    return quadratic * output_gain


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    add_engine_args(parser, default_preset="slots8192-scale40-levels7-int64")
    args = parser.parse_args()

    engine = make_engine(args)
    if engine.num_slots % _MATRIX_SIZE:
        parser.error("the engine slot count must be divisible by matrix size")
    repeats = engine.num_slots // _MATRIX_SIZE

    # 1. Trace into the canonical mixed-dialect Program. Tracing does not run a
    # lowering pipeline and does not establish execution readiness.
    captured = jit.trace(
        square_matvec_quadratic,
        inputs={
            "x": jit.encrypted(),
            "weight": jit.message(),
            "bias": jit.message(),
            "output_gain": jit.plaintext(),
            "matrix_size": jit.static(_MATRIX_SIZE),
            "repeats": jit.static(repeats),
        },
    )

    traced_readiness = captured.program.readiness(captured.workspace)
    if traced_readiness.runnable:
        raise RuntimeError("the traced semantic Program was unexpectedly ready")

    # 2. Select and run the ordinary lowering policy. A pipeline
    # transforms one clone of the canonical xDSL Program. The same retained
    # Workspace carries caller policy, graph-external materials, and runtime
    # capabilities without embedding live objects in the graph.
    lowered = jit.default_pipeline().run(
        captured.program,
        captured.workspace,
    )
    program = lowered.program
    workspace = lowered.workspace
    if workspace is not captured.workspace:
        raise RuntimeError("the pass pipeline did not retain its Workspace")

    # 3. Analyze requirements without materializing keys. Readiness is a
    # separate readiness check and is expected to fail before runtime bindings exist.
    key_requirements = jit.analyze_evaluation_key_requirements(program)
    before_bindings = program.readiness(workspace)
    if before_bindings.runnable:
        raise RuntimeError(
            "the lowered CKKS Program was ready without an engine"
        )

    evaluation_keys = fh.EvaluationKeySet(
        rotations=fh.RotationKeySet(
            {
                step: engine.rotation_key(step)
                for step in key_requirements.rotation_steps
            }
        ),
        relinearization=(
            engine.relinearization_key
            if key_requirements.requires_relinearization
            else None
        ),
    )
    workspace.update(
        {
            "engine": engine,
            "evaluation_keys": evaluation_keys,
        }
    )
    ready = program.readiness(workspace)
    if not ready.runnable:
        detail = "; ".join(item.message for item in ready.diagnostics)
        raise RuntimeError(f"the bound JIT Program is not ready: {detail}")

    index = torch.arange(_MATRIX_SIZE, dtype=torch.float64)
    application_x = 0.025 * torch.sin(0.7 * index) + 0.01 * torch.cos(
        1.3 * index
    )
    rows = index[:, None]
    columns = index[None, :]
    weight = (
        0.06 * torch.sin((rows + 1.0) * (columns + 2.0))
        + 0.03 * torch.cos(rows - 2.0 * columns)
        + 0.28 * torch.eye(_MATRIX_SIZE, dtype=torch.float64)
    )
    bias = 0.012 * torch.cos(0.9 * index)
    packed_x = application_x.repeat(repeats)
    output_gain_value = 0.75
    output_gain_plaintext = engine.plaintext(
        output_gain_value,
        level=2,
        scale=engine.config.default_scale,
    )

    reference = captured.reference(
        packed_x,
        weight,
        bias,
        output_gain_value,
    )
    clear_affine = application_x @ weight.T + bias
    clear_block = (
        (clear_affine + 0.25) * (clear_affine - 0.5) * output_gain_value
    )
    torch.testing.assert_close(
        reference[:_MATRIX_SIZE],
        clear_block,
        rtol=0.0,
        atol=0.0,
    )

    encrypted_x = engine.encrypt_message(packed_x, engine.public_key)
    # 4. Run exactly the Program whose readiness was checked. Supplying an
    # already encrypted value keeps online encryption outside this run request.
    result = jit.run(
        program,
        encrypted_x,
        weight,
        bias,
        output_gain_plaintext,
        workspace=workspace,
    )
    decoded = engine.decrypt_message(
        result,
        engine.secret_key,
        is_real=True,
    )
    error = torch.abs(decoded - reference)
    max_abs_error = float(error.max())
    if max_abs_error > _VALIDATION_ATOL:
        raise RuntimeError(
            "JIT execution exceeded its fixed CKKS validation threshold: "
            f"max_abs_error={max_abs_error:.3e}, "
            f"atol={_VALIDATION_ATOL:.3e}"
        )

    print_table(
        ["stage", "operations", "runnable", "diagnostics"],
        [
            [
                "trace",
                len(traced_readiness.requirements.operations),
                traced_readiness.runnable,
                ", ".join(item.code for item in traced_readiness.diagnostics),
            ],
            [
                "lowered, unbound",
                len(before_bindings.requirements.operations),
                before_bindings.runnable,
                ", ".join(item.code for item in before_bindings.diagnostics),
            ],
            [
                "lowered, bound",
                len(ready.requirements.operations),
                ready.runnable,
                ", ".join(item.code for item in ready.diagnostics) or "none",
            ],
        ],
    )
    print()
    print_table(
        ["planned evaluation keys", "value"],
        [
            ["rotation steps", sorted(key_requirements.rotation_steps)],
            ["relinearization", key_requirements.requires_relinearization],
        ],
    )
    print()
    print("--- lowered mixed-dialect Program ---")
    print(program.to_text())
    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 lowered.reports
        ],
    )
    print()
    print_table(
        ["execution", "level", "scale", "max abs error"],
        [
            [
                "trace -> passes -> run",
                result.level,
                f"{result.scale:.6e}",
                f"{max_abs_error:.3e}",
            ]
        ],
    )


if __name__ == "__main__":
    main()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262

Continue

JIT programs defines the Program, workspace, validity, and control-level concepts.
JIT internals specifies xDSL, schemas, pass scope, and extension interfaces.
Textual JIT IR and custom JIT pipelines continue with independent advanced workflows.
Input-role API, Program API, execution/readiness API, and pass API provide exact current signatures.

Trace, transform, and run a JIT program ​

Run the complete example ​

1. Define one semantic function ​

2. Trace with typed input roles ​

3. Inspect the capture result ​

4. Apply a selected pass pipeline ​

Visualize the lowered Program ​

5. Analyze the transformed requirements ​

6. Populate the retained workspace ​

7. Check readiness before execution ​

8. Build semantic and encrypted inputs ​

9. Decrypt and apply the example's acceptance criterion ​

Alternative entry paths ​

Complete source ​

Continue ​

Trace, transform, and run a JIT program

Run the complete example

1. Define one semantic function

2. Trace with typed input roles

3. Inspect the capture result

4. Apply a selected pass pipeline

Visualize the lowered Program

5. Analyze the transformed requirements

6. Populate the retained workspace

7. Check readiness before execution

8. Build semantic and encrypted inputs

9. Decrypt and apply the example's acceptance criterion

Alternative entry paths

Complete source

Continue