Source file onnx.ml

(**************************************************************************)
(*                                                                        *)
(*  This file is part of CAISAR.                                          *)
(*                                                                        *)
(*  Copyright (C) 2023                                                    *)
(*    CEA (Commissariat à l'énergie atomique et aux énergies              *)
(*         alternatives)                                                  *)
(*                                                                        *)
(*  You can redistribute it and/or modify it under the terms of the GNU   *)
(*  Lesser General Public License as published by the Free Software       *)
(*  Foundation, version 2.1.                                              *)
(*                                                                        *)
(*  It is distributed in the hope that it will be useful,                 *)
(*  but WITHOUT ANY WARRANTY; without even the implied warranty of        *)
(*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the          *)
(*  GNU Lesser General Public License for more details.                   *)
(*                                                                        *)
(*  See the GNU Lesser General Public License version 2.1                 *)
(*  for more details (enclosed in the file licenses/LGPLv2.1).            *)
(*                                                                        *)
(**************************************************************************)

open Base
module Format = Stdlib.Format
module Fun = Stdlib.Fun
module Oproto = Onnx_protoc (* Autogenerated during compilation *)
module Oprotom = Oproto.Onnx.ModelProto
module NCFG = Ir.Nier_cfg
module G = NCFG.NierCFGFloat

exception ParseError of string

type t = {
  n_inputs : int; (* Number of inputs. *)
  n_outputs : int; (* Number of outputs. *)
  nier : (G.t, string) Result.t; (* Intermediate representation. *)
}

(* ONNX format handling. *)
type op_attribute = Oproto.Onnx.AttributeProto.t

type tensordata =
  | Raw of bytes
  | Float of float list

let (no_attr : op_attribute) =
  {
    name = None;
    ref_attr_name = None;
    doc_string = None;
    type' = None;
    f = None;
    i = None;
    s = None;
    t = None;
    g = None;
    floats = [];
    ints = [];
    strings = [];
    tensors = [];
    graphs = [];
    sparse_tensor = None;
    tp = None;
    sparse_tensors = [];
    type_protos = [];
  }

let get_nested_dims (s : Oproto.Onnx.ValueInfoProto.t list) =
  match List.nth s 0 with
  | Some { type' = Some { value = `Tensor_type { shape = Some v; _ }; _ }; _ }
    ->
    v
  | _ -> []

let flattened_dim (dim : Oproto.Onnx.TensorShapeProto.Dimension.t list) =
  List.fold ~init:1 dim ~f:(fun acc x ->
    match x.value with
    | `Dim_value v -> acc * v
    | `Dim_param _ -> acc
    | `not_set -> acc)

let get_input_output_dim (model : Oprotom.t) =
  let input_shape, output_shape =
    match model.graph with
    | Some g -> (get_nested_dims g.input, get_nested_dims g.output)
    | _ -> ([], [])
  in
  (* TODO: here we only get the flattened dimension of inputs and outputs, but
     more interesting parsing could be done later on. *)
  let input_flat_dim = flattened_dim input_shape in
  let output_flat_dim = flattened_dim output_shape in
  (input_flat_dim, output_flat_dim)

let produce_cfg (g : Oproto.Onnx.GraphProto.t) =
  let open Oproto.Onnx in
  let nodes = g.node
  and inputs = g.input
  and outputs = g.output
  and initi = g.initializer' in
  let fold_vip_names acc n =
    match n.ValueInfoProto.name with
    | Some str -> Some str :: acc
    | None -> None :: acc
  in
  let i_nodes, o_nodes =
    ( List.fold inputs ~init:[] ~f:fold_vip_names,
      List.fold outputs ~init:[] ~f:fold_vip_names )
  and c_nodes = List.init (List.length nodes) ~f:(fun _ -> None) in
  let fold_nodes_ops_cfg ns =
    let get_node_operator_cfg x =
      match x.NodeProto.op_type with
      | None -> NCFG.Node.NO_OP
      | Some o -> (
        match o with
        | "Add" -> NCFG.Node.Add
        | "Sub" -> NCFG.Node.Sub
        | "Mul" -> NCFG.Node.Mul
        | "Div" -> NCFG.Node.Div
        | "Relu" -> NCFG.Node.ReLu
        | "MatMul" -> NCFG.Node.Matmul
        | "Gemm" -> NCFG.Node.Gemm
        | "LogSoftmax" -> NCFG.Node.LogSoftmax
        | "Transpose" -> NCFG.Node.Transpose
        | "Squeeze" -> NCFG.Node.Squeeze
        | "MaxPool" -> NCFG.Node.MaxPool
        | "Constant" -> NCFG.Node.Constant
        | "Conv" -> NCFG.Node.Conv
        | "Reshape" -> NCFG.Node.Reshape
        | "Flatten" -> NCFG.Node.Flatten
        | "Identity" -> NCFG.Node.Identity
        | _ -> raise (ParseError ("Unsupported ONNX operator " ^ o)))
    in
    List.fold ~f:(fun acc n -> get_node_operator_cfg n :: acc) ~init:[] ns
  in
  let c_ops = List.rev @@ fold_nodes_ops_cfg nodes
  and i_ops, o_ops =
    ( List.init ~f:(fun _ -> NCFG.Node.NO_OP) (List.length i_nodes),
      List.init ~f:(fun _ -> NCFG.Node.NO_OP) (List.length o_nodes) )
  in
  let fold_nodes_attr ns =
    let get_node_attr n = n.NodeProto.attribute in
    List.fold ~f:(fun acc n -> get_node_attr n :: acc) ~init:[] ns
  in

  let c_attr = List.rev @@ fold_nodes_attr nodes
  and i_attr, o_attr =
    ( List.init ~f:(fun _ -> [ no_attr ]) (List.length i_nodes),
      List.init ~f:(fun _ -> [ no_attr ]) (List.length o_nodes) )
  in
  let c_nodes_inputs, c_nodes_outputs =
    List.unzip
    @@ List.fold
         ~f:(fun acc n -> (n.NodeProto.input, n.NodeProto.output) :: acc)
         ~init:[] (List.rev nodes)
  and i_nodes_inputs, i_nodes_outputs, o_nodes_inputs, o_nodes_outputs =
    ( List.init ~f:(fun _ -> [ "NO_INPUT" ]) (List.length i_nodes),
      List.init ~f:(fun _ -> [ "" ]) (List.length i_nodes),
      List.init ~f:(fun _ -> [ "" ]) (List.length o_nodes),
      List.init ~f:(fun _ -> [ "NO_OUTPUT" ]) (List.length o_nodes) )
  in
  let data_dict =
    let dict_tensors_cfg ts =
      let get_float_from_index index data sh =
        let index = Array.to_list index and sh = Array.to_list sh in
        let pop_sh = List.tl_exn sh @ [ 1 ] in
        (* Returns the factors by which multiply each coordinate *)
        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 =
          List.fold2_exn ~f:(fun x y z -> x + (y * z)) ~init:0 index factors
        in
        match data with
        | Raw raw ->
          let offset = 4 * coord_in_data in
          (* Each float is coded on 4 bytes*)
          let res = EndianBytes.LittleEndian.get_float raw offset in
          res
        | Float f -> List.nth_exn f coord_in_data
      in
      let build_tensor_from_data sh data =
        let open NCFG.Tensor in
        let sh = Array.of_list @@ sh in
        let tensor = create sh K_float in
        let coords = all_coords (get_shape tensor) in
        let rec init_tensor t idx r =
          match idx with
          | x :: y ->
            let value =
              get_float_from_index (Array.of_list x) r (get_shape t)
            in
            set t (Array.of_list x) value;
            init_tensor t y r
          | [] -> t
        in
        init_tensor tensor coords data
      in
      let t_name x =
        match x.TensorProto.name with Some n -> n | None -> "C_NODE"
      in
      let t_dim x = x.TensorProto.dims in
      let t_data x =
        match x.TensorProto.raw_data with
        | Some rd -> Some (build_tensor_from_data (t_dim x) (Raw rd))
        | None -> (
          match x.TensorProto.float_data with
          | [] -> None
          | f -> Some (build_tensor_from_data (t_dim x) (Float f)))
      in
      List.fold
        ~f:(fun m x -> Map.add_exn m ~key:(t_name x) ~data:(t_data x))
        ~init:(Map.empty (module String))
        ts
    in
    dict_tensors_cfg initi
  in
  let unpack v =
    match v with
    | Some v -> v
    | None -> failwith "Unpack found an unexpected None"
  in
  let tensor_list =
    List.init
      ~f:(fun i ->
        match Map.find data_dict (unpack (List.nth_exn i_nodes i)) with
        | Some v -> v
        | None -> None)
      (List.length i_nodes)
  in
  let tensor_list_full = Map.to_alist data_dict in
  let tensor_list_rev = List.rev tensor_list in
  let vip_dims v =
    let val_t =
      match v.ValueInfoProto.type' with
      | Some t -> t
      | None -> failwith "No type in value info"
    in
    let tns_t =
      match val_t.TypeProto.value with
      | `Tensor_type t -> t
      | `not_set ->
        failwith "No tensor type in value info"
        (* TODO: support more tensor types *)
      | _ -> raise (ParseError "Unknown tensor type")
    in
    let tns_s =
      match tns_t.shape with
      | Some s -> s
      | None -> failwith "No tensor shape in value info"
    in
    List.rev
    @@ List.fold tns_s ~init:[] ~f:(fun acc d ->
         match d.value with
         | `Dim_value d -> d :: acc
         | `not_set | _ -> 0 :: acc)
  in

  let c_tensordim_list = List.init (List.length nodes) ~f:(fun _ -> [])
  and c_tensorraw_list = List.init (List.length nodes) ~f:(fun _ -> None)
  and o_tensordim_list =
    List.fold ~f:(fun acc n -> vip_dims n :: acc) ~init:[] outputs
  and o_tensorraw_list = List.init (List.length o_nodes) ~f:(fun _ -> None)
  and i_tensordim_list =
    List.fold ~f:(fun acc n -> vip_dims n :: acc) ~init:[] inputs
  and i_tensorraw_list = tensor_list_rev in
  let nodes_names = i_nodes @ c_nodes @ o_nodes in
  let ops = i_ops @ c_ops @ o_ops in
  let attrs = i_attr @ c_attr @ o_attr in
  let prevs_list = i_nodes_inputs @ c_nodes_inputs @ o_nodes_inputs in
  let nexts_list = i_nodes_outputs @ c_nodes_outputs @ o_nodes_outputs in
  let tensor_dims = i_tensordim_list @ c_tensordim_list @ o_tensordim_list in
  let tensors = i_tensorraw_list @ c_tensorraw_list @ o_tensorraw_list in
  let operator_parameters (attr : AttributeProto.t list) op =
    match op with
    | NCFG.Node.Transpose ->
      let ints_params = Array.of_list (List.nth_exn attr 0).ints in
      Some (NCFG.Node.Transpose_params ints_params)
    (*TODO: maxpool and conv operators: match attr.name in attributes to
     * create the correct value for each attribute*)
    (* | NCFG.Vertex.MaxPool -> *)
    (* | NCFG.Vertex.Conv -> *)
    | _ -> None
  in
  let rec build_op_param_list attrs ops l =
    match (attrs, ops) with
    | a :: b, c :: d -> build_op_param_list b d (operator_parameters a c :: l)
    | [], [] ->
      List.rev l
      (*All other list constructions are folding right, so we need to put a
        final revert *)
    | _ ->
      raise (ParseError "Operator and attribute lists have not the same size")
  in
  let op_params_cfg = build_op_param_list attrs ops [] in
  let cfg = G.init_cfg in
  (* adding inputs, outputs and cnodes to the cfg *)
  let unkerasize l = List.map ~f:(fun x -> if x = 0 then 1 else x) l in
  for i = 0 to List.length nodes_names - 1 do
    let (v : G.V.t) =
      NCFG.Node.create ~id:i
        ~name:(List.nth_exn nodes_names i)
        ~sh:(Array.of_list @@ unkerasize (List.nth_exn tensor_dims i))
        ~op:(List.nth_exn ops i)
        ~op_p:(List.nth_exn op_params_cfg i)
        ~pred:(List.nth_exn prevs_list i)
        ~succ:(List.nth_exn nexts_list i)
        ~tensor:(List.nth_exn tensors i)
    in
    G.add_vertex cfg v
  done;
  (* Adding edges between vertices *)
  (* For each unnamed vertex (= a calculation node) in the cfg ... *)
  (* return true if l1 has at least one common element wih l2 *)
  let rec shared_elm l1 l2 =
    match l1 with
    | x :: y -> List.mem l2 x ~equal:String.equal || shared_elm y l2
    | [] -> false
  in
  List.iter
    ~f:(fun (v : G.V.t) ->
      match v.name with
      | None ->
        let pred = v.pred and succ = v.succ in
        let prev_v =
          (* ... get all vertices in cfg that have the current vertex preds
           * in their succ (at least one of their succ is inside our preds )*)
          G.find_vertices cfg (fun (x : G.V.t) ->
            if shared_elm pred x.succ then true else false)
        (* ... get all vertices in cfg that have the current vertex preds
         * in their name (they are named the same as one of our preds )*)
        and named_pred =
          G.find_vertices cfg (fun (x : G.V.t) ->
            match x.name with
            | Some name -> if shared_elm pred [ name ] then true else false
            | None -> false)
        (* ... get all vertices in cfg that have the current vertex succ
         * in their name (they are named the same as one of our succs )*)
        and named_succ =
          G.find_vertices cfg (fun (x : G.V.t) ->
            match x.name with
            | Some name -> if shared_elm succ [ name ] then true else false
            | None -> false)
        (* get all vertices in cfg that have the current vertex succs
         * in their preds (at least one of their preds is inside our succ )*)
        and next_v =
          G.find_vertices cfg (fun (x : G.V.t) ->
            if shared_elm succ x.pred then true else false)
        in
        (* add edges between current vertex and identified preds and succs *)
        let v_predecessors = prev_v @ named_pred
        and v_successors = next_v @ named_succ in
        let unpack_tname (x : G.V.t) =
          match x.NCFG.Node.name with Some n -> n | None -> ""
        in
        List.iter
          ~f:(fun (x : G.V.t) ->
            let label =
              match List.nth x.succ 0 with
              | Some "NO_OUTPUT" ->
                let pred_name = unpack_tname x in
                if List.mem ~equal:String.equal v.NCFG.Node.pred pred_name
                then pred_name
                else ""
              | Some l -> l
              | None -> ""
            in
            G.add_edge_e cfg (x, label, v))
          v_predecessors;
        (* add successors edges after filtering those
         * that are already an edge*)
        List.iter
          ~f:(fun (x : G.V.t) ->
            let all_preds = G.preds cfg x and all_succs = G.succs cfg x in
            if List.mem ~equal:NCFG.Node.equal all_preds v
               || List.mem ~equal:NCFG.Node.equal all_succs v
            then ()
            else
              let label =
                match List.nth_exn x.pred 0 with
                | "NO_INPUT" ->
                  let succ_name = unpack_tname x in
                  if List.mem ~equal:String.equal v.NCFG.Node.succ succ_name
                  then succ_name
                  else ""
                | l -> l
              in
              G.add_edge_e cfg (v, label, x))
          v_successors
      | _ -> ())
    (G.vertex_list cfg);

  (*rationale of the following:
   * PyTorch stores network nodes in the field "inputs" of
   * the ONNX graph g, and network parameters as a list of tensors
   * in the ONNX initializer_.
   * To make the two correspond, elements of g.inputs and g.initializer_
   * share the same value in the field "name".
   * In Keras, elements of g.initializer_ have a name, but they do not
   * correspond to any name in g.inputs.
   * What we did before was then to create the actual nier cfg following the
   * PyTorch way.
   * Below, we complete the cfg with keras data by doing the following:
   *  * create a node for NIER for each tensor in onnx initializer_
   *  * for each NIER node, check if there is a node sharing the same name
   *    pred
   *  * if yes, remove the one with highest ID (those are initi nodes, but since
   *  there is already a node in CFG with this name we do not
   *  need those)
   *  * if not, for each NIER node, chck if there is a node
   *  which name is contained in prevs. add it to the prev
   * *)

  (* adding initi vertices to the cfg *)
  for i = 0 to List.length tensor_list_full - 1 do
    let shape =
      match snd (List.nth_exn tensor_list_full i) with
      | Some t -> unkerasize (Array.to_list @@ NCFG.Tensor.get_shape t)
      | None -> []
    in
    let (v : G.V.t) =
      NCFG.Node.create
        ~id:(i + List.length nodes_names)
        ~name:(Some (fst (List.nth_exn tensor_list_full i)))
        ~sh:(Array.of_list @@ unkerasize shape)
        ~op:NO_OP ~op_p:None ~pred:[] ~succ:[]
        ~tensor:(snd (List.nth_exn tensor_list_full i))
    in
    G.add_vertex cfg v
  done;
  (* build a list of nodes
   * sharing name but with different ids *)
  let same_name_diff_ids =
    let aux (x : G.V.t) =
      G.fold_vertex
        (fun y acc ->
          match (x.name, y.name) with
          | Some xa, Some ya ->
            if (not (y.id = x.id)) && String.equal xa ya
            then (x, y) :: acc
            else acc
          | _ -> acc)
        cfg []
    in
    G.fold_vertex (fun x l -> aux x :: l) cfg []
  in
  let highest_ids =
    List.fold
      ~f:(fun acc x ->
        match x with
        | a :: _ ->
          let maxval = max (fst a).NCFG.Node.id (snd a).NCFG.Node.id in
          maxval :: acc
        | [] -> acc)
      ~init:[] same_name_diff_ids
  in
  (* (* removing nodes with highest id, those are the*) (* * ones we just added
     *)*)
  List.iter
    ~f:(fun x ->
      match x with
      | l :: _ ->
        let v1 = fst l in
        if List.mem ~equal:( = ) highest_ids v1.NCFG.Node.id
        then
          (* Printf.printf "Removing id %d \n%!" *)
          (*   v1.NCFG.Vertex.id; *)
          G.remove_vertex cfg v1
        else ()
      | [] -> ())
    same_name_diff_ids;
  (* Now it is Keras time.
   * Look for nodes sharing name and preds,
   * then create edge *)
  let shared_name_preds =
    let aux (x : G.V.t) =
      match x.name with
      (* look in other vertices if name is among
       * predecessors *)
      | Some n -> G.find_vertices cfg (fun x -> shared_elm [ n ] x.pred)
      | None -> []
    in
    G.fold_vertex (fun x l -> (x, aux x) :: l) cfg []
  in
  List.iter
    ~f:(fun x ->
      let orgn = fst x and to_edge = snd x in
      List.iter
        ~f:(fun t ->
          if not (G.mem_edge cfg orgn t)
          then G.add_edge_e cfg (orgn, unpack orgn.NCFG.Node.name, t)
          else ())
        to_edge)
    shared_name_preds;
  (* else (); *)
  cfg

let nier_of_onnx_protoc (model : Oprotom.t) =
  match model.graph with
  | Some g -> produce_cfg g
  | None -> raise (ParseError "No graph in ONNX input file found")

let parse_in_channel in_channel =
  let open Result in
  try
    let buf = Stdio.In_channel.input_all in_channel in
    let reader = Ocaml_protoc_plugin.Reader.create buf in
    match Oprotom.from_proto reader with
    | Ok r ->
      let n_inputs, n_outputs = get_input_output_dim r in
      let nier =
        try Ok (nier_of_onnx_protoc r) with
        | ParseError s | Sys_error s -> Error s
        | Failure msg -> Error (Format.sprintf "Unexpected error: %s" msg)
      in
      Ok { n_inputs; n_outputs; nier }
    | _ -> Error "Cannot read protobuf"
  with
  | Sys_error s -> Error s
  | Failure msg -> Error (Format.sprintf "Unexpected error: %s" msg)

let parse filename =
  let in_channel = Stdlib.open_in filename in
  Fun.protect
    ~finally:(fun () -> Stdlib.close_in in_channel)
    (fun () -> parse_in_channel in_channel)