Source file nier_cfg.ml

open Base
open Stdio
open Bigarray

module Tensor = struct
  type ('a, 'b) t = ('a, 'b, c_layout) Genarray.t
  type shape = int array [@@deriving show]

  type ('a, 'b) t_kind =
    | K_int : (int64, int64_elt) t_kind
    | K_float : (float, float64_elt) t_kind

  let create : type a b. shape -> (a, b) t_kind -> (a, b) t =
   fun shape -> function
    | K_float -> Genarray.create float64 c_layout shape
    | K_int -> Genarray.create int64 c_layout shape

  let unsqueeze ~sh1 ~sh2 =
    let sh1, sh2 = (Array.to_list sh1, Array.to_list sh2) in
    let longest, shortest =
      match List.length sh1 > List.length sh2 with
      | true -> (sh1, sh2)
      | false -> (sh2, sh1)
    in
    (*find the index of the potential additional dimension*)
    let where_zero =
      match List.nth_exn longest 0 with
      | 0 -> Some 0
      | _ -> (
        match List.last_exn longest with
        | 0 -> Some (List.length longest - 1)
        | _ -> None)
    in
    match where_zero with
    | Some idx -> (
      match List.sub longest ~pos:idx ~len:(List.length shortest) with
      | [] -> None
      | _ -> Some (Array.of_list longest))
    | None -> None

  let get t idx = Genarray.get t idx
  let set t idx v = Genarray.set t idx v

  let all_coords sh =
    let sh = Array.to_list sh in
    let rec ranges acc shape =
      match shape with
      | x :: y -> ranges (List.init x ~f:(fun i -> i) :: acc) y
      | [] -> acc
      (* a list containing a list of all possible indexes, for each dimension *)
    in
    let xxs = ranges [] sh in
    (* add to each element of the list of all possible coordinates all*)
    (* * possible indexes ... *)
    let aux acc xs =
      List.concat
      @@ List.map xs ~f:(fun x -> List.map ~f:(fun lt -> x :: lt) acc)
      (* ... for each dimension, starting from an empty list of*)
      (* * possible coordinates *)
    in
    List.fold xxs ~init:[ [] ] ~f:aux

  let flatten t =
    let shape = Genarray.dims t in
    let all_idxs = all_coords shape in
    List.init (List.length all_idxs) ~f:(fun i ->
      get t (Array.of_list @@ List.nth_exn all_idxs i))

  let get_shape t = Genarray.dims t

  let equal f t1 t2 =
    let t1_sh = get_shape t1 and t2_sh = get_shape t2 in
    if Array.equal ( = ) t1_sh t2_sh
    then
      let all_idxs = all_coords (get_shape t1) in
      List.fold
        ~f:(fun acc x ->
          if acc
          then f (get t1 (Array.of_list x)) (get t2 (Array.of_list x))
          else false)
        all_idxs ~init:true
    else false

  let num_neurons sh = Array.fold ~init:1 ~f:(fun x y -> x * y) sh

  let get_flatnd_idx ~idx ~sh flt =
    let sh = Array.to_list sh in
    let pop_sh = List.tl_exn sh @ [ 1 ] in
    let rec get_factors_from_sh sh_f l =
      match sh_f with
      | [] -> List.rev l
      | _ ->
        get_factors_from_sh (List.tl_exn sh_f)
          (List.fold ~f:(fun x y -> x * y) ~init:1 sh_f :: l)
    in
    let factors = get_factors_from_sh pop_sh [] in
    let coord_in_data =
      match
        List.fold2
          ~f:(fun x y z -> x + (y * z))
          ~init:0 (Array.to_list idx) factors
      with
      | List.Or_unequal_lengths.Ok i -> i
      | List.Or_unequal_lengths.Unequal_lengths ->
        failwith "Unequal lengths in get_flatnd_idx"
    in
    List.nth_exn flt coord_in_data

  let transpose_2d _t = assert false
end

(* TODO: maybe add markers for special nodes, to reflect they are the inputs and
   outputs of the neural network? *)
