Module Ir.Low_level

A for-loop-based array language and backend-agnostic optimization

Global references

module Scope_id : sig ... end
type scope_id = Scope_id.t = {
  1. tn : Tnode.t;
  2. scope_id : Base.int;
}
val sexp_of_scope_id : scope_id -> Sexplib0.Sexp.t
val equal_scope_id : scope_id -> scope_id -> Base.bool
val hash_fold_scope_id : Ppx_hash_lib.Std.Hash.state -> scope_id -> Ppx_hash_lib.Std.Hash.state
val hash_scope_id : scope_id -> Ppx_hash_lib.Std.Hash.hash_value
val compare_scope_id : scope_id -> scope_id -> Base.int
val get_scope : Tnode.t -> scope_id

Low-level representation

type t =
  1. | Noop
  2. | Comment of Base.string
  3. | Staged_compilation of Base.unit -> PPrint.document
  4. | Seq of t * t
  5. | For_loop of {
    1. index : Indexing.symbol;
    2. from_ : Base.int;
    3. to_ : Base.int;
    4. body : t;
    5. trace_it : Base.bool;
    }
  6. | Zero_out of Tnode.t
  7. | Set of {
    1. tn : Tnode.t;
    2. idcs : Indexing.axis_index Base.array;
    3. llsc : scalar_t;
    4. mutable debug : Base.string;
    }
  8. | Set_from_vec of {
    1. tn : Tnode.t;
    2. idcs : Indexing.axis_index Base.array;
    3. length : Base.int;
    4. vec_unop : Ops.vec_unop;
    5. arg : scalar_arg;
    6. mutable debug : Base.string;
    }
  9. | Set_local of scope_id * scalar_t
  10. | Declare_local of {
    1. id : scope_id;
    2. needs_init : Base.bool;
    }

Cases: t -- code, scalar_t -- single number at some precision.

and scalar_t =
  1. | Local_scope of {
    1. id : scope_id;
    2. body : t;
    3. orig_indices : Indexing.axis_index Base.array;
    }
  2. | Get_local of scope_id
  3. | Get of Tnode.t * Indexing.axis_index Base.array
  4. | Get_dynamic of {
    1. tn : Tnode.t;
      (*

      The gathered table; treated as a read of tn, like Get.

      *)
    2. idcs : Indexing.axis_index Base.array;
      (*

      Static everywhere except dyn_axis.

      *)
    3. dyn_axis : Base.int;
      (*

      Which idcs slot is replaced by dyn_value at codegen time.

      *)
    4. dyn_value : scalar_arg;
      (*

      Integer-valued index spliced into the row-major offset at dyn_axis. gh-343: produced only by rewrite_one_hot_reductions; never escapes low-level / backend codegen.

      *)
    }
  5. | Get_merge_buffer of Tnode.t * Indexing.axis_index Base.array
  6. | Ternop of Ops.ternop * scalar_arg * scalar_arg * scalar_arg
  7. | Binop of Ops.binop * scalar_arg * scalar_arg
  8. | Unop of Ops.unop * scalar_arg
  9. | Constant of Base.float
  10. | Constant_bits of Base.int64
    (*

    Direct bit representation, primarily for uint4x32

    *)
  11. | Embed_index of Indexing.axis_index
and scalar_arg = scalar_t * Ops.prec

The argument precision is preserved in heterogeneous precision operation arguments, and is ignored (overridden) in homogeneous precision operations.

val sexp_of_t : t -> Sexplib0.Sexp.t
val sexp_of_scalar_t : scalar_t -> Sexplib0.Sexp.t
val sexp_of_scalar_arg : scalar_arg -> Sexplib0.Sexp.t
val equal : t -> t -> Base.bool
val equal_scalar_t : scalar_t -> scalar_t -> Base.bool
val equal_scalar_arg : scalar_arg -> scalar_arg -> Base.bool
val compare : t -> t -> Base.int
val compare_scalar_t : scalar_t -> scalar_t -> Base.int
val compare_scalar_arg : scalar_arg -> scalar_arg -> Base.int
val scalar_precision : scalar_t -> Ops.prec
val apply_op : Ops.op -> scalar_t Base.array -> scalar_t
val flat_lines : t Base.list -> t Base.list
val unflat_lines : t Base.list -> t
val loop_over_dims : Base.int Base.array -> body:(Indexing.axis_index Base.array -> t) -> t
val unroll_dims : Base.int Base.array -> body:(Indexing.axis_index Base.array -> offset:Base.int -> t) -> t
val loop_over_padding_region : dims:Base.int Base.array -> padding:Ops.axis_padding Base.array -> body:(Indexing.axis_index Base.array -> t) -> t

