package arrayjit

  1. Overview
  2. Docs
package arrayjit
Legend:
Page
Library
Module
Module type
Parameter
Class
Class type
Source

Module Arrayjit.Low_levelSource

Sourcemodule Lazy = Utils.Lazy

The code for operating on n-dimensional arrays.

Sourcemodule Nd = Ndarray
Sourcemodule Tn = Tnode
Sourcemodule Debug_runtime = Utils.Debug_runtime
Sourceval _get_local_debug_runtime : unit -> (module Minidebug_runtime.Debug_runtime)
Sourcemodule Scope_id : sig ... end
Sourcetype scope_id = Scope_id.t = {
  1. tn : Tn.t;
  2. scope_id : Base.int;
}
Sourceval sexp_of_scope_id : scope_id -> Sexplib0.Sexp.t
Sourceval equal_scope_id : scope_id -> scope_id -> Base.bool
Sourceval compare_scope_id : scope_id -> scope_id -> Base.int

*** Low-level representation.

Sourceval get_scope : Tn.t -> scope_id
Sourcetype 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 Tn.t
  7. | Set of {
    1. tn : Tn.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.

Sourceand 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 Tn.t * Indexing.axis_index Base.array
  5. | Binop of Ops.binop * float_t * float_t
  6. | Unop of Ops.unop * float_t
  7. | Constant of Base.float
  8. | Embed_index of Indexing.axis_index
Sourceval sexp_of_t : t -> Sexplib0.Sexp.t
Sourceval sexp_of_float_t : float_t -> Sexplib0.Sexp.t
Sourceval equal : t -> t -> Base.bool
Sourceval equal_float_t : float_t -> float_t -> Base.bool
Sourceval compare : t -> t -> Base.int
Sourceval compare_float_t : float_t -> float_t -> Base.int
Sourceval binop : op:Ops.binop -> rhs1:float_t -> rhs2:float_t -> float_t
Sourceval unop : op:Ops.unop -> rhs:float_t -> float_t
Sourceval flat_lines : t Base.List.t -> t Base.List.t
Sourceval unflat_lines : t list -> t
Sourceval comment_to_name : string -> string

*** Optimization ***

Sourcetype 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;
}
Sourceval virtualize_settings : virtualize_settings
Sourcetype visits =
  1. | Visits of Base.int
  2. | Recurrent
    (*

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

    *)
Sourceval visits_of_sexp : Sexplib0.Sexp.t -> visits
Sourceval sexp_of_visits : visits -> Sexplib0.Sexp.t
Sourceval equal_visits : visits -> visits -> Base.bool
Sourceval visits : Base.int -> visits
Sourceval recurrent : visits
Sourceval is_visits : visits -> bool
Sourceval is_recurrent : visits -> bool
Sourceval visits_val : visits -> Base.int Option.t
Sourceval recurrent_val : visits -> unit Option.t
Sourcemodule Variants_of_visits : sig ... end
Sourcetype traced_array = {
  1. tn : Tn.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;
    (*

    For dynamic indexes, we take a value of 0. This leads to an overestimate of visits, which is safe.

    *)
  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.

    *)
}
Sourceval sexp_of_traced_array : traced_array -> Sexplib0.Sexp.t
Sourceval partition_tf_with_comment : t Base.Array.t -> f:(t -> bool) -> t Base.Array.t * t Base.Array.t
Sourceval visit : is_assigned:Base.bool -> visits option -> visits
Sourceval is_constexpr_comp : (Tn.t, traced_array) Base.Hashtbl.t -> float_t -> Base.bool
Sourceval is_scalar_dims : Tn.t -> bool
Sourceval visit_llc : (Tn.t, traced_array) Base.Hashtbl.t -> merge_node_id:Base.int option Base.ref -> (Indexing.symbol, Tn.t) Base.Hashtbl.t -> max_visits:int -> t -> unit
Sourceval check_and_store_virtual : traced_array -> Indexing.static_symbol Base.List.t -> t -> unit
Sourceval optimize_integer_pow : bool Base.ref
Sourceval unroll_pow : base:float_t -> exp:Base.int -> float_t
Sourceval cleanup_virtual_llc : (Indexing.symbol, Tn.t) Base.Hashtbl.t -> static_indices:Indexing.static_symbol Base.List.t -> t -> t
Sourceval substitute_float : var:float_t -> value:float_t -> float_t -> float_t
Sourceval substitute_proc : var:float_t -> value:float_t -> t -> t
Sourceval simplify_llc : t -> t
Sourcetype traced_store = (Tn.t, traced_array) Base.Hashtbl.t
Sourceval sexp_of_traced_store : traced_store -> Sexplib0.Sexp.t
Sourcetype optimized = {
  1. traced_store : traced_store;
  2. llc : t;
  3. merge_node : Tn.t Base.option;
}
Sourceval sexp_of_optimized : optimized -> Sexplib0.Sexp.t
Sourceval code_hum_margin : int Base.ref
Sourceval pp_comma : Format.formatter -> unit -> unit
Sourceval pp_symbol : Format.formatter -> Indexing.symbol -> unit
Sourceval pp_static_symbol : Format.formatter -> Indexing.static_symbol -> unit
Sourceval fprint_function_header : ?name:string -> ?static_indices:Indexing.static_symbol list -> unit -> Format.formatter -> unit
Sourceval get_ident_within_code : ?no_dots:bool -> t Base.Array.t -> Tn.t -> Base.String.t
Sourceval fprint_hum : ?name:string -> ?static_indices:Indexing.static_symbol list -> unit -> Format.formatter -> t -> unit
Sourceval loop_over_dims : int Base.Array.t -> body:(Indexing.axis_index Base.Array.t -> t) -> t