- `reinitialize` a backend,

- `finalize` a context (freeing all of its arrays that don't come from its parent context).

We use keys of the `Low_level.traced_store` containers assuming that they are precisely the tensor nodes used in the compiled code -- and the `Virtual` nodes are the ones optimized-away. The context can contain nodes from the parent context corresponding to tensors only needed by parent or ancestor context's computations. The `get_ident` function (e.g. provided by `C_syntax`) returns a human-readable identifier that's un-ambiguous in the context of the compiled code (shared within `compile_batch`).

type code [@@deriving sexp_of]

type code_batch [@@deriving sexp_of]

val compile_batch :

names:string option array ->

Indexing.unit_bindings ->

Low_level.optimized option array ->

type code [@@deriving sexp_of]

type code_batch [@@deriving sexp_of]

val compile_batch :

?names:string array ->

?occupancy:(name:string -> src_n:int -> bool) ->

Indexing.unit_bindings ->

Assignments.comp array ->

code_batch

(Backend : Lowered_no_device_backend) : Lowered_backend = struct

include Backend

[@@deriving sexp_of]

include Add_scheduler (Backend)

let link context (code : code) ctx_arrays =

let runner_label = get_name context.stream in

let merge_buffer = context.stream.merge_buffer in

let schedule =

Task.enschedule ~schedule_task ~get_stream_name:get_name context.stream to_schedule

in

let link_batch context (code_batch : code_batch) ctx_arrays =

let runner_label = get_name context.stream in

let merge_buffer = context.stream.merge_buffer in

let bindings, schedules =

| Some proc ->

let ctx_arrays = Option.value_exn ctx_arrays.(i) in

let bindings', to_schedule =

}

[@@deriving sexp_of]

let name, lowered = lower_assignments ?name bindings comp.Assignments.asgns in

let from_prior_context =

Set.diff (Assignments.context_nodes ~use_host_memory comp.asgns) comp.embedded_nodes

in

{ from_prior_context; name; lowered; code; expected_merge_node = lowered.Low_level.merge_node }

let names, lowereds =

lower_batch_assignments ?names ?occupancy bindings

@@ Array.map comps ~f:(fun c -> c.Assignments.asgns)

in

let from_prior_context =

from_prior_context_batch ~use_host_memory

@@ Array.mapi lowereds ~f:(fun i -> Option.map ~f:(fun _ -> comps.(i)))

| _ -> ("(" ^ typ_of_prec to_ ^ ")(", ")")

(* TODO: The following link seems to claim it's better to expand into loops than use memset.

let module Syntax = C_syntax.C_syntax (C_syntax_config (struct

let ptx = cuda_to_ptx ~name @@ Buffer.contents b in

{ traced_store; ptx; params; bindings; name }

let module Syntax = C_syntax.C_syntax (C_syntax_config (struct

let procs = Array.filter_map lowereds ~f:(Option.map ~f:(fun lowereds -> (lowereds, None)))

end)) in

let sexp_of_code_batch _code_batch = failwith "Backend missing -- install the corresponding library"

failwith "Backend missing -- install the corresponding library"

failwith "Backend missing -- install the corresponding library"

let link _context _code = failwith "Backend missing -- install the corresponding library"

Tensor.default_grad_prec := grad_prec;

Utils.settings.output_debug_files_in_build_directory <- true;

(* This will only log from routines if log-level is high enough. *)

Rand.init (* seed *) 0;

let hid_dim_1 = 16 in

let hid_dim_2 = 8 in

let hid_dim_3 = 4 in

(* TINY for debugging: *)

(* let hid_dim = 2 in *)

let data_len = 3 * 5 * 1024 in

(* TINY for debugging: *)

(* let data_len = 3 * 4 in *)

let flat_len = data_len / 2 in

(* Note: [minibatch_size = batch_size / num_streams] is the actual per-device batch used. *)

(* let epochs = 100 in *)

(* let epochs = 50 in *)

(* TINY for debugging: *)

(* let epochs = 2 in *)

(* let epochs = 1 in *)

(* let init_lr = 0.1 in *)

Stdlib.Format.printf "Initial backend global debug info: %a\n%!" Sexp.pp_hum

@@ Backend.get_global_debug_info ();

let per_batch_callback ~at_batch:_ ~at_step:_ ~learning_rate:_ ~batch_loss:_ ~epoch_loss:_ =

if Option.is_none !start_time then start_time := Some (Time_now.nanoseconds_since_unix_epoch ())

in

(* Tn.print_accessible_headers (); *)

let per_epoch_callback ~at_step ~at_epoch ~learning_rate ~epoch_loss =

in

Backend.initialize Train.BT.Most_parallel_streams;

Stdio.print_endline "\n******** mlp_result **********";

Tensor.print_tree ~with_id:true ~with_grad:false ~depth:9 model_result;

Stdio.printf "\n********\n%!";

let callback (x, y) = Float.((infer_callback [| x; y |]).(0) >= 0.) in

in

Stdio.printf "\nHalf-moons scatterplot and decision boundary:\n%!";

Stdio.printf "\nBatch Log-loss:\n%!";

let plot_loss =

let open PrintBox_utils in

}

in

Stdio.printf "\n\n%!";

Stdlib.Format.printf "Final backend global debug info: %a\n%!" Sexp.pp_hum

@@ Backend.get_global_debug_info ();

result

]))))))

let _cuda_parallel_benchmarks =

List.concat_map

~f:(fun batch_size ->

List.concat_map [ (* 1; 3; *) 7 (* *) ] ~f:(fun seed ->

List.concat_map [ (* CDSL.double; *) CDSL.single (* ; CDSL.half *) ]

~f:(fun value_prec ->

[

data.tracking <- Some (Cu.Delimited_event.record context.stream.cu_stream));

]}

Old context finalizer from the cuda backend:

{[

let%track3_sexp finalize (ctx : context) : unit =

Array.mapi ctxs ~f:(fun dst_n ctx ->

if occupancy_dst ~dst_n then

snd

else [||])

in

(* We can cache scheduling, because merging and copying does not depend on static indexing. *)

type example_train_result = {

inputs : Tensor.t;

set_hosted learning_rate.value;

let sgd = sgd_update ~learning_rate ~weight_decay update in

let grad_updates = Array.map contexts ~f:(fun ctx -> Backend.link ctx grad_update) in

let sgd_update = to_routine (module Backend) grad_updates.(0).context bindings sgd in

Tensor.log_debug_info ~from_log_level:2 inputs;