Module Arrayjit.Assignments

The code for operating on n-dimensional arrays.

module Lazy = Utils.Lazy
module Tn = Tnode
val _get_local_debug_runtime : unit -> (module Minidebug_runtime.Debug_runtime)
type buffer =
  1. | Node of Tn.t
  2. | Merge_buffer of Tn.t
val sexp_of_buffer : buffer -> Sexplib0.Sexp.t
type fetch_op =
  1. | Constant of Base.float
  2. | Imported of Ops.global_identifier
  3. | Slice of {
    1. batch_idx : Indexing.static_symbol;
    2. sliced : Tn.t;
    }
  4. | Embed_symbol of Indexing.static_symbol

Resets a array by performing the specified computation or data fetching.

and t =
  1. | Noop
  2. | Seq of t * t
  3. | Block_comment of Base.string * t
    (*

    Same as the given code, with a comment.

    *)
  4. | Accum_ternop of {
    1. initialize_neutral : Base.bool;
    2. accum : Ops.binop;
    3. op : Ops.ternop;
    4. lhs : Tn.t;
    5. rhs1 : buffer;
    6. rhs2 : buffer;
    7. rhs3 : buffer;
    8. projections : Indexing.projections Lazy.t;
    }
  5. | Accum_binop of {
    1. initialize_neutral : Base.bool;
    2. accum : Ops.binop;
    3. op : Ops.binop;
    4. lhs : Tn.t;
    5. rhs1 : buffer;
    6. rhs2 : buffer;
    7. projections : Indexing.projections Lazy.t;
    }
  6. | Accum_unop of {
    1. initialize_neutral : Base.bool;
    2. accum : Ops.binop;
    3. op : Ops.unop;
    4. lhs : Tn.t;
    5. rhs : buffer;
    6. projections : Indexing.projections Lazy.t;
    }
  7. | Fetch of {
    1. array : Tn.t;
    2. fetch_op : fetch_op;
    3. dims : Base.int Base.array Lazy.t;
    }
val sexp_of_fetch_op : fetch_op -> Sexplib0.Sexp.t
val sexp_of_t : t -> Sexplib0.Sexp.t
type comp = {
  1. asgns : t;
  2. embedded_nodes : Base.Set.M(Tn).t;
    (*

    The nodes in asgns that are not in embedded_nodes need to already be in contexts linked with the comp.

    *)
}

Computations based on assignments. Note: the arrayjit library makes use of, but does not produce nor verify the embedded_nodes associated to some given asgns.

val sexp_of_comp : comp -> Sexplib0.Sexp.t
val to_comp : t -> comp
val empty_comp : comp
val get_name_exn : t -> Base.String.t
val is_total : initialize_neutral:Base.bool -> projections:Indexing.projections -> Base.bool
val context_nodes : use_host_memory:'a Base.option -> t -> Tn.t_set

Returns materialized nodes in the sense of Tnode.is_in_context_force. NOTE: it must be called after compilation; otherwise, it will disrupt memory mode inference.

val sequential : t Base.List.t -> t
val sequence : comp Base.List.t -> comp
val to_low_level : t -> Low_level.t
val flatten : t -> t Base__List.t
val is_noop : t -> bool
val get_ident_within_code : ?no_dots:bool -> t -> Tn.t -> Base.String.t
val fprint_hum : ?name:Base.string -> ?static_indices:Indexing.static_symbol Base.list -> unit -> Stdlib.Format.formatter -> t -> unit
val lower : unoptim_ll_source:Stdlib.Format.formatter Base.option -> ll_source:Stdlib.Format.formatter Base.option -> cd_source:Stdlib.Format.formatter option -> name:Base.string -> Indexing.static_symbol Base.list -> t -> Low_level.optimized