Source file node.ml

(**************************************************************************)
(*                                                                        *)
(*  This file is part of CAISAR.                                          *)
(*                                                                        *)
(*  Copyright (C) 2024                                                    *)
(*    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

type ty =
  | Float
  | Int64
[@@deriving show]

(** TODO: add the information needed to compute the shape *)
type descr =
  | Constant of { data : Gentensor.t }
  | Add of {
      input1 : t;
      input2 : t;
    }
  | Sub of {
      input1 : t;
      input2 : t;
    }
  | Mul of {
      input1 : t;
      input2 : t;
    }
  | Div of {
      input1 : t;
      input2 : t;
    }
  | Matmul of {
      input1 : t;
      input2 : t;
    }
  | Gemm of {
      inputA : t;
      inputB : t;
      inputC : t option;
      alpha : float;
      beta : float;
      transA : int;
      transB : int;
    }
  | LogSoftmax
  | ReLu of { input : t }
  | Transpose of {
      input : t;
        (* called "data" in ONNX documentation :
           https://onnx.ai/onnx/operators/onnx__Transpose.html*)
      perm : int list;
    }
  | Squeeze of {
      data : t;
      axes : t option; (* data int64 *)
    }
  | MaxPool
  | Conv
  | Reshape of {
      input : t;
      shape : t; (* data int64 *)
    }
  | Flatten of {
      input : t;
      axis : int;
    }
  | Identity of { input : t }
  | Input of { shape : Shape.t }
  | RW_Linearized_ReLu
  | Concat of {
      inputs : t list;
      axis : int;
    }
  | Gather of {
      input : t;
      indices : t;
      axis : int;
    }
  | ReduceSum of {
      input : t;
      axes : t option;
      keepdims : int;
      noop_with_empty_axes : int;
    }
  | GatherND of {
      data : t;
      indices : t;
      batch_dims : int;
    }
  | RandomNormal of {
      dtype : int;
      mean : float;
      scale : float;
      seed : float;
      shape : int array;
    }
  | Abs of { input : t }
  | Log of { input : t }

and t = {
  id : int;
  descr : descr; [@printer fun fmt d -> pp_descr fmt d]
  shape : Shape.t;
  ty : ty;
}

let pp_descr fmt p =
  match p with
  | Input { shape } -> Fmt.pf fmt "Input: %a" Shape.pp shape
  | Transpose { perm; _ } ->
    Fmt.pf fmt "Transpose: [%a]" Fmt.(list ~sep:semi int) perm
  | Constant { data = Int64 b } when Shape.size (Tensor.shape b) < 3 ->
    Fmt.pf fmt "Constant[%a]" Fmt.(list ~sep:comma int64) (Tensor.flatten b)
  | Constant _ -> Fmt.pf fmt "Constant"
  | Add _ -> Fmt.pf fmt "Add"
  | Sub _ -> Fmt.pf fmt "Sub"
  | Mul _ -> Fmt.pf fmt "Mul"
  | Div _ -> Fmt.pf fmt "Div"
  | Matmul _ -> Fmt.pf fmt "Matmul"
  | Gemm _ -> Fmt.pf fmt "Gemm"
  | LogSoftmax -> Fmt.pf fmt "LogSoftmax"
  | ReLu _ -> Fmt.pf fmt "ReLu"
  | Squeeze _ -> Fmt.pf fmt "Squeeze"
  | MaxPool -> Fmt.pf fmt "MaxPool"
  | Conv -> Fmt.pf fmt "Conv"
  | Reshape _ -> Fmt.pf fmt "Reshape"
  | Flatten _ -> Fmt.pf fmt "Flatten"
  | Identity _ -> Fmt.pf fmt "Identity"
  | RW_Linearized_ReLu -> Fmt.pf fmt "RW_Linearized_ReLu"
  | Concat { axis; _ } -> Fmt.pf fmt "Concat[%i]" axis
  | Gather _ -> Fmt.pf fmt "Gather"
  | ReduceSum _ -> Fmt.pf fmt "ReduceSum"
  | GatherND _ -> Fmt.pf fmt "GatherND"
  | RandomNormal _ -> Fmt.pf fmt "RandomNormal"
  | Abs _ -> Fmt.pf fmt "Abs"
  | Log _ -> Fmt.pf fmt "Log"

let show_descr t = Fmt.str "%a" pp_descr t
let compare { id = id1; _ } { id = id2; _ } = Int.compare id1 id2
let equal { id = id1; _ } { id = id2; _ } = Int.equal id1 id2
let hash { id; _ } = id
let sexp_of_t node = Base.Int.sexp_of_t node.id
let pp fmt n = Fmt.pf fmt "@[%i: %a@]" n.id pp_descr n.descr
let show n = Fmt.str "%a" pp n

include Base.Comparator.Make (struct
  type nonrec t = t

  let compare = compare
  let sexp_of_t = sexp_of_t
end)

let rec compute_shape n = n.shape

and compute_shape_descr = function
  | Add { input1; _ }
  | Div { input1; _ }
  | Mul { input1; _ }
  | Sub { input1; _ } ->
    compute_shape input1
  | Flatten { input; axis } ->
    (* (d_0 X d_1 … d_(axis-1), d_axis X d_(axis+1) … X dn). *)
    let shape = compute_shape input in
    let d1 = ref 1 in
    let d2 = ref 1 in
    for i = 0 to axis - 1 do
      d1 := !d1 * Shape.get shape i
    done;
    for i = axis to Shape.rank shape - 1 do
      d2 := !d2 * Shape.get shape i
    done;
    Shape.of_list [ !d1; !d2 ]
  | Input { shape } -> shape
  | ReLu { input } -> compute_shape input
  | Transpose { input; perm = [] } ->
    compute_shape input |> Shape.to_list |> List.rev |> Shape.of_list
  | Transpose { input; perm } ->
    let shape = compute_shape input in
    let rank = Shape.rank shape in
    assert (Int.equal rank (List.length perm));
    let shape' = Array.create ~len:rank 0 in
    let shape = Shape.to_array shape in
    Base.List.iteri perm ~f:(fun i j -> Array.set shape' i (Array.get shape j));
    Shape.of_array shape'
  | Constant { data } -> Gentensor.shape data
  | Concat { inputs; axis } ->
    let shapes = List.map ~f:compute_shape inputs in
    let shape = List.hd_exn shapes in
    let axis = if axis < 0 then Shape.rank shape + axis else axis in
    let l = List.map ~f:(fun s -> Shape.get s axis) shapes in
    let i = List.reduce_exn ~f:( + ) l in
    Shape.set shape axis i
  | Gather { input; indices; axis } -> (
    let input_shape = compute_shape input in
    let indices_shape = compute_shape indices in
    let axis = if axis < 0 then Shape.rank input_shape + axis else axis in
    match List.split_n (Shape.to_list input_shape) axis with
    | _, [] -> failwith "axis is bigger than shape rank"
    | before, _ :: after ->
      Shape.of_list (before @ Shape.to_list indices_shape @ after))
  | Matmul { input1; input2 } ->
    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
    (* 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
     *)
    let check_matmul_size_ab ~a_sh ~b_sh =
      let adim2 = pad_left a_sh in
      let bdim2 = pad_right b_sh in
      let adim = one_padding adim2 (List.length bdim2 - List.length adim2) in
      let bdim = one_padding bdim2 (List.length adim2 - List.length bdim2) in
      let rec infer_csize acc ad bd =
        match (ad, bd) with
        | [ m; n ], [ nn; p ] ->
          if nn = n
          then List.rev_append 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 [] adim bdim
    in
    (* TODO: in case of pad_left and pad_right remove the added dimension. But
       it is not clear what must be done when broadcasting is done *)
    Shape.of_list
      (check_matmul_size_ab
         ~a_sh:(Shape.to_list (compute_shape input1))
         ~b_sh:(Shape.to_list (compute_shape input2)))
  | Reshape { input; shape; _ } ->
    let shape =
      match shape.descr with
      | Constant { data = Int64 a } ->
        List.map ~f:Int64.to_int_exn (Tensor.flatten a)
      | _ ->
        (* Some constant propagation could be useful in some cases eg. patch-1
           VNNcomp *)
        failwith "non-constant shape in reshape not supported"
    in
    List.iter shape ~f:(function
      | -1 | 0 -> failwith "not implemented 0 -1 in shape for reshape"
      | _ -> ());
    let out = Shape.of_list shape in
    if Shape.size out <> Shape.size input.shape
    then
      failwith
        "Reshape: shape of input and shape given have not the same number of \
         elements";
    out
  | Gemm { inputA; inputB; inputC = _; alpha = _; beta = _; transA; transB } ->
    let rank2 i =
      match Shape.to_array_unsafe i.shape with
      | [| k; n |] -> (k, n)
      | _ -> failwith "Gemm input must be of size 2"
    in
    let tr trans (k, n) = if trans = 1 then (n, k) else (k, n) in
    let a1, a2 = tr transA @@ rank2 inputA in
    let b1, b2 = tr transB @@ rank2 inputB in
    if not (Int.equal a2 b1)
    then
      Fmt.failwith "Gemm (M:%i,K:%i) (K:%i,N:%i) -> (M:%i,N:%i)" a1 a2 b1 b2 a1
        b2;
    Shape.of_array [| a1; b2 |]
  | ( LogSoftmax | Squeeze _ | MaxPool | Conv | Identity _ | RW_Linearized_ReLu
    | ReduceSum _ | GatherND _ | RandomNormal _ | Abs _ | Log _ ) as n ->
    failwith (Fmt.str "todo compute shape : %a" pp_descr n)

let compute_ty : _ -> ty = function
  | Constant { data = Float _ } -> Float
  | Constant { data = Int64 _ } -> Int64
  | _ -> Float

let create =
  let c = ref (-1) in
  fun descr ->
    Int.incr c;
    { id = !c; descr; shape = compute_shape_descr descr; ty = compute_ty descr }

let constant_int_array a =
  create (Constant { data = Gentensor.of_int64_array a })

let reshape shape node =
  if Shape.equal node.shape shape
  then node
  else
    create
      (Reshape
         {
           input = node;
           shape =
             constant_int_array
               (Array.map ~f:Int64.of_int @@ Shape.to_array shape);
         })

let gather_int_as_matmul input i =
  let input1 = reshape (Shape.of_array [| 1; Shape.size input.shape |]) input in
  let selector = Array.create ~len:(Shape.size input1.shape) Float.zero in
  Array.set selector i Float.one;
  let selector =
    Gentensor.Float
      (Tensor.of_array1
         (Shape.of_array [| Array.length selector; 1 |])
         (Bigarray.Array1.of_array Float64 C_layout selector))
  in
  let input2 = create (Constant { data = selector }) in
  let result = create (Matmul { input1; input2 }) in
  reshape (Shape.of_array [| 1 |]) result

let gather_int ?(encode = true) input i =
  if encode
  then gather_int_as_matmul input i
  else
    let indices =
      create (Constant { data = Gentensor.create_1_int64 (Int64.of_int i) })
    in
    create (Gather { input; indices; axis = 0 })

let mul_float input f =
  let input1 = reshape (Shape.of_array [| 1; 1 |]) input in
  let f = Array.create ~len:1 f in
  let f =
    Gentensor.Float
      (Tensor.of_array1
         (Shape.of_array [| Array.length f; 1 |])
         (Bigarray.Array1.of_array Float64 C_layout f))
  in
  let input2 = create (Constant { data = f }) in
  let result = create (Matmul { input1; input2 }) in
  reshape (Shape.of_array [| 1 |]) result

let div_float ?(encode = true) input f =
  if encode
  then
    let f = Float.one /. f in
    mul_float input f
  else
    let input1 = reshape (Shape.of_array [| 1; 1 |]) input in
    let f = Array.create ~len:1 f in
    let f =
      Gentensor.Float
        (Tensor.of_array1
           (Shape.of_array [| Array.length f; 1 |])
           (Bigarray.Array1.of_array Float64 C_layout f))
    in
    let input2 = create (Constant { data = f }) in
    let result = create (Div { input1; input2 }) in
    reshape (Shape.of_array [| 1 |]) result

let concat_0 = function
  | [ n ] -> n
  | [] -> failwith "empty concat"
  | inputs -> create (Concat { inputs; axis = 0 })

let preds node =
  match node.descr with
  | Constant _ | Input _ -> []
  | Add { input1; input2 }
  | Sub { input1; input2 }
  | Mul { input1; input2 }
  | Div { input1; input2 }
  | Matmul { input1; input2 } ->
    [ input1; input2 ]
  | Gather { input; indices; axis = _ } -> [ input; indices ]
  | GatherND { data; indices; batch_dims = _ } -> [ data; indices ]
  | ReLu { input } | Abs { input } | Log { input } -> [ input ]
  | Concat { inputs; axis = _ } -> inputs
  | ReduceSum { input; axes = Some x; _ } -> [ input; x ]
  | ReduceSum { input; axes = None; _ } -> [ input ]
  | RandomNormal _ -> []
  | Transpose { input; _ } -> [ input ]
  | Flatten { input; _ } -> [ input ]
  | Identity { input } -> [ input ]
  | Gemm { inputA; inputB; inputC = Some x; _ } -> [ inputA; inputB; x ]
  | Gemm { inputA; inputB; inputC = None; _ } -> [ inputA; inputB ]
  | Squeeze { data; _ } -> [ data ]
  | Reshape { input; shape; _ } -> [ input; shape ]
  | LogSoftmax | MaxPool | Conv | RW_Linearized_ReLu -> []

let map f n =
  match n.descr with
  | Constant _ | Input _ -> n
  | Add { input1; input2 } ->
    create (Add { input1 = f input1; input2 = f input2 })
  | Sub { input1; input2 } ->
    create (Sub { input1 = f input1; input2 = f input2 })
  | Mul { input1; input2 } ->
    create (Mul { input1 = f input1; input2 = f input2 })
  | Div { input1; input2 } ->
    create (Div { input1 = f input1; input2 = f input2 })
  | Matmul { input1; input2 } ->
    create (Matmul { input1 = f input1; input2 = f input2 })
  | ReLu { input } -> create (ReLu { input = f input })
  | Abs { input } -> create (Abs { input = f input })
  | Log { input } -> create (Log { input = f input })
  | RandomNormal _ as descr -> create descr
  | ReduceSum { input; axes; keepdims; noop_with_empty_axes } ->
    create (ReduceSum { input = f input; axes; keepdims; noop_with_empty_axes })
  | Gather { input; indices; axis } ->
    create (Gather { input = f input; indices = f indices; axis })
  | GatherND { data; indices; batch_dims } ->
    create (GatherND { data = f data; indices = f indices; batch_dims })
  | Transpose t -> create (Transpose { t with input = f t.input })
  | Flatten t -> create (Flatten { t with input = f t.input })
  | Identity { input } -> create (Identity { input = f input })
  | Concat { inputs; axis } ->
    create (Concat { inputs = List.map ~f inputs; axis })
  | Gemm t ->
    create
      (Gemm
         {
           t with
           inputA = f t.inputA;
           inputB = f t.inputB;
           inputC = Base.Option.map t.inputC ~f;
         })
  | Squeeze t -> create (Squeeze { t with data = f t.data })
  | Reshape t -> create (Reshape { t with input = f t.input })
  | LogSoftmax | MaxPool | Conv | RW_Linearized_ReLu -> n (* todo *)

(* let map_rec f node = let h = Base.Hashtbl.create (module Base.Int) in let rec
   aux n = Base.Hashtbl.find_or_add h n.id ~default:(fun () -> f (map aux n)) in
   aux node *)

let replace_input f node =
  let h = Base.Hashtbl.create (module Base.Int) in
  let rec aux n =
    Base.Hashtbl.find_or_add h n.id ~default:(fun () ->
      match n.descr with Input _ -> f () | _ -> map aux n)
  in
  aux node

(* iter on the nodes accessible from [node] ([node] comprised) without
   repetition *)
let map_rec f node =
  let h = Base.Hashtbl.create (module Base.Int) in
  let rec aux n =
    Base.Hashtbl.find_or_add h n.id ~default:(fun () -> f (map aux n))
  in
  aux node

let iter_rec f node =
  let h = Base.Hashtbl.create (module Base.Int) in
  let rec aux n =
    Base.Hashtbl.find_or_add h n.id ~default:(fun () ->
      List.iter ~f:aux (preds n);
      f n)
  in
  aux node