Skip to content

backend

Structural typing for memory-planner backends.

The codegen pipeline interacts with the memory planner through a narrow contract:

  • plan(model, memory_constraints, tensor_constraints) -> MemoryPlan

Concrete planners may live in third-party packages and need not inherit from :class:AirMemoryPlanner (the bundled abstract base class). To make backend swapping a typing-level concept rather than an inheritance one, this module defines :class:PlannerBackend as a runtime-checkable :class:typing.Protocol. Any class that exposes a matching plan signature (and optionally NAME) satisfies it.

The codegen consumes a planner backend through this protocol; the :class:MemoryPlan it returns is then projected into the render-side DTO (:mod:helia_aot.aot.render_plan) before reaching the templates. That two-layer split is the decoupling: planner backends own allocation policy, the DTO owns emit-ready facts.

Classes

PlannerBackend

Structural contract for a memory-planner backend.

A backend is any object with a plan method matching the signature below. The bundled :class:AirMemoryPlanner subclasses satisfy this protocol; third-party planners that prefer not to inherit from the abstract base can satisfy it duck-typed.

The protocol intentionally requires only plan: a registry name is useful for CLI-driven configs (see :data:MEMORY_PLANNER_REGISTRY) but backends instantiated and injected directly need not advertise one.

Functions

plan 
plan(model: 'AirModel', memory_constraints: list['MemoryConstraint'] | None = None, tensor_constraints: list['AttributeRuleset'] | None = None) -> 'MemoryPlan'

Produce a :class:MemoryPlan for model.

Implementations must populate every field that the render DTO builder relies on:

  • arena_usages for scratch arenas
  • persistent_arenas for whole-program-live arenas
  • constant_arenas for read-only / staged-constant arenas
  • tensor_allocs[tensor.id] for every tensor referenced by the model, with a populated binding (role, memory, offset, source_memory)