`ren.jit` replay¶

Overview¶

The @ren.jit decorator can make many computations faster. For example:

from time import perf_counter
from ren.schemes.ckks import CKKSParams
from ren.schemes.ckks.ckks import Ciphertext, encode, encrypt, keygen
from ren.schemes.ckks.keys import context
import ren
import random


def time_fn(f):
    start = perf_counter()
    f()
    end = perf_counter()
    t = int((end - start) * 1e6)
    print(f"{t:6} μs")


def random_ciphertext():
    data = [random.uniform(-1, 1) for _ in range(context().params.slot_count)]
    return encrypt(encode(data, msg_bound=1.0))


def warmup():
    x = random_ciphertext()
    without_jit(x).realize()


def without_jit(x: Ciphertext) -> Ciphertext:
    for _ in range(10):
        x = 2 * x * x - 1
    return x


@ren.jit
def with_jit(x: Ciphertext) -> Ciphertext:
    for _ in range(10):
        x = 2 * x * x - 1
    return x


ctx = keygen(CKKSParams.default())
with ctx:
    warmup()

    funcs = ["without_jit", "with_jit"]
    for func_name in funcs:
        print(func_name)
        for _ in range(5):
            x = random_ciphertext()
            func = globals()[func_name]
            time_fn(lambda: func(x).realize())
        print()

without_jit
216695 μs
217473 μs
223730 μs
213738 μs
223736 μs

with_jit
232799 μs
  7567 μs
  6482 μs
  6375 μs
  8065 μs

Why are the second and subsequent calls to the function so much faster when it is decorated with @ren.jit? When the user performs a polynomial computation, the ren engine creates a plan for executing that computation, as explained in detail in the Execution pipeline and Runtime sections. By default, the engine uses this plan once and then forgets about it. The @ren.jit decorator allows the plan to be stored for later reuse.

It is important to recognize that although the @ren.jit decorator is applied to a function, what it actually captures is a polynomial computation graph. In particular, the decorator only understands objects of type ren.poly.Poly, and of types like ren.schemes.ckks.Plaintext and ren.schemes.ckks.Ciphertext that are built from Poly.

import ren


@ren.jit
def non_poly(x: float) -> float:
    for _ in range(10):
        x = 2 * x * x - 1
    return x


try:
    non_poly(0.5)
except RuntimeError as e:
    print(repr(e))

RuntimeError('ren.jit: function must receive at least one Poly argument')

Therefore, the @ren.jit decorator should only be used on functions that can be represented by a polynomial computation graph. Applying the decorator to other functions can lead to runtime errors or even silent incorrect behavior. The decorator generally should not be applied to functions with side effects.

Usage and limitations¶

Function arguments¶

The @ren.jit decorator has limited support for arguments of types that are not derived from Poly.

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


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


def random_poly():
    arr = [random.randint(0, 10) for _ in range(8)]
    return Poly(arr, ring)


@ren.jit
def scalar_mul(x: Poly, y: int) -> Poly:
    print(f"building graph with y={y}")
    return x * y


x1 = random_poly()
# this call builds a graph for scalar_mul with y=2
scalar_mul(x1, 2).realize()
x2 = random_poly()
# this call reuses the graph
scalar_mul(x2, 2).realize()
x3 = random_poly()
try:
    # error: the graph cannot be reused because y changed
    scalar_mul(x3, 3).realize()
except RuntimeError as e:
    print(repr(e))

building graph with y=2
RuntimeError('ren.jit: non-Poly scalar arg changed (captured (2,), got (3,)); jit freezes it at capture. Wrap varying scalars as Plaintext.const(...), or clear_jit() to re-capture.')

The ren engine does not have the ability to build a graph that computes scalar_mul(x,y) for arbitrary x and y. However, it does know how to create a graph for any fixed y.

If you plan to call scalar_mul with a few different y-values (and you expect to use each y-value more than once), you can use clear_jit to forget the graph, or use functools.partial to create multiple functions and apply jit to each one.

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

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


def random_poly():
    arr = [random.randint(0, 10) for _ in range(8)]
    return Poly(arr, ring)


@ren.jit
def scalar_mul(x: Poly, y: int) -> Poly:
    print(f"building graph with y={y}")
    return x * y


