docs menu

Module Tolk.Jit

Captured schedules.

A captured schedule is every call a function makes, in order, over the buffers it was given. Lowering it once gives a schedule that runs again on other buffers: its inputs become parameters, which each run binds, its intermediates share a few arenas, and its kernels are compiled. The engine links the result once and runs it on each call.

#val width : ?beam:int -> unit -> int

width ~beam () is the width jit_lower ~beam searches at: beam, or else the setting JITBEAM's, or else BEAM's.

#val shaping : ?beam:int -> unit -> (string * string) list

shaping ~beam () is Setting.shaping as jit_lower ~beam compiles under it: with BEAM at the width the lowering searches, and JITBEAM, which the lowering has read by then, unset (Setting.shaping_under). Two lowerings whose shaping is equal compile alike, whatever BEAM and JITBEAM were. It binds no setting.

val jit_lower : 
  ?beam:int ->
  ?search:(int -> Postrange.Scheduler.t -> Postrange.Scheduler.t) ->
  targets:(string -> Helpers.Target.t) ->
  held_bufs:Ops.t list ->
  inputs:Ops.t list ->
  input_vars:(string * int) list ->
  Ops.t ->
  Ops.t

jit_lower ~beam ~search ~targets ~held_bufs ~inputs ~input_vars linear is the captured schedule linear, an Op.t.Linear of calls, ready to link:

the ith buffer of inputs is replaced wherever linear reaches it by the parameter of slot i (Shape.param), of its type, device, size, alignment and phase (Shape.storage_phase);
each variable named in input_vars, the variables of the capture, is unbound wherever linear binds it: a run binds it. A variable bound under another name keeps its binding;
its buffers are placed in arenas (Memory.memory_plan_rewrite), except held_bufs, whose contents outlive a run, such as the buffers a caller keeps or that hold constants;
each buffer not in held_bufs, which holds only what one run computes, becomes storage of the program (Shape.alloc) of its size, type and device, which the program's link allocates;
it is compiled (Realize.compile_linear, with search and targets), with the beam width beam, or else Setting.jitbeam's, or else Setting.beam's, as BEAM, and JITBEAM unset.

Raises as Realize.compile_linear does.

#val source_digest : string

source_digest is the digest of this library's sources, which a key of what they make depends on.