docs menu

on this page

Module Postrange.Scheduler

Schedulers.

A scheduler holds a kernel and the optimisations applied to it so far. It is mutable: applying an optimisation rewrites its kernel in place.

#type t

The type for schedulers.

#val v : Ops.t -> Renderer.t -> t

v ast ren schedules the kernel ast, the sink of a kernel (Ops.kernel_info), for ren, with the ranges its reductions and ends close flattened (Simplify.pm_flatten_range). Its applied optimisations are those ast's argument lists.

#val copy : t -> t

copy k is a scheduler of k's kernel and optimisations, which applying an optimisation to either does not change in the other.

#val ast : t -> Ops.t

ast k is k's kernel.

#val ren : t -> Renderer.t

ren k is the renderer k optimises for.

#val applied_opts : t -> Opt.t list

applied_opts k is the optimisations applied to k, in order.

Axes

#val rngs : t -> Ops.t list

rngs k is k's axes.

#val shape_len : t -> int

shape_len k is the number of k's axes.

#val full_shape : t -> Ops.sint list

full_shape k is the size of each of k's axes, simplified.

#val axis_types : t -> Ops.Axis_type.t list

axis_types k is the role of each of k's axes.

#val reduce_ranges : t -> Ops.t list

reduce_ranges k is the ranges k's reductions close, each once.

#val axes_of : ?reduce:bool -> t -> Ops.Axis_type.t list -> int list

axes_of ~reduce k types is the positions of k's axes of a role in types, or of any role if types is [], that are reduce axes if reduce is true and are not if it is false. Without reduce, any.

#val upcast_size : t -> Ops.sint

upcast_size k is the product of the sizes of k's Ops.Axis_type.t.Upcast axes.

#val read_by_bounds : t -> Ops.t list

read_by_bounds k is the ranges that the bound of one of k's ranges reads, as i in a loop j < x[i]. A loop has one count of trips, so none of them can become a thread's lanes.

#val opaque : t -> Ops.t list

opaque k is k's opaque operations, but those of the bodies it calls: its calls (Ops.call), of code in the target's language among them (Ops.custom_function). Each runs once for each point of the ranges its sources read, and takes one pointer or value for each source.

#val read_by_opaque : t -> Ops.t list

read_by_opaque k is the ranges that a source of one of k's opaque operations reads, as i in f(&x[i]) or popcount(x[i]). Each takes one pointer or value for each source, so none of them can become a thread's lanes, which would pass it a vector.

#val unread_by_opaque : t -> Ops.t list

unread_by_opaque k is k's ranges that the sources of one of its opaque operations do not read, as j for f(&x[i]) beside a loop over j. Every thread of a launch runs each such operation, so none of them can take threads, which would run it once in each.

#val threadless : t -> Ops.t list

threadless k is k's ranges that take no threads: those unread_by_opaque, and the weak ranges whose steps share storage with the code outside them (Rewriting).

#val upcastable_dims : ?reduce:bool -> t -> int list

upcastable_dims ~reduce k is the positions of k's global, local and weak axes of a known size that no range's bound reads (read_by_bounds), no opaque operation reads (read_by_opaque) and no register store or loop's condition depends on (Rewriting), that are reduce axes if reduce is true and are not if it is false (default).

#val group_for_reduces : t -> int

group_for_reduces k is the number of k's reduce axes that are warp or local axes.

#val reduceops : t -> Ops.t list

reduceops k is k's reductions, in Ops.backward_slice's order.

#val reduceop : t -> Ops.t option

reduceop k is the first of reduceops k.

#val bufs : t -> Ops.t list

bufs k is k's indexed accesses (Op.t.Index), in the reverse of Ops.backward_slice's order.

#val colored_shape : t -> string

colored_shape k is the size of each of k's axes, right-aligned in four columns and coloured by its role (Ops.Axis_type.color).

Rewriting

