arrayjit 0.3.3 (latest) · OCaml Package

module Lazy = Utils.Lazy

module Debug_runtime = Utils.Debug_runtime

val optimization_level : int Base.ref

type mem_properties =

| Local_only
(*
The array is only needed for a local computation, is allocated on the stack.
*)
| From_context
(*
The array has a copy allocated per-cpu-device, may or may not exist on the host.
*)
| Constant_from_host
(*
The array is read directly from the host.
*)

val mem_properties_of_sexp : Sexplib0.Sexp.t -> mem_properties

val sexp_of_mem_properties : mem_properties -> Sexplib0.Sexp.t

val equal_mem_properties : mem_properties -> mem_properties -> Base.bool

val compare_mem_properties : mem_properties -> mem_properties -> Base.int

val local_only : mem_properties

val from_context : mem_properties

val constant_from_host : mem_properties

val is_local_only : mem_properties -> bool

val is_from_context : mem_properties -> bool

val is_constant_from_host : mem_properties -> bool

val local_only_val : mem_properties -> unit Option.t

val from_context_val : mem_properties -> unit Option.t

val constant_from_host_val : mem_properties -> unit Option.t

module Variants_of_mem_properties : sig ... end

val root_ctx : Gccjit.context option Base.ref

module Tn = Tnode

type context = {

}

val sexp_of_context : context -> Sexplib0.Sexp.t

val unsafe_cleanup : ?unsafe_shutdown:bool -> unit -> Base.unit

val is_initialized : unit -> bool

val initialize : unit -> Base.unit

val finalize : context -> unit

val init : label:Base.string -> context

type ndarray = {

nd : Tn.t;
(*
The original array.
*)
mutable ptr : Gccjit.rvalue Lazy.t;
(*
Pointer to the first value of the associated array.
- if mem = Constant_from_host, the pointer to the first element of the hosted Ndarray,
- if mem = From_context, either a pointer to Ndarray from context.arrays when ~shared:false, or the function parameter when ~shared:true,
- if mem = Local_only, the address of the on-the-stack array.
*)
mem : mem_properties;
dims : Base.int Base.array;
size_in_bytes : Base.int;
num_typ : Gccjit.type_;
(*
The type of the stored values: short (precision Half), float (precision Single), double (precision Double).
*)
is_double : Base.bool;

}

val sexp_of_ndarray : ndarray -> Sexplib0.Sexp.t

type info = {

ctx : Gccjit.context;
func : Gccjit.function_;
traced_store : Low_level.traced_store;
init_block : Gccjit.block;
arrays : (Tn.t, ndarray) Base.Hashtbl.t;

}

val sexp_of_info : info -> Sexplib0.Sexp.t

type param_source =

| Log_file_name
| Param_ptr of Tn.t
| Static_idx of Indexing.static_symbol

val sexp_of_param_source : param_source -> Sexplib0.Sexp.t

type routine = {

info : info;
bindings : Indexing.unit_bindings;
name : Base.string;
result : Gccjit.result;
opt_ctx_arrays : Ndarray.t Base.Map.M(Tn).t Base.option;
params : param_source Base.list;

}

val sexp_of_routine : routine -> Sexplib0.Sexp.t

type code =

| Postponed of {
1. compiled : Low_level.traced_store * Low_level.t;
2. bindings : Indexing.unit_bindings;
3. name : Base.string;
}
| Jitted of routine

val sexp_of_code : code -> Sexplib0.Sexp.t

type gccjit_param = Gccjit.param

type gccjit_lvalue = Gccjit.lvalue

type gccjit_rvalue = Gccjit.rvalue

val sexp_of_gccjit_param : Gccjit.param -> Base.Sexp.t

val sexp_of_gccjit_lvalue : Gccjit.lvalue -> Base.Sexp.t

val sexp_of_gccjit_rvalue : Gccjit.rvalue -> Base.Sexp.t

type ctx_nodes =

| Ctx_arrays of Ndarray.t Base.Map.M(Tn).t Base.ref
| Param_ptrs of (gccjit_param * param_source) Base.list Base.ref

val sexp_of_ctx_nodes : ctx_nodes -> Sexplib0.Sexp.t

val jit_array_offset : 
  Gccjit.context ->
  idcs:Gccjit.rvalue Base.Array.t ->
  dims:int Base.Array.t ->
  Gccjit.rvalue

val zero_out : Gccjit.context -> Gccjit.block -> ndarray -> unit