module Node = struct
  type shape = int array

  let show_shape sh =
    let sh = Array.to_list sh in
    match sh with
    | [] -> "[]"
    | x :: y ->
      "[" ^ Int.to_string x
      ^ List.fold ~init:"" ~f:(fun str e -> str ^ ";" ^ Int.to_string e) y
      ^ "]"

  type operator =
    | Add
    | Sub
    | Mul
    | Div
    | Matmul
    | Gemm
    | LogSoftmax
    | ReLu
    | Transpose
    | Squeeze
    | MaxPool
    | Conv
    | Reshape
    | Flatten
    | Identity
    | Constant
    | NO_OP
    | RW_Linearized_ReLu

  let str_op o =
    match o with
    | Add -> "Add"
    | Sub -> "Sub"
    | Mul -> "Mul"
    | Div -> "Div"
    | Matmul -> "Matmul"
    | Gemm -> "Gemm"
    | LogSoftmax -> "LogSoftmax"
    | ReLu -> "ReLu"
    | Transpose -> "Transpose"
    | Squeeze -> "Squeeze"
    | MaxPool -> "MaxPool"
    | Conv -> "Conv"
    | Reshape -> "Reshape"
    | Flatten -> "Flatten"
    | Identity -> "Identity"
    | Constant -> "Constant"
    | NO_OP -> "NO_OP"
    | RW_Linearized_ReLu -> "RW_Linearized_ReLu"

  type ksize = Ksize of shape
  type stride = Stride of shape
  type pads = Pads of shape
  type dilations = Dilations of shape

  type operator_parameters =
    | Pool_params of (ksize * stride option * pads option * dilations option)
    | Conv_params of (ksize * stride option * pads option * dilations option)
    | Transpose_params of shape
    | RW_Linearized_ReLu_params of
        (bool list list * ((string, float) Base.Hashtbl.t list * int))

  let str_op_params p =
    match p with
    | Transpose_params s ->
      let str_sh = show_shape s in
      "Transpose params: " ^ str_sh
    | Pool_params (Ksize k, s, p, d) | Conv_params (Ksize k, s, p, d) ->
      let str_k = show_shape k
      and str_s = match s with None -> "" | Some (Stride ss) -> show_shape ss
      and str_p = match p with None -> "" | Some (Pads pp) -> show_shape pp
      and str_d =
        match d with None -> "" | Some (Dilations dd) -> show_shape dd
      in
      "Pool params: KSIZE: " ^ str_k ^ ", Pads: " ^ str_p ^ ", Stride: " ^ str_s
      ^ ", Dilations: " ^ str_d
    | RW_Linearized_ReLu_params l ->
      (* Only displays the activation scheme on the ReLU node *)
      let activations = fst l in
      let act' =
        List.map
          ~f:(fun l1 ->
            List.map
              ~f:(fun b -> match b with true -> "true" | false -> "false")
              l1)
          activations
      in
      let act'' =
        List.map ~f:(fun l -> "[" ^ String.concat ~sep:";" l ^ "]") act'
      in
      let act''' = "[" ^ String.concat ~sep:";" act'' ^ "]" in
      "RW_Linearized_ReLu_params: " ^ act'''

  type ('a, 'b) t = {
    id : int;
    name : string option;
    shape : shape;
    operator : operator;
    operator_parameters : operator_parameters option;
    pred : string list;
    succ : string list;
    tensor : ('a, 'b) Tensor.t option;
  }

  let compare v1 v2 = Stdlib.compare v1.id v2.id
  let hash (v : ('a, 'b) t) = v.id
  let equal v1 v2 = v1.id = v2.id

  let create ~id ~name ~sh ~op ~op_p ~pred ~succ ~tensor =
    {
      id;
      name;
      shape = sh;
      operator = op;
      operator_parameters = op_p;
      pred;
      succ;
      tensor;
    }

  let get_name t = match t.name with Some n -> n | None -> "C_NODE"
  let get_shape t = t.shape
  let get_op t = t.operator
  let get_tensor t = t.tensor
  let get_pred_list t = t.pred
  let get_succ_list t = t.succ
  let is_data_node t = match get_tensor t with None -> false | Some _ -> true

  (* TODO: some flags on the node would be cleaner than this*)
  let is_input_node t = List.equal String.equal t.pred [ "NO_INPUT" ]
  let is_output_node t = List.equal String.equal t.succ [ "NO_OUTPUT" ]

  let num_neurons t =
    match get_shape t with
    | [||] -> 0
    | l -> Array.fold ~init:1 ~f:(fun x acc -> x * acc) l

  let show n f =
    let id = Int.to_string n.id in
    let name = get_name n
    and operator = str_op n.operator
    and operator_parameters =
      match n.operator_parameters with
      | Some p -> str_op_params p
      | None -> "no parameters"
    and shape = show_shape n.shape
    and prevs =
      List.fold_left ~f:(fun x y -> x ^ "," ^ y) ~init:"" (get_pred_list n)
    and nexts =
      List.fold_left ~f:(fun x y -> x ^ "," ^ y) ~init:"" (get_succ_list n)
    and tensor =
      match n.tensor with
      (*limit of size for tensor strings, complying with
       * dot string size limit of 16Ko *)
      | Some t ->
        let display_indices =
          let all_indices = Tensor.all_coords (Tensor.get_shape t) in
          if List.length all_indices > 10
          then
            let rec firstk k xs =
              match xs with
              | [] -> failwith "firstk"
              | x :: xs -> if k = 1 then [ x ] else x :: firstk (k - 1) xs
            in
            firstk 10 all_indices
          else all_indices
        in
        let t_value_string f =
          List.fold_left
            ~f:(fun acc l ->
              acc
              ^ show_shape (Array.of_list l)
              ^ ": "
              ^ f (Tensor.get t (Array.of_list l))
              ^ "\n")
            ~init:"" display_indices
        in
        "Tensor value\n: " ^ t_value_string f ^ "\nShape: "
        ^ show_shape (Tensor.get_shape t)
      | None -> "No tensor in node"
    in
    "ID :" ^ id ^ "\nNAME: " ^ name ^ "\nOP: " ^ operator ^ "\nOP PARAMS:"
    ^ operator_parameters ^ "\nSHAPE: " ^ shape ^ "\nPREVS: " ^ prevs
    ^ "\nNEXTS: " ^ nexts ^ "\nTENSORS INFOS:" ^ tensor