A split into lanes (Opt.target.Upcast) runs points of an axis in one thread, which holds one copy of each register: points that write a register cell different values leave one of them. A step of a weak axis reads a register cell only after writing it, or one no step writes (Ops.Axis_type.t.Weak), so an axis takes lanes unless a register store's cells, value or gate, or the condition that repeats a loop (Op.t.Backedge), depends on it in data: a value depends on what its sources depend on, a register load on its offsets and gate alone, since its cell holds what the same step stored, an Op.t.After on its first source, and a reduction on its value but the axes it closes.

Threads have their own registers, but run unordered with the code outside their axis, and an effect that does not depend on a thread axis runs in one of its threads (Ops.Axis_type). So a weak axis takes no threads, global or local, when an effect that depends on it and one that does not access one storage of global or workgroup memory and one of them writes it, or when one that does not depend on it reads a register one depending on it writes, which holds the last step's value. The effects are stores, calls and code in the target's language: each reads what its data loads and the pointers it passes, and writes what it stores into, a call the arguments its body writes (Shape.written_params) and code every pointer it is passed.

#val convert_loop_to_global : t -> unit

convert_loop_to_global k makes global each weak range that an end of k's sink closes and no reduction does, whose bound a launch computes (Ops.known_at_launch), that every buffered value (Op.t.Stage) runs inside and that take threads (threadless), if k's renderer has local indices.

val shift_to : 
  ?top:bool ->
  ?new_rng:Ops.t ->
  t ->
  Ops.t ->
  int ->
  Opt.target ->
  Ops.t * Ops.t

shift_to k r amount target splits the axis r into (r', s): r' of r's size divided by amount, and s, of size amount and target's role, a new range unless new_rng is given. r becomes r' * amount + s, or s * size r' + r' if top (default false), and the ranges reductions and ends close stay flat.

Raises Invalid_argument if amount is not greater than 1, if r's role is not among split_targets target, if r's trips are not known at launch (Ops.known_at_launch), if target is Opt.target.Upcast and r takes no lanes (read_by_bounds, Rewriting), or Opt.target.Local and r takes no threads (threadless), or if amount does not divide r's size.

val apply_opt : 
  ?append_opt:bool ->
  t ->
  Opt.t ->
  (Ops.t list, string) Stdlib.result

apply_opt k opt applies opt to k, records it among k's applied optimisations if append_opt (default true), and is Ok the axes it made: a split's [r'; s] (shift_to), a pad's grown axis, a tensor core's three axes N, M and K, and [] for a swap.

It is Error with the reason if opt does not apply, leaving k as it was unless a tensor core's reduction is found ambiguous after its axes are split. It does not apply when its axis is out of range, a split's target cannot come from its axis's role or its axis's trips are not known at launch (Ops.known_at_launch), as j < 4 (i + 1)'s, its axis takes no lanes for an upcast or no threads for a local split (Rewriting), a split exceeds what the renderer allows (32 lanes of a reduce axis or 16 of another, locals only with local indices, the workgroup's shared memory), an upcast of an axis a range's bound reads (read_by_bounds) or an opaque operation reads (read_by_opaque), a local split of an axis an opaque operation does not read (unread_by_opaque), a local split of a reduce axis where a range's bound reads an axis or whose reduction runs inside an unrolled reduce axis, a pad's axis is not of a constant size, would more than quadruple, or holds an effect a guard cannot make a no-op: a call or a store not through an Op.t.Index; or no tensor core fits, as none does a kernel where a range's bound reads an axis, or one with an opaque operation. Raises Invalid_argument if the shared memory a split needs depends on a symbolic size and cannot be decided.

#val get_optimized_ast : ?name_override:string -> t -> Ops.t

get_optimized_ast ~name_override k is k's kernel with its argument listing the applied optimisations and marking it optimised (optimized). It is named name_override, or else "r" for a kernel that reduces and "E" otherwise, followed by the sizes of its hardware indices and axes, each after an underscore. The name holds no colour.