val get_c_ptr : 
  Gccjit.context ->
  Gccjit.type_ ->
  ('a, 'b, 'c) Bigarray_compat.Genarray.t ->
  Gccjit.rvalue

val prepare_array : 
  debug_log_zero_out:(Gccjit.block -> ndarray -> 'a) ->
  Gccjit.context ->
  (Tn.t, ndarray) Base.Hashtbl.t ->
  (Arrayjit.Low_level.Tn.t, Low_level.traced_array) Base.Hashtbl.t ->
  ctx_nodes ->
  (Gccjit.block -> Gccjit.function_ -> Base.unit) list Base.ref ->
  Tn.t ->
  unit

val prec_to_kind : Ops.prec -> Gccjit.Type.type_kind

val prec_is_double : Ops.prec -> bool

val is_builtin_op : Ops.binop -> bool

val builtin_op : Ops.binop -> Gccjit.binary_op

val arr_debug_name : ndarray -> Base.String.t

val debug_log_zero_out : 
  Gccjit.context ->
  (Gccjit.rvalue * Gccjit.function_ * Gccjit.function_) option Lazy.t ->
  Gccjit.block ->
  ndarray ->
  unit

val debug_log_index : 
  Gccjit.context ->
  (Gccjit.rvalue * Gccjit.function_ * Gccjit.function_) option ->
  Gccjit.block ->
  string ->
  Gccjit.rvalue ->
  unit

val prepare_arrays : 
  Gccjit.context ->
  log_functions:
    (Gccjit.rvalue * Gccjit.function_ * Gccjit.function_) option Lazy.t ->
  (Tn.t, ndarray) Base.Hashtbl.t ->
  (Arrayjit.Low_level.Tn.t, Low_level.traced_array) Base.Hashtbl.t ->
  ctx_nodes ->
  (Gccjit.block -> Gccjit.function_ -> Base.unit) list Base.ref ->
  Low_level.t ->
  unit

val jit_func : 
  name:Base.String.t ->
  opt_ctx_arrays:
    (Tn.t, Ndarray.t, Tn.comparator_witness) Base.Map.t Base.Option.t ->
  Gccjit.context ->
  'a Indexing.bindings ->
  ((Arrayjit.Low_level.Tn.t, Low_level.traced_array) Base.Hashtbl.t
   * Low_level.t) ->
  info * Ndarray.t Base.Map.M(Tn).t * (gccjit_param * param_source) Base.list

val header_sep : Re.re

val maybe_jit : 
  name:Base.string ->
  opt_ctx_arrays:
    (Tn.t, Ndarray.t, Tn.comparator_witness) Base.Map.t Base.Option.t ->
  (Base.unit -> Base.unit) Indexing.bindings ->
  (Low_level.traced_store * Low_level.t) ->
  routine

val maybe_jit_batch : 
  names:Base.string Base.array ->
  opt_ctx_arrays:
    (Tn.t, Ndarray.t, Tn.comparator_witness) Base.Map.t Base.Option.t ->
  (Base.unit -> Base.unit) Indexing.bindings ->
  (Low_level.traced_store * Low_level.t) Base.array ->
  routine Base.Array.t

val jit_routine : 
  context ->
  routine ->
  context
  * (Indexing.static_symbol * Base.int Base.ref) Base.List.t
  * (unit ->
  Tn.work)
  * Base.string

val jit_code : 
  context ->
  code ->
  context
  * (Indexing.static_symbol * Base.int Base.ref) Base.List.t
  * (unit ->
  Tn.work)
  * Base.string

val from_host : context -> Tn.t -> Base.bool

val to_host : context -> Tn.t -> Base.bool

val merge_from_global : 
  unoptim_ll_source:Format.formatter Base.Option.t ->
  ll_source:Format.formatter Base.Option.t ->
  name:Base.string ->
  dst:Arrayjit.Low_level.Tn.t ->
  accum:Ops.binop ->
  src:Ops.voidptr ->
  (Arrayjit.Low_level.Tn.t, Low_level.traced_array) Base.Hashtbl.t
  * Low_level.t

val merge : 
  ?name_prefix:Base.String.t ->
  Tn.t ->
  accum:Ops.binop ->
  src:context ->
  (Base.unit -> Base.unit) Indexing.bindings ->
  routine Base.Option.t

val merge_batch : 
  ?name_prefixes:Base.string Base.array ->
  occupancy:(Tn.t -> src_n:int -> src:context -> Utils.requirement) ->
  Tn.t Base.List.t ->
  accum:Ops.binop ->
  srcs:context Base.array ->
  (Base.unit -> Base.unit) Indexing.bindings ->
  (Tn.t, routine option Base.Array.t) Base.Hashtbl.t

package arrayjit