Source file writer.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

(** The name of a node in the onnx is directly its id *)
let get_name (v : Nir.Node.t) = Int.to_string v.id

let value_info_from_tensor_shape (node : Nir.Node.t) =
  let open Oproto.Onnx in
  let dim =
    List.map (Nir.Shape.to_list node.shape) ~f:(fun i ->
      TensorShapeProto.Dimension.make ~value:(`Dim_value (Int64.of_int i)) ())
  in
  let shape = TensorShapeProto.make ~dim () in
  let elem_type =
    match node.ty with
    | Nir.Node.BFloat16 -> TensorProto.DataType.BFLOAT16
    | Float16 -> FLOAT16
    | Float -> FLOAT
    | UInt8 -> UINT8
    | Int8 -> INT8
    | Int32 -> INT32
    | Int64 -> INT64
  in
  let value = `Tensor_type (TypeProto.Tensor.make ~elem_type ~shape ()) in
  let type' = TypeProto.make ~value () in
  let value_info = ValueInfoProto.make ~name:(get_name node) ~type' () in
  value_info

let convert_into_tensor ?name (t : Nir.Gentensor.t) =
  let mk data_type =
    Oproto.Onnx.TensorProto.make ~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) ()
  | Int8 t ->
    mk INT8
      ~int32_data:(List.map ~f:Int32.of_int_exn @@ Nir.Tensor.flatten t)
      ()
  | UInt8 t ->
    mk UINT8
      ~int32_data:(List.map ~f:Int32.of_int_exn @@ Nir.Tensor.flatten t)
      ()
  | Int32 t -> mk INT32 ~int32_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 protocs, input =
    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
      let mk_bool name b =
        AttributeProto.make ~name ~type':INT
          ~i:(Int64.of_int @@ if b then 1 else 0)
          ()
      in
      match v.descr with
      | LogSoftmax | Transpose _ | Squeeze _ | MaxPool | Conv | Identity _
      | RW_Linearized_ReLu | GatherND _ | ReduceSum _ ->
        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" []
      | QLinearMatMul _ -> make "QLinearMatMul" []
      | ReLu _ -> make "Relu" []
      | Input _ -> g_input := Some v
      | 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;
          ]
      | QGemm { alpha; transA; transB; _ } ->
        make "QGemm"
          [
            mk_float "alpha" alpha;
            mk_int "transA" transA;
            mk_int "transB" transB;
          ]
      | Sign _ -> make "Sign" []
      | ArgMax { axis; keepdims; _ } ->
        make "ArgMax" [ mk_int "axis" axis; mk_bool "keepdims" keepdims ]
      | Pow _ -> make "Pow" []
      | QuantizeLinear { axis; _ } ->
        make "QuantizeLinear" [ mk_int "axis" axis ]
      | DequantizeLinear { axis; _ } ->
        make "DequantizeLinear" [ mk_int "axis" axis ]
    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 input ] in
  let output =
    let output = Nir.Ngraph.output nir in
    [ value_info_from_tensor_shape output ]
  in
  let value_info =
    List.map ~f:value_info_from_tensor_shape (Nir.Ngraph.nodes nir)
  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:Config.version ~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