Skip to content

JIT guide

ren.jit is a replay cache for realized polynomial programs. It is useful when the same computation runs many times over arguments that contain Poly values with the same shape, ring, device, dtype, and domain.

The first call captures the program. Later calls with the same argument spec replay the captured exec units with new input buffers.

flowchart TB
  A["Call jitted function"]
  B["Cache lookup"]
  C["Capture on miss"]
  D["Replay on hit"]
  E["Return captured output structure"]

  A --> B
  B -- "miss" --> C
  B -- "hit" --> D
  C --> E
  D --> E

Minimal example

import ren
from ren.poly import Poly
from ren.ring import RingSpec

ring = RingSpec(n=8, moduli=(17, 97))

@ren.jit
def step(x: Poly, y: Poly) -> Poly:
    return (x + y) * (x - 1)

out1 = step(
    Poly([1, 2, 3, 4, 5, 6, 7, 8], ring),
    Poly([8, 7, 6, 5, 4, 3, 2, 1], ring),
)

out2 = step(
    Poly([2, 2, 2, 2, 2, 2, 2, 2], ring),
    Poly([1, 1, 1, 1, 1, 1, 1, 1], ring),
)

The first call builds and captures the schedule. The second call reuses the captured plan because the realized input specs match.

The inputs do not have to be direct Poly arguments. JIT walks ordinary containers and dataclasses, so higher-level objects that contain polys also work. For example, CKKS Ciphertext is a dataclass with two polynomial components:

import ren
from ren.schemes.ckks import Ciphertext

@ren.jit
def encrypted_step(x: Ciphertext) -> Ciphertext:
    return x + 1

Call this inside an active CKKS context, the same way you would call normal CKKS code.

What is in the cache key

JIT collects contained Poly inputs from ordinary containers and dataclasses, realizes them, then builds a cache key from their specs.

Realized inputs are keyed by:

  • ring
  • size
  • dtype
  • device
  • domain

Constant polynomial inputs are keyed by:

  • ring
  • constant value
  • device
  • domain

That means changing the ring, device, realized domain, buffer shape, dtype, or constant value can create a new captured plan.

Capture boundary

On a cache miss, JIT:

realize input Polys
call the user function
collect output Polys
realize outputs while recording ExecUnits
optionally prune input-independent work
replan captured intermediate memory
build input replacement tables
try to create a backend graph runner
store the captured plan

On a cache hit, JIT swaps the current input buffers into the captured plan, runs the backend graph runner if one exists, otherwise runs the recorded exec units, then clears bound inputs so stale buffers are not retained.

Dos

  • Use ren.jit around a stable polynomial computation that you call repeatedly.
  • Keep argument specs stable if you want replay instead of recapture.
  • Pass Poly values directly, or inside ordinary containers/dataclasses such as CKKS ciphertexts.
  • Realize or inspect outputs after each call when you need concrete values.
  • Use clear_jit() when you intentionally want to drop captured plans and recapture.
  • Consider @ren.jit(prune=True) when the function has expensive input-independent setup work.

Don'ts

  • Do not use JIT for a function whose arguments contain no Poly inputs.
  • Do not return only non-polynomial data.
  • Do not call another jitted function during capture.
  • Do not expect changed constants to reuse the same plan.
  • Do not expect one JIT object to hold unbounded shapes/specs; it stores a fixed number of captured plans.
  • Do not rely on each replay returning a fresh Python output object. The captured output structure is reused.

Common failure modes

No polynomial input

@ren.jit
def bad(scale: int) -> int:
    return scale + 1

This fails because the function receives no Poly, either directly or inside a supported argument structure. JIT needs at least one polynomial input to define the replay boundary.

No polynomial output

@ren.jit
def also_bad(x: Poly) -> int:
    return x.ring.n

This fails on capture because outputs must contain at least one Poly.

Accidental recapture

@ren.jit
def add_const(x: Poly, c: Poly) -> Poly:
    return x + c

add_const(x, Poly.const(7, x.ring))  # captures one plan
add_const(x, Poly.const(8, x.ring))  # captures another plan

Constant polynomial values are part of the cache key. If the value changes, Ren captures another plan.

Nested capture

@ren.jit
def inner(x: Poly) -> Poly:
    return x + 1

@ren.jit
def outer(x: Poly) -> Poly:
    return inner(x) * 2

This can fail when outer is capturing and inner tries to capture too. Keep the jitted boundary around the whole stable computation, or call non-jitted helpers inside it.

Feedback aliasing

@ren.jit
def step(x: Poly) -> Poly:
    return (x + 1).contiguous()

x = step(x)
x = step(x)

This pattern is supported: replay protects read-before-write aliasing when a previous output is fed back as a later input. It can still be surprising because the output object is reused across replays.