Source file writer.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
module Format = Stdlib.Format
module Fun = Stdlib.Fun
module Oproto = Onnx_protoc (* Autogenerated during compilation *)
module Oprotom = Oproto.Onnx.ModelProto

let value_info_from_tensor_shape ~name ~shape =
  let open Oproto.Onnx in
  let dim =
    List.map (Nir.Shape.to_list shape) ~f:(fun i ->
      TensorShapeProto.Dimension.make ~value:(`Dim_value (Int64.of_int i)) ())
  in
  let shape = TensorShapeProto.make ~dim () in
  let value =
    `Tensor_type
      (TypeProto.Tensor.make
         ~elem_type:AttributeProto.AttributeType.(to_int FLOAT)
         ~shape ())
  in
  let type' = TypeProto.make ~value () in
  let value_info = ValueInfoProto.make ~name ~type' () in
  value_info

let convert_into_tensor ?name (t : Nir.Gentensor.t) =
  let mk data_type =
    Oproto.Onnx.TensorProto.make
      ~data_type:Oproto.Onnx.TensorProto.DataType.(to_int data_type)
      ~dims:
        (List.map ~f:Int64.of_int @@ Nir.Shape.to_list @@ Nir.Gentensor.shape t)
      ?name
  in
  match t with
  | Float t -> mk FLOAT ~float_data:(Nir.Tensor.flatten t) ()
  | Int64 t -> mk INT64 ~int64_data:(Nir.Tensor.flatten t) ()

let default_opset_import =
  let open Oproto.Onnx in
  let onnx_domain = "" in
  OperatorSetIdProto.make ~domain:onnx_domain ~version:13L ()

let nir_to_onnx_protoc (nir : Nir.Ngraph.t) =
  let open Oproto.Onnx in
  let get_name (v : Nir.Node.t) = Int.to_string v.id in
  let protocs, (input, input_shape) =
    let acc = Queue.create () in
    let g_input = ref None in
    let vertex_to_protoc (v : Nir.Node.t) =
      let name = get_name v in
      let input = List.map ~f:get_name (Nir.Node.preds v) in
      let output = [ name ] in
      let make op_type attribute =
        Queue.enqueue acc
          (Oproto.Onnx.NodeProto.make ~input ~output ~name ~op_type ~attribute
             ~doc_string:"" ())
      in
      let mk_int name i =
        AttributeProto.make ~name ~type':INT ~i:(Int64.of_int i) ()
      in
      let mk_ints name ints =
        AttributeProto.make ~name ~type':INTS
          ~ints:(List.map ~f:Int64.of_int ints)
          ()
      in
      let mk_float name f = AttributeProto.make ~name ~type':FLOAT ~f () in
      let mk_tensor name t = AttributeProto.make ~name ~type':TENSOR ~t () in
      match v.descr with
      | LogSoftmax | Transpose _ | Squeeze _ | MaxPool | Conv | Identity _
      | RW_Linearized_ReLu | GatherND _ | ReduceSum _ ->
        Caisar_logging.Logging.not_implemented_yet (fun m ->
          m "Operator %a not implemented yet." Nir.Node.pp_descr v.descr)
      | Reshape _ -> make "Reshape" []
      | Flatten { axis; _ } -> make "Flatten" [ mk_int "axis" axis ]
      | Constant { data } ->
        let data = convert_into_tensor data in
        make "Constant" [ mk_tensor "value" data ]
      | Add _ -> make "Add" []
      | Sub _ -> make "Sub" []
      | Mul _ -> make "Mul" []
      | Div _ -> make "Div" []
      | Matmul _ -> make "MatMul" []
      | ReLu _ -> make "Relu" []
      | Input { shape } -> g_input := Some (v, shape)
      | Concat { axis; _ } -> make "Concat" [ mk_int "axis" axis ]
      | Gather { axis; _ } -> make "Gather" [ mk_int "axis" axis ]
      | Abs _ -> make "Abs" []
      | Log _ -> make "Log" []
      | RandomNormal { dtype; mean; scale; seed; shape } ->
        make "RandomNormal"
          [
            mk_int "dtype" dtype;
            mk_float "mean" mean;
            mk_float "scale" scale;
            mk_float "seed" seed;
            mk_ints "shape" (Array.to_list shape);
          ]
      | Gemm { alpha; beta; transA; transB; _ } ->
        make "Gemm"
          [
            mk_float "alpha" alpha;
            mk_float "beta" beta;
            mk_int "transA" transA;
            mk_int "transB" transB;
          ]
    in
    Nir.Ngraph.iter_vertex vertex_to_protoc nir;
    (Queue.to_list acc, Option.value_exn !g_input)
  in
  let docstr =
    "This ONNX model was generated from the Neural Intermediate Representation \
     of CAISAR"
  in
  let input =
    [ value_info_from_tensor_shape ~name:(get_name input) ~shape:input_shape ]
  in
  let output =
    let output = Nir.Ngraph.output nir in
    [
      value_info_from_tensor_shape ~name:(get_name output)
        ~shape:(Nir.Node.compute_shape output);
    ]
  in
  let protog =
    GraphProto.make ~name:"ONNX CAISAR Export" ~node:protocs ~initializer':[]
      ~sparse_initializer:[] ~doc_string:"ONNX graph generated from CAISAR NIR"
      ~input ~output ~value_info:[] ~quantization_annotation:[] ()
  in
  let protom =
    Oproto.Onnx.ModelProto.make ~ir_version:8L
      ~opset_import:[ default_opset_import ] ~producer_name:"CAISAR"
      ~producer_version:"1.0" ~domain:"" ~model_version:(-1L) ~doc_string:docstr
      ~graph:protog ~metadata_props:[] ~training_info:[] ~functions:[] ()
  in
  let writer = Oprotom.to_proto protom in
  Ocaml_protoc_plugin.Writer.contents writer

let write_to_onnx nir out_channel =
  let onnx = nir_to_onnx_protoc nir in
  Stdio.Out_channel.output_string out_channel onnx

let to_file nir filename =
  let out_chan = Stdlib.open_out filename in
  write_to_onnx nir out_chan;
  Stdlib.close_out out_chan