for y in (2, 3):
    for _ in range(5):
        x = random_poly()
        scalar_mul(x, y).realize()
    scalar_mul.clear_jit()


def scalar_mul_no_jit(x: Poly, y: int) -> Poly:
    print(f"building graph with y={y}")
    return x * y


for y in (2, 3):
    jitted = ren.jit(functools.partial(scalar_mul_no_jit, y=y))
    for _ in range(5):
        x = random_poly()
        jitted(x).realize()

building graph with y=2
building graph with y=3
building graph with y=2
building graph with y=3

Overwritten output¶

As of this writing, each call to a jitted function overwrites previous outputs of the function. See issue #372.

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


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


def constant_poly(c: int) -> Poly:
    arr = [c] + [0] * 7
    return Poly(arr, ring)


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


z = constant_poly(0)
one = constant_poly(1)
two = constant_poly(2)
z += mul(one, one)
z += mul(two, two)
# incorrect result: calling mul(two, two) overwrites the output of mul(one, one),
# but the computation z += mul(one, one) has not been realized yet
print(z.tolist())

z = constant_poly(0)
one = constant_poly(1)
two = constant_poly(2)
z += mul(one, one)
z.realize()
z += mul(two, two)
# correct result: mul(one, one) added to x before calling mul(two, two)
print(z.tolist())

# The engine correctly handles the special case where the input from one computation
# is fed into the next


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


z = constant_poly(0)
one = constant_poly(1)
two = constant_poly(2)
z = mul_add(one, one, z)
z = mul_add(two, two, z)
# correct result: special handling by engine avoids previous issue
print(z.tolist())

[8, 0, 0, 0, 0, 0, 0, 0]
[5, 0, 0, 0, 0, 0, 0, 0]
[5, 0, 0, 0, 0, 0, 0, 0]

Unsupported operations¶

The .realize() function should not be called in a jitted function, because jit operates on a computation graph, and .realize() will break the graph into pieces:

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

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


# Not allowed; .realize() breaks the computation graph into two pieces
@ren.jit
def add_two(x: Poly) -> Poly:
    y = x + 1
    y.realize()
    return y + 1


x = Poly([1, 2, 3, 4, 0, 0, 0, 0], ring)
try:
    add_two(x).realize()
except RuntimeError as e:
    print(repr(e))

RuntimeError('ren.jit: realize should not be called in a jitted function')

Similarly, calling a jitted function from inside another jitted function is not allowed.

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


ring = RingSpec(n=8, moduli=(17, 97))
x = Poly([1, 2, 3, 4, 0, 0, 0, 0], ring)


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


# nested jit is not allowed
@ren.jit
def outer(x: Poly) -> Poly:
    return inner(x) * 2


try:
    outer(x).realize()
except RuntimeError as e:
    print(repr(e))


# use .fn to get the non-jitted version of the inner function
@ren.jit
def outer_fixed(x: Poly) -> Poly:
    return inner.fn(x) * 2  # pyright: ignore


outer_fixed(x).realize()

RuntimeError('ren.jit: nested jit is not supported')

As of this writing, @ren.jit does not support operations that use randomness. See issue #413. Note that ren.schemes.ckks.encrypt uses randomness.

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

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


# Not allowed; uses randomness
@ren.jit
def add_noise(x: Poly) -> Poly:
    return x + Poly.ternary(ring)


x = Poly([1, 2, 3, 4, 0, 0, 0, 0], ring)
try:
    add_noise(x).realize()
except RuntimeError as e:
    print(repr(e))

RuntimeError('Random number generation during jit capture is not supported')

Plan reuse¶

If a jitted function is called twice with arguments from two separate rings, it will create two different plans. Less intuitively, jit also creates separate plans for constant and non-constant polynomials. Moreover, for constant polynomials, it creates a new plan every time the constant changes.

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


ring = RingSpec(n=8, moduli=(17, 97))
x = Poly([1, 2, 3, 4, 0, 0, 0, 0], ring)
y = Poly([5, 6, 7, 8, 0, 0, 0, 0], ring)
z = Poly([9, 0, 1, 2, 0, 0, 0, 0], ring)


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


old_num_plans = 0


