OCANNL classifies tensor nodes according to their memory properties:

```ocaml

Contexts track (or store) the on-device arrays corresponding to tensor nodes. Contexts form a hierarchy: linking takes a parent context and outputs a child context. Related contexts that use a tensor node must use the same on-device array for the tensor node. If two unrelated contexts are on the same device, i.e. have a common ancestor, and use the same tensor node that is not part of the most recent common ancestor, the behavior is undefined.

During the compilation process, the old context cannot be available when `compile` is handled. Currently, all backends generate context-and-device-independent kernels, that refer to context arrays via parameters.

let is_total ~initialize_neutral ~projections =

initialize_neutral && Indexing.is_bijective projections

let open Utils.Set_O in

let empty = Set.empty (module Tn) in

let one tn =

in

let of_node = function Node rhs -> one rhs | Merge_buffer _ -> empty in

let rec loop = function

module type No_device_buffer_and_copying = sig

include Alloc_buffer with type stream := unit

val get_used_memory : unit -> int

(** Returns (an upper bound of) the memory used for arrays, in bytes. *)

No_device_buffer_and_copying with type buffer_ptr = unit Ctypes.ptr = struct

type buffer_ptr = unit Ctypes.ptr

let sexp_of_buffer_ptr = Ops.sexp_of_voidptr

include Buffer_types (struct

Ctypes_memory_stubs.memcpy

~dst:(Ndarray.get_fatptr_not_managed dst)

~src ~size:(Ndarray.size_in_bytes dst)

module Device_types (Device_config : Device_config) = struct

module type Backend_impl_common = sig

include Buffer

module type Buffer = sig

type buffer_ptr [@@deriving sexp_of]

include module type of Buffer_types (struct

type nonrec buffer_ptr = buffer_ptr [@@deriving sexp_of]

end)

let ordinal_of ctx = ctx.stream.device.ordinal in

let name_of ctx = Backend.(get_name ctx.stream) in

let same_device = ordinal_of dst = ordinal_of src in

then false

else

match Map.find src.ctx_arrays tn with

let%debug3_sexp verify_prior_context ~use_host_memory ~ctx_arrays ~from_prior_context : unit =

Set.iter from_prior_context ~f:(fun tn ->

if

&& not (Option.is_some @@ Map.find ctx_arrays tn)

then raise @@ Utils.User_error ("The linked context lacks node " ^ Tnode.debug_name tn))

}

let%track3_sexp alloc_if_needed (stream : stream) ~key ~data:node ctx_arrays =

[%log Tn.debug_name key];

[%log (key : Tnode.t)];

let default () =

alloc_zero_init_array (Lazy.force key.prec) ~dims:(Lazy.force key.dims) stream

in

let add_new () = Map.add_exn ctx_arrays ~key ~data:(default ()) in

let device = stream.device in

if Tn.known_non_cross_stream key then add_new ()

if Hashtbl.mem device.cross_stream_candidates key then

Map.add_exn ctx_arrays ~key ~data)

if Hashtbl.mem device.owner_stream key then (

if not (equal_stream stream (Hashtbl.find_exn device.owner_stream key)) then

raise

(** The low-level prcedure to compile, and the arrays of the context it will be linked to if not

val logs_to_stdout : bool

val main_kernel_prefix : string

val kernel_prep_line : string

let get_ident =

Low_level.get_ident_within_code ~no_dots:true @@ Array.map B.procs ~f:(fun l -> l.llc)

let pp_zero_out ppf tn =

Stdlib.Format.fprintf ppf "@[<2>memset(%s, 0, %d);@]@ " (get_ident tn) @@ Tn.size_in_bytes tn

}

[@@deriving sexp_of]

let get_global_run_id =

let next_id = ref 0 in

fun () ->

end) =

let procs = Input.procs

let use_host_memory = use_host_memory

let logs_to_stdout = false

let main_kernel_prefix = ""

let%track3_sexp link_compiled ~merge_buffer ~runner_label ctx_arrays (code : procedure) =

let name : string = code.name in

let log_file_name = Utils.diagn_log_file [%string "debug-%{runner_label}-%{code.name}.log"] in

let run_variadic =

[%log_level

Param_2f (get_ptr, merge_buffer, link bs ps Ctypes.(ptr void @-> cs))

| bs, Param_ptr tn :: ps ->

let c_ptr =

in

Param_2 (ref (Some c_ptr), link bs ps Ctypes.(ptr void @-> cs))

in

in

let%diagn_l_sexp work () : unit =

[%log_result name];

Indexing.apply run_variadic ();

if Utils.debug_log_from_routines () then (

Utils.log_trace_tree (Stdio.In_channel.read_lines log_file_name);

type buffer_ptr = Cu.Deviceptr.t

let sexp_of_buffer_ptr ptr = Sexp.Atom (Cu.Deviceptr.string_of ptr)

include Buffer_types (struct

type nonrec buffer_ptr = buffer_ptr [@@deriving sexp_of]

end)

module Device_config = struct

include Backend_buffer

let cuda_properties =

let cache =

in

if not @@ is_initialized () then invalid_arg "cuda_properties: CUDA not initialized";

let cache = Lazy.force cache in

Lazy.force cache.(device.ordinal)

let suggested_num_streams device =

match !global_config with

(* Map.iteri ctx_arrays ~f:(fun ~key ~data:ptr -> if key.Low_level.zero_initialized then

[%log "launching the kernel"];

(if Utils.debug_log_from_routines () then

Utils.add_log_processor ~prefix:log_id_prefix @@ fun _output ->

[%log_block

}

[@@deriving sexp_of]

let gcc_typ_of_prec =

let open Gccjit in

let ptr_typ = Type.pointer num_typ in

let ident = get_ident tn in

let hosted = Tn.is_hosted_force tn 344 in

let ptr =

match (in_ctx, hosted) with

let p = Param.create ctx ptr_typ ident in

param_ptrs := (p, Param_ptr tn) :: !param_ptrs;

Lazy.from_val (RValue.param p)

let addr arr =

Lazy.from_val @@ get_c_ptr ctx num_typ @@ Ctypes.bigarray_start Ctypes_static.Genarray arr

in

| Some (Single_nd arr) -> addr arr

| Some (Double_nd arr) -> addr arr

| None -> assert false)

let arr_typ = Type.array ctx num_typ size_in_elems in

let v = ref None in

let initialize _init_block func = v := Some (Function.local func arr_typ ident) in

let%track3_sexp link_compiled ~merge_buffer ~runner_label ctx_arrays (code : procedure) =

let name : string = code.name in

List.iter code.params ~f:(function

| _ -> ());

let log_file_name = Utils.diagn_log_file [%string "debug-%{runner_label}-%{code.name}.log"] in

let run_variadic =

Param_1 (ref (Some log_file_name), link bs ps Ctypes.(string @-> cs))

| bs, Param_ptr tn :: ps ->

let c_ptr =

in

Param_2 (ref (Some c_ptr), link bs ps Ctypes.(ptr void @-> cs))

| bs, Merge_buffer :: ps ->

type runner

type event

let sexp_of_dev _dev = failwith "Backend missing -- install the corresponding library"

let sexp_of_runner _runner = failwith "Backend missing -- install the corresponding library"

let sexp_of_event _event = failwith "Backend missing -- install the corresponding library"

let get_name _stream = failwith "Backend missing -- install the corresponding library"

let sexp_of_buffer_ptr _buffer_ptr = failwith "Backend missing -- install the corresponding library"

type nonrec buffer = buffer_ptr Backend_intf.buffer