- Huge refactoring of backend internal interfaces and API (not repeating same code).

- Built per-tensor-node stream-to-stream synchronization into copying functions.

- Re-introduced whole-device blocking synchronization, which now is just a slight optimization as it also cleans up event book-keeping.

type procedure [@@deriving sexp_of]

val compile_batch :

names:string option array ->

Indexing.unit_bindings ->

Low_level.optimized option array ->

procedure option array

val link_compiled :

merge_buffer:buffer option ref ->

[@@deriving sexp_of]

let compile ~name bindings lowered : code =

{ lowered; proc }

let compile_batch ~names bindings lowereds : code_batch =

{ lowereds; procs }

include Add_scheduler (Backend)

module Tn = Tnode

module C_syntax (B : sig

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

val use_host_memory : bool

val logs_to_stdout : bool

val main_kernel_prefix : string

val kernel_prep_line : string

val typ_of_prec : Ops.prec -> string

val binop_syntax : Ops.prec -> Ops.binop -> string * string * string

val unop_syntax : Ops.prec -> Ops.unop -> string * string

val convert_precision : from:Ops.prec -> to_:Ops.prec -> string * string

end) =

let get_ident =

let in_ctx tn = B.(Tn.is_in_context ~use_host_memory tn)

let pp_zero_out ppf tn =

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

open Indexing.Pp_helpers

let pp_array_offset ppf (idcs, dims) =

(* let compute_array_offset ~idcs ~dims = Array.fold2_exn idcs dims ~init:0 ~f:(fun offset idx dim

(B.typ_of_prec (Lazy.force tn.Tn.prec) ^ " *" ^ get_ident tn, Param_ptr tn) :: params

result

module C_syntax_config (Input : sig

end) =

let procs = Input.procs

let use_host_memory = use_host_memory

let logs_to_stdout = false

let main_kernel_prefix = ""

let kernel_prep_line = ""

let typ_of_prec = Ops.c_typ_of_prec

let binop_syntax = Ops.binop_c_syntax

let unop_syntax = Ops.unop_c_syntax

let convert_precision = Ops.c_convert_precision

let module Syntax = C_syntax.C_syntax (C_syntax_config (struct

end)) in

(* FIXME: do we really want all of them, or only the used ones? *)

let idx_params = Indexing.bound_symbols bindings in

let pp_file = Utils.pp_file ~base_name:name ~extension:".c" in

pp_file.finalize ();

let result = c_compile_and_load ~f_name:pp_file.f_name in

{ result; params; bindings; name }

let module Syntax = C_syntax.C_syntax (C_syntax_config (struct

end)) in

(* FIXME: do we really want all of them, or only the used ones? *)

let idx_params = Indexing.bound_symbols bindings in

@@ common_prefix (Array.to_list @@ Array.concat_map ~f:Option.to_array names))

in

let pp_file = Utils.pp_file ~base_name ~extension:".c" in

let params =

Array.mapi lowereds ~f:(fun i lowered ->

Option.map2 names.(i) lowered ~f:(fun name lowered ->

in

pp_file.finalize ();

let result = c_compile_and_load ~f_name:pp_file.f_name in

ptx

module C_syntax_config (Input : sig

end) =

type nonrec buffer_ptr = buffer_ptr [@@deriving sexp_of]

let procs = Input.procs

let use_host_memory = use_host_memory

let logs_to_stdout = true

let main_kernel_prefix = "extern \"C\" __global__"

let kernel_prep_line =

"/* FIXME: single-threaded for now. */if (threadIdx.x != 0 || blockIdx.x != 0) { return; }"

let typ_of_prec = function

| Ops.Byte_prec _ -> "unsigned char"

(* 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

end)) in

let idx_params = Indexing.bound_symbols bindings in

let b = Buffer.create 4096 in

let ppf = Stdlib.Format.formatter_of_buffer b in

if Utils.debug_log_from_routines () then

Stdlib.Format.fprintf ppf "@,__device__ int printf (const char * format, ... );@,";

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

{ traced_store; ptx; params; bindings; name }

let compile_batch ~names bindings lowereds =

let module Syntax = C_syntax.C_syntax (C_syntax_config (struct

end)) in

let idx_params = Indexing.bound_symbols bindings in

let b = Buffer.create 4096 in

let ppf = Stdlib.Format.formatter_of_buffer b in

let params_and_names =

Array.map2_exn names lowereds

~f:

(Option.map2 ~f:(fun name lowered ->

in

let name : string =

String.(

let get_c_ptr ctx num_typ ptr = Gccjit.(RValue.ptr ctx (Type.pointer num_typ) ptr)

let open Gccjit in

let traced = Low_level.(get_node traced_store tn) in

let dims = Lazy.force tn.dims in

let hosted = Tn.is_hosted_force tn 344 in

let in_ctx = Tn.is_in_context ~use_host_memory tn in

let ptr =

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

loop_proc ~toplevel:true ~name ~env body;

!current_block

Low_level.{ traced_store; llc = proc; merge_node } =

let open Gccjit in

let c_index = Type.get ctx Type.Int in

let data =

prepare_node

~debug_log_zero_out:(debug_log_zero_out ctx log_functions get_ident)

in

Hashtbl.add_exn nodes ~key:tn ~data);

let params : (gccjit_param * param_source) list = !param_ptrs in

Block.return_void after_proc;

(ctx_info, params)

let get_ident = Low_level.get_ident_within_code ~no_dots:true [| lowered.llc |] in

let open Gccjit in

if Option.is_none !root_ctx then initialize ();

(* if Utils.settings.with_debug && Utils.settings.output_debug_files_in_build_directory then (

(if Utils.settings.output_debug_files_in_build_directory then

let f_name = Utils.build_file @@ name ^ "-gccjit-debug.c" in

Context.dump_to_file ctx ~update_locs:true f_name);

Context.release ctx;

{ info; result; bindings; name; params = List.map ~f:snd params }

(lowereds : Low_level.optimized option array) =

let get_ident =

Low_level.get_ident_within_code ~no_dots:true

let funcs =

Array.mapi lowereds ~f:(fun i lowered ->

match (names.(i), lowered) with

| Some name, Some lowered ->

Some (info, params)

| _ -> None)

in