end

module type VInput = sig
  type l
  type r

  val convert_f : l -> string
end

module MakeVertex (I : VInput) = struct
  type t = (I.l, I.r) Node.t

  let compare = Node.compare
  let hash = Node.hash
  let equal = Node.equal
  let convert_f = I.convert_f

  type label = string

  let label (n : t) = match n.Node.name with Some n -> n | None -> ""
  let create _name = assert false
end

module Edge = struct
  type t = string

  let compare = Stdlib.compare
  let equal = phys_equal
  let default = ""
end

module NierCFG (I : VInput) = struct
  module Vertex = MakeVertex (I)
  include Graph.Imperative.Digraph.ConcreteBidirectionalLabeled (Vertex) (Edge)

  let convert_f = Vertex.convert_f
  let vertex_list g = fold_vertex (fun x l -> x :: l) g []

  let input_nodes g =
    let input_criterion (v : ('a, 'b) Node.t) acc =
      match v.id with 0 -> Some v | _ -> acc
    in
    match fold_vertex (fun v acc -> input_criterion v acc) g None with
    | Some r -> [ r ]
    | None -> failwith "Something strange, no node for describing inputs found"

  let preds g v = pred g v

  let preds_names g v =
    let preds_list = pred_e g v in
    List.fold ~init:[] ~f:(fun acc (_, n, _) -> n :: acc) preds_list

  let succs_names g v =
    let succs_list = succ_e g v in
    List.fold ~init:[] ~f:(fun acc (_, n, _) -> n :: acc) succs_list

  let succs g v = succ g v
  let init_cfg = create ()

  let find_vertices g f =
    fold_vertex (fun x l -> if f x then x :: l else l) g []

  let data_node_of n g =
    fold_pred (fun v _ -> if Node.is_data_node v then Some v else None) g n None

  let infer_shape g n in_shape ~on_backward =
    let op = Node.get_op n in
    match op with
    | Node.Add -> (
      match data_node_of n g with
      | Some d_n -> Node.get_shape d_n
      | None -> failwith "Error, Add operator lacks a data node")
    | Node.ReLu -> in_shape
    | Node.Matmul ->
      let pad_left = function
        | [] -> failwith "Impossible to pad empty shape"
        | [ a ] -> [ 1; a ]
        | x -> x
      in
      let pad_right = function
        | [] -> failwith "Impossible to pad empty shape"
        | [ a ] -> [ a; 1 ]
        | x -> x
      in
      let rec one_padding l i =
        if i <= 0 then l else one_padding (1 :: l) (i - 1)
      in
      let dn_shape =
        match data_node_of n g with
        | Some dn -> Node.get_shape dn
        | None -> failwith "Error, MatMul operator lacks a data node"
      in
      (* Expected semantic:
       * Matrix multiplication C = AB
       * A (shape [n;m]); B (shape [m;p]); C (shape [n;p])
       * shape of b: b_sh
       * shape of a: a_sh
       * shape of c: c_sh
       * It is expected here that B is the shape of the node
       * yielding the data tensor in the NIER
       *)
      let check_matmul_size_ba ~b_sh ~a_sh =
        let bdim2 = pad_left b_sh in
        let adim2 = pad_right a_sh in
        let bdim = one_padding bdim2 (List.length adim2 - List.length bdim2) in
        let adim = one_padding adim2 (List.length bdim2 - List.length adim2) in
        let rec infer_csize acc ad bd =
          match (ad, bd) with
          | [ m; n ], [ nn; p ] ->
            if nn = n
            then (n, List.append (List.rev acc) [ m; p ])
            else failwith "size of matrices not adequate"
          | a :: la, b :: lb ->
            if a = b
            then infer_csize (a :: acc) la lb
            else if a = 1
            then infer_csize (b :: acc) la lb
            else if b = 1
            then infer_csize (a :: acc) la lb
            else failwith "Checking matmul_size failed: one discordance"
          | _, _ -> failwith "Checking matmul_size failed"
        in
        infer_csize [] bdim adim
      in
      let check_matmul_size_bc ~b_sh ~c_sh =
        let bdim2 = pad_left b_sh in
        let cdim2 = pad_right c_sh in
        let bdim = one_padding bdim2 (List.length cdim2 - List.length bdim2) in
        let cdim = one_padding cdim2 (List.length bdim2 - List.length cdim2) in
        let rec infer_asize acc bd cd =
          match (bd, cd) with
          | [ m; p ], [ n; pp ] ->
            if pp = p
            then (n, List.append (List.rev acc) [ n; m ])
            else failwith "size of matrices not adequate"
          | b :: lb, c :: lc ->
            if b = c
            then infer_asize (b :: acc) lb lc
            else if b = 1
            then infer_asize (b :: acc) lb lc
            else if c = 1
            then infer_asize (c :: acc) lb lc
            else failwith "Checking matmul_size failed: one discordance"
          | _, _ -> failwith "Checking matmul_size failed"
        in
        infer_asize [] bdim cdim
      in
      if on_backward
      then
        Array.of_list
        @@ snd
             (check_matmul_size_bc ~b_sh:(Array.to_list dn_shape)
                ~c_sh:(Array.to_list in_shape))
      else
        Array.of_list
        @@ snd
             (check_matmul_size_ba ~b_sh:(Array.to_list in_shape)
                ~a_sh:(Array.to_list dn_shape))
    | a -> failwith (Printf.sprintf "operator %s not supported" (Node.str_op a))
end

module NierCFGInt = NierCFG (struct
  type l = int64
  type r = int64_elt

  let convert_f = Int64.to_string
end)

module NierCFGFloat = NierCFG (struct
  type l = float
  type r = float64_elt

  let convert_f = Float.to_string
end)

module NierCFGDot = Graph.Graphviz.Dot (struct
  include NierCFGFloat (* use the graph module from above *)

  let node_label (v : vertex) = Node.show v convert_f
  let edge_attributes (_, e, _) = [ `Label e; `Color 4711 ]
  let default_edge_attributes _ = []
  let get_subgraph _ = None
  let vertex_attributes v = [ `Shape `Box; `Label (node_label v) ]
  let vertex_name (v : vertex) = Int.to_string v.id
  let default_vertex_attributes _ = []
  let graph_attributes _ = []
end)

let print_cfg_graph g = NierCFGDot.fprint_graph Stdlib.Format.std_formatter g

let out_cfg_graph g =
  let file = Out_channel.create "cfg.dot" in
  NierCFGDot.output_graph file g