Skip to content

Execution pipeline

Connect lazy Poly graphs to the scheduler, memory planner, backend, and JIT capture/replay paths. Read Engine overview first if you need the shorter map.

flowchart LR A["Lazy Poly graph"] B["Realization boundary"] C["Schedule, plan memory, lower"] D["Run ExecUnits"] E["Realized buffers"] A --> B --> C --> D --> E

JIT wraps the same realization machinery with a capture path on cache miss and a replay path on cache hit:

flowchart LR A["Call jitted function"] B["Realize input Polys"] C{"Plan cached?"} D["Capture:<br/>run function, realize outputs,<br/>record ExecUnits, store plan"] E["Replay:<br/>bind input buffers,<br/>run captured plan"] F["Return output structure"] A --> B --> C C -- "miss" --> D --> F C -- "hit" --> E --> F

Eager realization schedules, plans, lowers, and runs the lazy graph each time a concrete value is needed. JIT capture uses that same realization machinery on a cache miss while recording executable units; replay skips scheduling and lowering for the captured computation when the input specs match a captured plan.

Lazy graph

Poly wraps a Node. Poly operations add nodes without launching backend kernels.

  • a + b, a - b, and a * b create arithmetic nodes.
  • a.ntt() and a.intt() create domain conversion nodes.
  • a.to_ring(ring) creates a TO_RING node, which may become either a real conversion or a buffer view.
  • a.rescale(moduli) creates a RESCALE node.
  • a.rotate(k) creates a high-level ROTATE node that is lowered later.
  • a.to(device) creates a COPY node.

A Node has an Op, a dtype, zero or more sources in src, and optional operation metadata in arg. Nodes form a DAG, not a tree. If the same polynomial expression is reused, downstream nodes can point at the same producer node.

Ring semantics live on metadata nodes. AS_RING attaches a RingSpec and a domain to storage or computation: Domain.COEFF means coefficient representation, and Domain.NTT means NTT representation.

Creation

Creating a public polynomial from coefficients stores the input coefficient limbs in a Python buffer, then converts them into the ring's Montgomery RNS basis with TO_RNS:

coefficients
  -> packed u32 limb array
  -> BUFFER(device="python")
  -> TO_RNS(ring, limb_count)
  -> Poly

Random constructors such as uniform(), uniform_on_interval(), ternary(), and sparse_ternary() create lazy RANDOM nodes with ring metadata. Realization materializes the random polynomial. Constant constructors are different: Poly.const() creates a symbolic CONST node and attaches ring metadata. Constants stay allocation-free until a later pass needs to lower them into coefficient or NTT form.

Realization

Realization turns lazy polynomial IR into concrete buffers. realize(), tolist(), tobytes(), serialization, decryption, and JIT capture all force realization.

During Poly.realize(), Ren gathers the requested output nodes into a single SINK:

SINK(out0, out1, ...)

The single sink gives the scheduler one root that covers every output requested together. If two output polynomials use the same intermediate, scheduling can see that shared dependency in one graph.

The normal realization path is:

Poly.realize()
  -> Poly._schedule_and_lower()
  -> schedule(SINK(...))
  -> plan_memory(...)
  -> make_exec_units(...)
  -> run_exec_units(...)

After scheduling, Ren records a becomes_map from old lazy outputs to their new buffer-backed outputs. Ren updates existing Poly objects whose graphs overlap the realized computation so they point at realized BUFFER or CONST nodes.

Scheduling

Scheduling turns a high-level polynomial DAG into a list of ScheduledKernel objects. Ren applies this pass sequence:

1. Simplify algebra, ring rules, constants, and copies
2. Lower high-level ring ops such as ROTATE
3. Inject NTT or INTT nodes for domain requirements
4. Simplify domain conversions
5. Decide materialization boundaries
6. Lower BUFFERIZE markers to BUFFER + STORE + AFTER
7. Build the becomes_map for realized outputs
8. Lower automorphisms and ring-specific reduction rules
9. Lower ring constants into explicit coefficient or NTT expressions
10. Turn eligible quotient views into BUFFER_VIEW
11. Indexify elementwise computation with GIDX, INDEX, LOAD, and STORE
12. Kernelize top-level stores into scheduled kernel roots

The scheduler inserts representation changes based on operation requirements. MUL and ROTATE require NTT inputs, RESCALE requires coefficient-domain input, and some TO_RING conversions require coefficient-domain input. After insertion, domain simplification removes round trips such as NTT(INTT(x)).

Some expressions need materialized buffers. Ren marks materialization boundaries for outputs, conversion operations, copies, contiguous requests, and stores. The scheduler first represents a boundary as BUFFERIZE(x), then lowers it into BUFFER, STORE, and AFTER. AFTER keeps dependency ordering explicit without changing the value that flows through the graph.

Ren lowers eligible TO_RING operations to views. If the destination moduli are a contiguous subset of the source basis and the layout/domain constraints allow it, the scheduler creates a BUFFER_VIEW with an element offset. If a view feeds an elementwise compute chain, the scheduler can also fuse the view into indexed loads instead of creating a standalone view kernel.

After indexing, top-level STORE nodes become KERNEL nodes. Concrete buffers are replaced by abstract SLOT parameters in the kernel body. Copy and buffer-view items are represented as lightweight scheduled items instead of normal compute kernels.

Memory planning

Before scheduled kernels become executable units, Ren plans intermediate memory. The memory planner computes first and last use for buffers in the scheduled kernel list. For optimizable intermediate buffers, it simulates a TLSF-style allocator and packs them into one arena per device.

The goal is to reuse storage for temporaries whose live ranges do not overlap. Ren excludes buffers from this optimization when identity matters: final outputs, already allocated buffers, copy and buffer-view buffers, and explicit no-opt buffers. The memory planner makes allocation decisions once from the schedule. Runtime execution reuses arena offsets instead of allocating and freeing each temporary.

Backend lowering and execution

make_exec_units() lowers each scheduled item to an ExecUnit. Compute kernels go through the selected backend: the Python backend interprets the lowered kernel AST with NumPy and Montgomery helpers, while the CUDA backend renders kernel code, compiles it, builds launch specs, and runs it on CUDA queues.

Copies lower to either device transfer or host-mediated copy depending on allocator support. Buffer views lower to validation-only exec units that bind a view to its base.

Execution follows a short loop:

for each exec unit:
  allocate required buffers
  run the unit
  let Python refcounts release intermediates after their final use

JIT capture and replay

ren.jit wraps a Python function whose arguments and outputs contain Poly objects. Inputs can be direct arguments or nested inside ordinary containers and dataclasses such as CKKS Plaintext and Ciphertext objects. Top-level scalar arguments are frozen into the captured plan; if they change on replay, Ren raises instead of running a stale branch.

On a cache miss, JIT realizes inputs, runs the user function under capture, realizes outputs while recording exec units, replans captured intermediate memory, builds input replacement tables, tries to create a backend graph runner, and stores the captured plan. Explicit realize() calls and random polynomial generation inside the jitted body fail early. On a cache hit, JIT swaps the current input buffers into the captured plan and replays it.

Eager lazy execution builds and schedules the graph on each realization. JIT pays that cost once per argument shape/spec and reuses the captured execution plan. See ren.jit replay for examples and failure modes.

Source map