Source file reader.ml

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

exception ParseError of string

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

(* ONNX format handling. *)
module Convert : sig
  val nir_of_onnx_protoc : Oproto.Onnx.ModelProto.t -> Nir.Ngraph.t
  val get_input_output_dim : Oproto.Onnx.ModelProto.t -> int * int
end = struct
  let get_shape_of_dims (s : Oproto.Onnx.TensorShapeProto.t) =
    Nir.Shape.of_list
    @@ List.map s ~f:(function
         | { value = `Dim_value v; _ } -> Int64.to_int_exn v
         | { value = `Dim_param s; _ } ->
           failwith (Fmt.str "Parameteric shape %s" s)
         | { value = `not_set; _ } -> failwith "Part of a shape not set")

  let get_shape_of_value (s : Oproto.Onnx.ValueInfoProto.t) =
    match s with
    | { type' = Some { value = `Tensor_type { shape = Some v; _ }; _ }; _ } ->
      get_shape_of_dims v
    | _ -> failwith "Value as not shape"

  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 (s : Oproto.Onnx.TensorShapeProto.Dimension.t list) =
    Nir.Shape.size (get_shape_of_dims s)

  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 convert_tensor (ts : Oproto.Onnx.TensorProto.t) : Nir.Gentensor.t =
    let dims = Nir.Shape.of_list @@ List.map ~f:Int64.to_int_exn ts.dims in
    let size = Nir.Shape.size dims in
    let read_raw ~get kind =
      match ts.raw_data with
      | None ->
        failwith "TensorProto have no data field for the given data type"
      | Some data ->
        let t = Bigarray.(Array1.create kind c_layout size) in
        for i = 0 to size - 1 do
          let v = get data i in
          Bigarray.Array1.set t i v
        done;
        Nir.Tensor.of_array1 dims t
    in
    let read_gen ~get ~conv elt_size kind custom_data =
      match custom_data with
      | [] ->
        read_raw kind ~get:(fun raw coord_in_data ->
          let offset = elt_size * coord_in_data in
          get raw offset)
      | l when List.length l <> size ->
        failwith "not enough data according to dimension"
      | l ->
        let t = Bigarray.(Array1.create kind c_layout size) in
        List.iteri l ~f:(fun i f -> Bigarray.Array1.set t i (conv f));
        Nir.Tensor.of_array1 dims t
    in
    match ts.data_type with
    | None -> failwith "TensorProto should have a type"
    | Some ty -> (
      match ty with
      | UNDEFINED -> failwith "Invalid UNDEFINED data type"
      | FLOAT ->
        (*ambiguity depending on the handled IR whether float are float32 or
          float64: we represent both as float64 *)
        Float
          (read_gen 4 Float64 ts.float_data
             ~get:EndianBytes.LittleEndian.get_float ~conv:Fun.id)
      | UINT8 ->
        UInt8
          (read_gen 1 Int8_unsigned ts.int32_data
             ~get:EndianBytes.LittleEndian.get_uint8 ~conv:Int32.to_int_exn)
      | INT8 ->
        Int8
          (read_gen 1 Int8_signed ts.int32_data
             ~get:EndianBytes.LittleEndian.get_int8 ~conv:Int32.to_int_exn)
      | INT32 ->
        Int32
          (read_gen 4 Int32 ts.int32_data
             ~get:EndianBytes.LittleEndian.get_int32 ~conv:Fun.id)
      | INT64 ->
        Int64
          (read_gen 8 Int64 ts.int64_data
             ~get:EndianBytes.LittleEndian.get_int64 ~conv:Fun.id)
      | DOUBLE ->
        Float
          (read_gen 4 Float64 ts.float_data
             ~get:EndianBytes.LittleEndian.get_float ~conv:Fun.id)
      | UINT16 | INT16 | STRING | BOOL | FLOAT16 | UINT32 | UINT64 | COMPLEX64
      | COMPLEX128 | BFLOAT16 ->
        failwith "Unsupported data type")

  type value =
    | Node of Oproto.Onnx.NodeProto.t
    | Tensor of Oproto.Onnx.TensorProto.t

  let produce_cfg (g : Oproto.Onnx.GraphProto.t) =
    let open Oproto.Onnx in
    let converted = Hashtbl.create (module String) in
    (* Associate output to node or initializer *)
    let of_output_value = Hashtbl.create (module String) in
    List.iter g.node ~f:(fun n ->
      match n.output with
      | [] -> failwith "Node without output"
      | [ o ] ->
        (* outputs must be uniq *)
        Hashtbl.add_exn of_output_value ~key:o ~data:(Node n)
      | _ -> failwith "Node with multiple outputs are not handled");
    List.iter g.initializer' ~f:(fun t ->
      match t.name with
      | None -> failwith "Initializer must have a name"
      | Some o ->
        (* outputs must be uniq *)
        Hashtbl.add_exn of_output_value ~key:o ~data:(Tensor t));
    assert (List.is_empty g.sparse_initializer);
    (* compute main output and input *)
    let output, output_shape =
      match g.output with
      | [] -> failwith "graph without output"
      | [ output ] -> (Option.value_exn output.name, get_shape_of_value output)
      | _ -> failwith "graph with more than one output"
    in
    (* Add input in already converted *)
    let input_name, input_shape =
      let input =
        List.filter_map g.input ~f:(fun i ->
          let name = Option.value_exn i.name in
          if Hashtbl.mem of_output_value name
          then None
          else Some (name, get_shape_of_value i))
      in
      match input with
      | [] -> failwith "graph without input, can be accepted"
      | [ input ] -> input
      | _ -> failwith "graph with more than one input node (unsupported)"
    in
    Hashtbl.add_exn converted ~key:input_name
      ~data:(Nir.Node.create (Input { shape = input_shape }));
    (* converter *)
    let rec convert output =
      Hashtbl.findi_or_add ~default:convert_aux converted output
    and convert_aux output =
      let value = Hashtbl.find_exn of_output_value output in
      match value with
      | Node n ->
        let one_arg = function
          | [ input ] -> input
          | _ -> failwith "should have one argument"
        in
        let two_arg = function
          | [ input1; input2 ] -> (input1, input2)
          | _ -> failwith "should have two arguments"
        in
        let attrs =
          Hashtbl.of_alist_exn
            (module String)
            (List.map ~f:(fun a -> (Option.value_exn a.name, a)) n.attribute)
        in
        let get_attr ?default name m =
          match Hashtbl.find attrs name with
          | Some v -> m v
          | None -> (
            match default with
            | Some v -> v
            | None -> Fmt.failwith "Required attribute %s missing" name)
        in
        let get_float ?default name : float =
          get_attr ?default name (function
            | { type' = Some AttributeProto.AttributeType.FLOAT; f = Some f; _ }
              ->
              f
            | _ -> failwith "Attribute wrongly typed")
        in
        let get_int ?default name : int =
          get_attr ?default name (function
            | { type' = Some AttributeProto.AttributeType.INT; i = Some i; _ }
              ->
              Int64.to_int_exn i
            | _ -> failwith "Attribute wrongly typed")
        in
        let get_ints ?default name : int list =
          get_attr ?default name (function
            | { type' = Some AttributeProto.AttributeType.INTS; ints = l; _ } ->
              List.map ~f:Int64.to_int_exn l
            | _ -> failwith "Attribute wrongly typed")
        in
        let get_tensor ?default name : Nir.Gentensor.t =
          get_attr ?default name (function
            | {
                type' = Some AttributeProto.AttributeType.TENSOR;
                t = Some t;
                _;
              } ->
              convert_tensor t
            | _ -> failwith "Attribute wrongly typed")
        in
        let get_bool ?default name = get_int ?default name = 1 in
        let n' =
          match n.op_type with
          | None -> failwith "Node without op_type (No-op?)"
          | Some s -> (
            match s with
            | "Add" ->
              let input1, input2 = two_arg n.input in
              Nir.Node.Add { input1 = convert input1; input2 = convert input2 }
            | "Sub" ->
              let input1, input2 = two_arg n.input in
              Nir.Node.Sub { input1 = convert input1; input2 = convert input2 }
            | "Mul" ->
              let input1, input2 = two_arg n.input in
              Nir.Node.Mul { input1 = convert input1; input2 = convert input2 }
            | "Div" ->
              let input1, input2 = two_arg n.input in
              Nir.Node.Div { input1 = convert input1; input2 = convert input2 }
            | "Pow" ->
              let input1, input2 = two_arg n.input in
              Nir.Node.Pow { input1 = convert input1; input2 = convert input2 }
            | "Relu" ->
              let input1 = one_arg n.input in
              Nir.Node.ReLu { input = convert input1 }
            | "MatMul" ->
              let input1, input2 = two_arg n.input in
              Nir.Node.Matmul
                { input1 = convert input1; input2 = convert input2 }
            | "QLinearMatMul" ->
              let ( inputA,
                    inputA_scale,
                    inputA_zero_point,
                    inputB,
                    inputB_scale,
                    inputB_zero_point,
                    y_scale,
                    y_zero_point ) =
                (* Awful pattern-matching to avoid multiple linear list
                   scanning. Is this actually needed/useful? *)
                match n.input with
                | [
                 inputA;
                 inputA_scale;
                 inputA_zero_point;
                 inputB;
                 inputB_scale;
                 inputB_zero_point;
                 y_scale;
                 y_zero_point;
                ] ->
                  ( inputA,
                    inputA_scale,
                    inputA_zero_point,
                    inputB,
                    inputB_scale,
                    inputB_zero_point,
                    y_scale,
                    y_zero_point )
                | _ -> failwith "QLinearMatMul must have 8 inputs"
              in
              Nir.Node.QLinearMatMul
                {
                  inputA = convert inputA;
                  inputA_scale = convert inputA_scale;
                  inputA_zero_point = convert inputA_zero_point;
                  inputB = convert inputB;
                  inputB_scale = convert inputB_scale;
                  inputB_zero_point = convert inputB_zero_point;
                  y_scale = convert y_scale;
                  y_zero_point = convert y_zero_point;
                }
            | "Gemm" ->
              let inputA, inputB, inputC =
                match n.input with
                | [ inputA; inputB ] -> (inputA, inputB, None)
                | [ inputA; inputB; inputC ] -> (inputA, inputB, Some inputC)
                | _ -> failwith "Gemm must have 2 or 3 inputs"
              in
              Nir.Node.Gemm
                {
                  inputA = convert inputA;
                  inputB = convert inputB;
                  inputC = Option.map ~f:convert inputC;
                  alpha = get_float ~default:1.0 "alpha";
                  beta = get_float ~default:1.0 "beta";
                  transA = get_int ~default:0 "transA";
                  transB = get_int ~default:0 "transB";
                }
            | "QGemm" ->
              let ( inputA,
                    inputA_scale,
                    inputA_zero_point,
                    inputB,
                    inputB_scale,
                    inputB_zero_point,
                    inputC,
                    y_scale,
                    y_zero_point ) =
                (* Awful pattern-matching to avoid multiple linear list
                   scanning. Is this actually needed/useful? *)
                match n.input with
                | [
                 inputA;
                 inputA_scale;
                 inputA_zero_point;
                 inputB;
                 inputB_scale;
                 inputB_zero_point;
                ] ->
                  ( inputA,
                    inputA_scale,
                    inputA_zero_point,
                    inputB,
                    inputB_scale,
                    inputB_zero_point,
                    None,
                    None,
                    None )
                | [
                 inputA;
                 inputA_scale;
                 inputA_zero_point;
                 inputB;
                 inputB_scale;
                 inputB_zero_point;
                 inputC;
                ] ->
                  ( inputA,
                    inputA_scale,
                    inputA_zero_point,
                    inputB,
                    inputB_scale,
                    inputB_zero_point,
                    Some inputC,
                    None,
                    None )
                | [
                 inputA;
                 inputA_scale;
                 inputA_zero_point;
                 inputB;
                 inputB_scale;
                 inputB_zero_point;
                 inputC;
                 y_scale;
                ] ->
                  ( inputA,
                    inputA_scale,
                    inputA_zero_point,
                    inputB,
                    inputB_scale,
                    inputB_zero_point,
                    Some inputC,
                    Some y_scale,
                    None )
                | [
                 inputA;
                 inputA_scale;
                 inputA_zero_point;
                 inputB;
                 inputB_scale;
                 inputB_zero_point;
                 inputC;
                 y_scale;
                 y_zero_point;
                ] ->
                  ( inputA,
                    inputA_scale,
                    inputA_zero_point,
                    inputB,
                    inputB_scale,
                    inputB_zero_point,
                    Some inputC,
                    Some y_scale,
                    Some y_zero_point )
                | _ -> failwith "QGemm must have between 6 and 9 inputs"
              in
              let inputA = convert inputA in
              let inputA_scale = convert inputA_scale in
              let inputA_zero_point = convert inputA_zero_point in
              let inputB = convert inputB in
              let inputB_scale = convert inputB_scale in
              let inputB_zero_point = convert inputB_zero_point in
              let inputC = Option.map ~f:convert inputC in
              let y_scale = Option.map ~f:convert y_scale in
              let y_zero_point = Option.map ~f:convert y_zero_point in
              let alpha = get_float ~default:1.0 "alpha" in
              let transA = get_int ~default:0 "transA" in
              let transB = get_int ~default:0 "transB" in
              Nir.Node.QGemm
                {
                  inputA;
                  inputA_scale;
                  inputA_zero_point;
                  inputB;
                  inputB_scale;
                  inputB_zero_point;
                  inputC;
                  y_scale;
                  y_zero_point;
                  alpha;
                  transA;
                  transB;
                }
            | "LogSoftmax" -> Nir.Node.LogSoftmax
            | "Transpose" ->
              Nir.Node.Transpose
                { input = convert (one_arg n.input); perm = get_ints "perm" }
            | "Squeeze" ->
              let data, axes =
                match n.input with
                | [ data ] -> (convert data, None)
                | [ data; axes ] -> (convert data, Some (convert axes))
                | _ -> failwith "Squeeze must have 1 or 2 inputs"
              in
              Nir.Node.Squeeze { data; axes }
            | "MaxPool" -> MaxPool
            | "Constant" -> Constant { data = get_tensor "value" }
            | "Conv" -> Conv
            | "Flatten" ->
              Flatten
                { input = convert @@ one_arg n.input; axis = get_int "axis" }
            | "Reshape" ->
              let input, shape = two_arg n.input in
              Reshape { input = convert input; shape = convert shape }
            | "Abs" -> Nir.Node.Abs { input = convert @@ one_arg n.input }
            | "Log" -> Nir.Node.Log { input = convert @@ one_arg n.input }
            | "RandomNormal" ->
              Nir.Node.RandomNormal
                {
                  dtype = get_int "dtype";
                  mean = get_float "mean";
                  scale = get_float "scale";
                  seed = get_float "seed";
                  shape = Array.of_list (get_ints "shape");
                }
            | "ArgMax" ->
              ArgMax
                {
                  input = convert @@ one_arg n.input;
                  axis = get_int "axis";
                  keepdims = get_bool ~default:1 "keepdims";
                }
            | "Concat" ->
              Concat
                {
                  inputs = n.input |> List.map ~f:(fun i -> convert i);
                  axis = get_int "axis";
                }
            | "Sign" -> Sign { input = convert @@ one_arg n.input }
            | "QuantizeLinear" ->
              let x, y_scale, y_zero_point =
                match n.input with
                | [ x; y_scale ] -> (x, y_scale, None)
                | [ x; y_scale; y_zero_point ] -> (x, y_scale, Some y_zero_point)
                | _ -> failwith "QuantizeLinear must have 2 or 3 inputs"
              in
              Nir.Node.QuantizeLinear
                {
                  x = convert x;
                  y_scale = convert y_scale;
                  y_zero_point = Option.map ~f:convert y_zero_point;
                  axis = get_int ~default:1 "axis";
                }
            | "DequantizeLinear" ->
              let x, x_scale, x_zero_point =
                match n.input with
                | [ x; x_scale ] -> (x, x_scale, None)
                | [ x; x_scale; x_zero_point ] -> (x, x_scale, Some x_zero_point)
                | _ -> failwith "DequantizeLinear must have 2 or 3 inputs"
              in
              Nir.Node.DequantizeLinear
                {
                  x = convert x;
                  x_scale = convert x_scale;
                  x_zero_point = Option.map ~f:convert x_zero_point;
                  axis = get_int ~default:1 "axis";
                }
            (* TODO: ReduceSum, GatherND *)
            | s -> failwith (Printf.sprintf "Unknown operators %s" s))
        in
        Nir.Node.create n'
      | Tensor t -> Nir.Node.create (Constant { data = convert_tensor t })
    in
    let output' = convert output in
    assert (Nir.Shape.equal output'.shape output_shape);
    Nir.Ngraph.create output'

  let nir_of_onnx_protoc (model : Oprotom.t) =
    (match model.ir_version with
    | None -> failwith "IR version not specified"
    | Some (3L | 4L | 5L | 6L | 7L | 8L | 9L | 10L) -> ()
    | Some i -> failwith (Printf.sprintf "Unsupported IR version %Li" i));
    assert (not (List.is_empty model.opset_import));
    if false
    then
      List.iter model.opset_import ~f:(fun opset ->
        Format.printf "opset:%s (%Li)@."
          (Option.value ~default:"" opset.domain)
          (Option.value_exn opset.version));
    match model.graph with
    | Some g -> produce_cfg g
    | None -> raise (ParseError "No graph in ONNX input file found")
end

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 = Convert.get_input_output_dim r in
      let nir =
        try Ok (Convert.nir_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; nir }
    | _ -> Error "Cannot read protobuf"
  with
  | Sys_error s -> Error s
  | Failure msg -> Error (Format.sprintf "Unexpected error: %s" msg)

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