Skip to content

Scheduling

Scheduling turns a high-level Node DAG into executable kernel work. The entry point is schedule(SINK(...)), and the pass sequence lives in ren/engine/schedule.py. Read Engine overview or Execution pipeline first if you want the wider realization path before the scheduler details.

Pass shape

Ren runs these scheduler passes:

Step Purpose
Simplify Algebra, ring rules, constants, and copies.
Lower high-level ops Lower high-level ring ops such as ROTATE.
Inject domains Insert NTT or INTT where an op requires a domain.
Simplify domains Remove redundant conversion round trips.
Plan materialization Choose nodes that need buffers.
Lower stores Convert BUFFERIZE(x) into BUFFER, STORE, and AFTER.
Build becomes_map Map lazy outputs to realized buffer-backed outputs.
Lower ring ops Enforce ring-specific reductions and conversions.
Lower constants Give ring constants explicit coefficient or NTT semantics.
Late views Turn eligible TO_RING nodes into BUFFER_VIEW.
Indexify Convert elementwise work into GIDX, INDEX, LOAD, STORE.
Kernelize Wrap top-level stores into scheduled kernels.

Each pass fixes representation decisions needed by later passes. Domain injection runs before materialization, and constant lowering waits until the graph has explicit ring/domain context.

Trace scheduler rewrites

apply_rewrites(...) accepts a trace callback. This example records a small summary of pass-level replacements.

from ren.engine.schedule import apply_rewrites
from ren.ir.node import Node
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([4, 3, 2, 1, 0, 0, 0, 0], ring)

expr = (x + y) * 2
root = Node.sink(expr.node)

events: list[tuple[str | None, str, str]] = []


def trace(pass_name: str | None, before: Node, after: Node) -> None:
    if before is after:
        return
    events.append((pass_name, before.op.name, after.op.name))


lowered, becomes_map = apply_rewrites(root, trace=trace)

print("initial root:", root.op.name)
print("lowered root:", lowered.op.name)
print("materialized outputs:", [node.base.op.name for node in becomes_map.values()])
print("rewrite events:", len(events))
for pass_name, before, after in events[:8]:
    print(f"{pass_name}: {before} -> {after}")

initial root: SINK
lowered root: SINK
materialized outputs: ['BUFFER']
rewrite events: 90
inject_domain: ADD -> NTT
inject_domain: AS_RING -> NTT
inject_domain: MUL -> MUL
inject_domain: SINK -> SINK
domain: NTT -> AS_RING
domain: MUL -> MUL
domain: SINK -> SINK
bufferize: MUL -> BUFFERIZE

The example schedules (x + y) * 2. The scalar 2 enters as a ring constant. Ren treats this as a normal MUL, so domain injection gives both operands NTT-domain semantics before constant lowering turns 2 into its NTT broadcast form. The bufferize: MUL -> BUFFERIZE row means the final product was marked as needing a concrete output buffer.

The trace stays compact by design. To debug a scheduling issue, filter by pass name first, then inspect the graph before and after that pass.

Materialization

Ren materializes values that need buffer identity. Common boundaries are:

  • final outputs;
  • conversion operations;
  • explicit contiguous() requests;
  • stores and copies;
  • views over realized sources that must remain addressable.

The AFTER node is a dependency edge. It keeps ordering explicit while letting the value continue to behave like the buffer it wraps.

Kernelization

Kernelization replaces concrete buffers in a store body with abstract SLOT parameters. The scheduled item keeps the lowered AST plus the buffers in slot order.

Compute kernels use this shape:

SINK(STORE(INDEX(SLOT(0), GIDX), value))

Copy and buffer-view work are scheduled as lightweight non-compute units.

Scheduler terms

Term Meaning
SINK Single root node used to schedule one or more requested outputs together.
BUFFERIZE Temporary marker saying a value needs buffer identity.
STORE Store a computed value into a buffer.
AFTER Dependency edge that preserves ordering while returning the wrapped value.
GIDX Global element index used after elementwise work is indexified.
INDEX Read a logical element from a node at an index.
LOAD Load from a buffer slot during indexed lowering.
SLOT Abstract kernel parameter that replaces a concrete buffer in the scheduled AST.
KERNEL Scheduled kernel wrapper around a lowered store/root.

Files to read