Module Arrayjit.Low_level

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

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

Low-level representation

type t =
  1. | Noop
  2. | Comment of Base.string
  3. | Staged_compilation of Base.unit -> Base.unit
  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. llv : float_t;
    4. mutable debug : Base.string;
    }
  8. | Set_local of scope_id * float_t

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

and float_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_global of Ops.global_identifier * Indexing.axis_index Base.array Base.option
  4. | Get of Tnode.t * Indexing.axis_index Base.array
  5. | Ternop of Ops.ternop * float_t * float_t * float_t
  6. | Binop of Ops.binop * float_t * float_t
  7. | Unop of Ops.unop * float_t
  8. | Constant of Base.float
  9. | Embed_index of Indexing.axis_index
val sexp_of_t : t -> Sexplib0.Sexp.t
val sexp_of_float_t : float_t -> Sexplib0.Sexp.t
val equal : t -> t -> Base.bool
val equal_float_t : float_t -> float_t -> Base.bool
val compare : t -> t -> Base.int
val compare_float_t : float_t -> float_t -> Base.int
val apply_op : Ops.op -> float_t Base.array -> float_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

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;
}
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. mutable computations : (Indexing.axis_index Base.array Base.option * t) Base.list;
    (*

    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.

    *)
  3. assignments : Base.int Base.array Base.Hash_set.t;
  4. accesses : (Base.int Base.array, visits) Base.Hashtbl.t;
  5. mutable zero_initialized : Base.bool;
  6. mutable zeroed_out : Base.bool;
  7. mutable read_before_write : Base.bool;
    (*

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

    *)
  8. mutable read_only : Base.bool;
  9. 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.

    *)
}
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 optimized = {
  1. traced_store : traced_store;
  2. llc : t;
  3. merge_node : Tnode.t Base.option;
}
val sexp_of_optimized : optimized -> Sexplib0.Sexp.t
val optimize : unoptim_ll_source:Stdlib.Format.formatter Base.option -> ll_source:Stdlib.Format.formatter Base.option -> name:Base.string -> Indexing.static_symbol Base.list -> t -> optimized
val input_and_output_nodes : optimized -> (Base.Set.M(Arrayjit.Tnode).t * Base.Set.M(Arrayjit.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 fprint_function_header : ?name:Base.string -> ?static_indices:Indexing.static_symbol Base.list -> Base.unit -> Stdlib.Format.formatter -> Base.unit
val get_ident_within_code : ?no_dots:Base.bool -> t Base.array -> Tnode.t -> Base.string
val fprint_cstyle : ?name:Base.string -> ?static_indices:Indexing.static_symbol Base.list -> Base.unit -> Stdlib.Format.formatter -> t -> Base.unit

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

val fprint_hum : ?name:Base.string -> ?static_indices:Indexing.static_symbol Base.list -> Base.unit -> Stdlib.Format.formatter -> t -> Base.unit

Adheres more to the %cd syntax, does not output implicit type casts.