def plan_added() -> bool:
    global old_num_plans
    added = len(add_poly._plans) > old_num_plans
    old_num_plans = len(add_poly._plans)
    return added


def print_plan_added(name):
    if plan_added():
        print(f"added plan for {name}")
    else:
        print(f"reused plan for {name}")


add_poly(x, y).realize()
print_plan_added("x + y")
add_poly(x, z).realize()
print_plan_added("x + z")
add_poly(x, Poly.const(7, ring))
print_plan_added("x + 7")
add_poly(x, Poly.const(8, ring))
print_plan_added("x + 8")

ring = RingSpec(n=4, moduli=(17, 97))
x = Poly([1, 2, 3, 4], ring)
y = Poly([5, 6, 7, 8], ring)
add_poly(x, y).realize()
print_plan_added("new ring")

added plan for x + y
reused plan for x + z
added plan for x + 7
added plan for x + 8
added plan for new ring

Flowchart¶

On the first matching call, Ren runs the Python function, realizes the output, and records the scheduled backend work. Ren pays that capture cost on the first call, so a one-off call is usually not worth jitting.

Later calls with the same polynomial input specs and top-level scalar arguments replay the recorded exec units with new input buffers. During replay, Ren reuses the captured schedule instead of rebuilding and rescheduling the graph, so subsequent matching calls run faster.

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

Cache key¶

JIT collects contained Poly inputs from ordinary containers and dataclasses, realizes them, then builds a cache key from their specs. It also records top-level non-Poly scalar arguments because replay does not re-enter the Python function.

Realized inputs are keyed by:

ring
size
dtype
device
domain

Constant polynomial inputs are keyed by:

ring
constant value
device
domain

Changing the ring, device, realized domain, buffer shape, dtype, or constant value can create a new captured plan. If the polynomial specs match an existing plan but a top-level scalar argument changes, Ren raises instead of recapturing, because that scalar may have selected a Python branch during capture.

Example¶

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


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


@ren.jit
def step(x: Poly, y: Poly, offset: Poly) -> Poly:
    global captures
    captures += 1
    return (x + y + offset).contiguous()


out1 = step(
    Poly([1, 2, 3, 4, 0, 0, 0, 0], ring),
    Poly([4, 3, 2, 1, 0, 0, 0, 0], ring),
    Poly([1, 1, 1, 1, 0, 0, 0, 0], ring),
)
print("captures after first call:", captures)
print("first output:", out1.to_centered_list())

out2 = step(
    Poly([2, 2, 2, 2, 0, 0, 0, 0], ring),
    Poly([1, 1, 1, 1, 0, 0, 0, 0], ring),
    Poly([0, 1, 0, 1, 0, 0, 0, 0], ring),
)
print("captures after second call:", captures)
print("second output:", out2.to_centered_list())

captures after first call: 1
first output: [6, 6, 6, 6, 0, 0, 0, 0]
captures after second call: 1
second output: [3, 4, 3, 4, 0, 0, 0, 0]

The first call builds and captures the schedule, so captures becomes 1. The second call has new polynomial values but the same realized input specs, so Ren replays the captured plan without entering the Python function again and captures stays 1.

Capture boundary¶

On a cache miss, JIT:

realize input Polys
enter capture mode and call the user function
collect output Polys
realize outputs while recording ExecUnits
prune input-independent work if enabled
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.

Only JIT should realize outputs during capture. Calling realize() yourself inside the jitted body raises, because it would split execution at the wrong boundary. Creating random polynomials inside the body also raises so replay cannot reuse capture-time randomness.

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 want to drop captured plans and recapture.
Consider @ren.jit(prune=True) when the function has expensive input-independent setup work.
Keep one JIT boundary around a stable computation when you do not need separate inner and outer caches.

Don'ts¶

Do not use JIT for a function whose arguments contain no Poly inputs.
Do not return only non-polynomial data.
Do not expect changed polynomial constant inputs to reuse the same plan; constant values are part of the cache key.
Do not expect top-level scalar arguments to vary across replays. Use a Poly value, a CKKS Plaintext.const(...), or call clear_jit() before capturing a different branch.
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.
Do not call a jitted function from inside another jitted function. Keep helper functions undecorated and put one @ren.jit boundary around the stable computation.

Source links¶

ren/engine/jit.py

ren.jit replay¶