Generate loops that iterate only over the padding margins of a tensor. For dimensions with padding, generates separate loops for left margin, middle (recursing), and right margin. The middle region continues recursing to find padding in other dimensions.

Optimization

type virtualize_settings = {
  1. mutable enable_device_only : Base.bool;
  2. mutable max_visits : Base.int;
  3. mutable max_tracing_dim : Base.int;
  4. mutable inline_scalar_constexprs : Base.bool;
  5. mutable inline_simple_computations : Base.bool;
  6. mutable inline_complex_computations : Base.bool;
}
val virtualize_settings : virtualize_settings
type visits =
  1. | Visits of Base.int
  2. | Recurrent
    (*

    A Recurrent visit is when there is an access prior to any assignment in an update.

    *)
val sexp_of_visits : visits -> Sexplib0.Sexp.t
val visits_of_sexp : Sexplib0.Sexp.t -> visits
val equal_visits : visits -> visits -> Base.bool
val visits : Base.int -> visits
val recurrent : visits
val is_visits : visits -> Base.bool
val is_recurrent : visits -> Base.bool
val visits_val : visits -> Base.int Base.option
val recurrent_val : visits -> Base.unit Base.option
module Variants_of_visits : sig ... end
type traced_array = {
  1. tn : Tnode.t;
  2. assignments : Base.int Base.array Base.Hash_set.t;
  3. accesses : (Base.int Base.array, visits) Base.Hashtbl.t;
  4. mutable zero_initialized_by_code : Base.bool;
  5. mutable zeroed_out : Base.bool;
  6. mutable read_before_write : Base.bool;
    (*

    The node is read before it is written (i.e. it is recurrent).

    *)
  7. mutable read_only : Base.bool;
    (*

    Surprisingly, the notions of read-only and of constant memory mode come apart: small hosted constants are not read-only because they are initialized on devices by being assigned to; and a volatile memory mode is read-only from the devices' perspective.

    *)
  8. mutable is_scalar_constexpr : Base.bool;
    (*

    True only if the tensor node has all axes of dimension 1, is either zeroed-out or assigned before accessed, is assigned at most once, and from an expression involving only constants or tensor nodes that were at the time is_scalar_constexpr.

    *)
  9. mutable is_accessing : Base.bool;
    (*

    False only if the tensor node is built from index embeddings and scalar constant expressions.

    *)
  10. mutable is_complex : Base.bool;
    (*

    True only if the tensor node is built from a genuinely complex scalar computation (one that accesses other non-constexpr computations). Sharing a loop symbol with another tensor does not, by itself, make a node complex (see #134).

    *)
  11. mutable prefers_virtual_one_hot : Base.bool;
    (*

    True when at least one setter for this tensor is a one-hot selector assignment, i.e. a Cmpeq between the embedded range iterator and a loop-variable-free expression. When has_non_one_hot_setter is false this tensor is exempt from the visit-count Never_virtual rule (task-73617488).

    *)
  12. mutable has_non_one_hot_setter : Base.bool;
    (*

    True when at least one setter is NOT a one-hot selector (including Set_from_vec). A tensor with prefers_virtual_one_hot && not has_non_one_hot_setter is the candidate for the one-hot virtualizer exemption.

    *)
  13. mutable is_range_producer : Base.bool;
    (*

    True when at least one Set assigns this tensor from a bare Embed_index scalar, i.e. the tensor is a Range_over_offsets producer. Used by the indirect arm of is_one_hot_selector_assignment to prove that a Get(rtn, [k]) will inline to Embed_index k rather than arbitrary values (task-73617488).

    *)
}
val sexp_of_traced_array : traced_array -> Sexplib0.Sexp.t
val get_node : (Tnode.t, traced_array) Base.Hashtbl.t -> Tnode.t -> traced_array
val optimize_integer_pow : Base.bool Base.ref
type traced_store = (Tnode.t, traced_array) Base.Hashtbl.t
val sexp_of_traced_store : traced_store -> Sexplib0.Sexp.t
type optimize_ctx = {
  1. computations : (Tnode.t, (Indexing.axis_index Base.array Base.option * t) Base.list) Base.Hashtbl.t;
    (*

    The computations (of the tensor node) are retrieved for optimization just as they are populated, so that the inlined code corresponds precisely to the changes to the arrays that would happen up till that point. Within the code blocks paired with an index tuple, all assignments and accesses must happen via the index tuple; if this is not the case for some assignment, the node cannot be virtual. Currently, we only allow for-loop symbols in assignment indices of virtual nodes.

    *)
}
val sexp_of_optimize_ctx : optimize_ctx -> Sexplib0.Sexp.t
type optimized = {
  1. traced_store : traced_store;
  2. optimize_ctx : optimize_ctx;
  3. llc : t;
  4. merge_node : Tnode.t Base.option;
}
val sexp_of_optimized : optimized -> Sexplib0.Sexp.t
val optimize : optimize_ctx -> unoptim_ll_source:(PPrint.document -> Base.unit) Base.option -> ll_source:(PPrint.document -> Base.unit) Base.option -> name:Base.string -> Indexing.static_symbol Base.list -> t -> optimized
val reads_scope_before_set : scope_id -> t -> Base.bool

reads_scope_before_set id body returns true if id is read (via Get_local) before the first definitely-executed Set_local id in body. Use this at code-generation time to decide whether a Local_scope or Declare_local declaration needs a zero initializer.

val rewrite_one_hot_reductions : t -> t

gh-343: rewrites the narrow one-hot embedding pattern -- an Add reduction over a loop variable k whose body selects an embedding-table row via k == index_expr (a logical one-hot) -- into a guarded dynamic gather (Get_dynamic) that reads the table row at index_expr directly, with an in-range guard returning 0 out of [0, vocab_size) to preserve the one-hot semantics. Unmatched or unsupported reductions are left unchanged. Called internally by optimize between simplify_llc and eliminate_common_subexpressions; exposed for testing.

val eliminate_common_subexpressions : t -> t

Eliminates common subexpressions within each statement's scalar expression tree. Replaces duplicate Local_scope nodes (structurally identical modulo scope_id) with Get_local references to the first occurrence. Called internally by optimize; exposed for testing.

val hoist_cross_statement_cse : t -> t

Hoists shared Local_scope computations from sibling statements to the enclosing scope. When two or more sibling statements share an alpha-equivalent Local_scope node, the computation is extracted as a Declare_local + body preceding the first user, and all occurrences are replaced with Get_local.

val input_and_output_nodes : optimized -> (Base.Set.M(Ir.Tnode).t * Base.Set.M(Ir.Tnode).t) * Tnode.t Base.option

Inputs are the materialized read-only and read-before-write (within the code) non-constant non-merge nodes. They are inputs in a broad sense, as they could be recurrent nodes or parameters. Outputs are all the materialized nodes written-to by the code. The last returned component is the input merge node, if used in the code.

Printing

val code_hum_margin : Base.int Base.ref
val function_header_doc : ?name:Base.string -> ?static_indices:Indexing.static_symbol Base.list -> Base.unit -> PPrint.document
val get_ident_within_code : ?no_dots:Base.bool -> ?blacklist:Base.string Base.list -> t Base.array -> Tnode.t -> Base.string
val to_doc_cstyle : ?name:Base.string -> ?static_indices:Indexing.static_symbol Base.list -> Base.unit -> t -> PPrint.document

Adheres more to the C syntax, outputs implicit type casts.

val to_doc : ?name:Base.string -> ?static_indices:Indexing.static_symbol Base.list -> Base.unit -> t -> PPrint.document

Adheres to the %cd syntax.