Source file `owl_computation_cpu_engine.ml`

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
# 1 "src/base/compute/owl_computation_cpu_engine.ml"
(*
 * OWL - OCaml Scientific Computing
 * Copyright (c) 2016-2022 Liang Wang <liang@ocaml.xyz>
 *)

open Owl_types

(* Functor of making a CPU-based engine to execute computation graphs. *)

module Make_Nested (Graph : Owl_computation_graph_sig.Sig) = struct
  module Graph = Graph
  open Graph.Optimiser.Operator.Symbol

  (* module aliases *)

  module CG_Init = Owl_computation_cpu_init.Make (Graph)
  module CG_Eval = Owl_computation_cpu_eval.Make (Graph)

  (* core interface *)

  (* ! Make sure the order of evaluation of init and eval is always the same. *)
  let eval_gen nodes =
    CG_Init._init_terms nodes;
    CG_Eval._eval_terms nodes


  let eval_elt xs = Array.map elt_to_node xs |> eval_gen

  let eval_arr xs = Array.map arr_to_node xs |> eval_gen

  let eval_graph graph =
    Graph.invalidate_rvs graph;
    Graph.get_outputs graph |> eval_gen
end

(* Functor of making CPU-based engine with unrolled module hierarchy *)

module Make (A : Ndarray_Mutable) = struct
  include
    Owl_computation_engine.Flatten
      (Make_Nested (Owl_computation_engine.Make_Graph (Owl_computation_cpu_device.Make (A))))
end

(* Make functor ends *)