(* SPDX-License-Identifier: MPL-2.0 *)

(* SPDX-FileCopyrightText: 2026 Jonathan D.A. Jewell *)

(** ONNX Backend (MVP — proof of wire format).

Emits a binary [.onnx] [ModelProto] from a strict subset of AffineScript:

- one entry function (named [graph], [main], or the first [fn])

- parameters typed [Array[Float]] are graph inputs (1-D, dimension param)

- return type [Array[Float]] is the graph output

- body is a chain of [let] bindings whose RHS is one of:

* a variable reference (passes the SSA value forward)

* a call to a recognised ONNX op by name

and a final expression naming the output

The recognised op set is small: arithmetic ([Add], [Sub], [Mul], [Div]),

common activations ([Relu], [Sigmoid], [Tanh]), and identity-shaped ops

([Identity]). Each is matched by AffineScript function name (case-

insensitive lowercased — [add] maps to [Add], etc.). Unrecognised calls

produce a hard error so the regression is loud.

What this does NOT do (intentional, MVP scope):

- tensor shape inference

- broadcasting rules

- attribute encoding (axes, alpha, etc.)

- initializer tensors / weights

- any verification beyond "the bytes parse as ONNX"

Validation strategy: round-trip the output through any ONNX reader

(oxionnx-proto, onnxruntime, tract). All three accept the same wire

format; if the bytes decode they're ONNX-conformant.

*)

open Ast

open Kernel_sublang

(* ============================================================================

Op recognition

Map an AffineScript function name to (ONNX op type, expected arity).

============================================================================ *)

let recognise_op (name : string) : (string * int) option =

match String.lowercase_ascii name with

| "add" -> Some ("Add", 2)

| "sub" -> Some ("Sub", 2)

| "mul" -> Some ("Mul", 2)

| "div" -> Some ("Div", 2)

| "relu" -> Some ("Relu", 1)

| "sigmoid" -> Some ("Sigmoid", 1)

| "tanh" -> Some ("Tanh", 1)

| "neg" -> Some ("Neg", 1)

| "abs" -> Some ("Abs", 1)

| "exp" -> Some ("Exp", 1)

| "log" -> Some ("Log", 1)

| "sqrt" -> Some ("Sqrt", 1)

| "identity" -> Some ("Identity", 1)

| _ -> None

(* ============================================================================

Type validation

Graph inputs and outputs must be [Array[Float]]. We accept the AffineScript

surface forms [Array[Float]], [ref Array[Float]], [mut Array[Float]].

============================================================================ *)

(* [strip_ownership] now comes from Kernel_sublang. *)

let is_array_float (te : type_expr) : bool =

match array_element te with

| Some "Float" -> true

| _ -> false

(* ============================================================================

Build the ONNX graph from an AffineScript function body

We lower a chain of [let v = call(...)] bindings into a list of ONNX

nodes. The graph inputs are the function parameters, the output is the

final expression (which must be a variable reference).

============================================================================ *)

type build_state = {

mutable nodes : Onnx_proto.node list;

mutable next_id : int;

}

let fresh_node_name (st : build_state) (op : string) : string =

let id = st.next_id in

st.next_id <- id + 1;

Printf.sprintf "%s_%d" op id

(** Lower a value-producing expression into either a variable name (already

in scope) or a new node whose output we name and return. *)

let rec lower_expr (st : build_state) (e : expr) : string =

match e with

| ExprVar id -> id.name

| ExprSpan (inner, _) -> lower_expr st inner

| ExprApp (callee, args) ->

let fn_name = match callee with

| ExprVar id -> id.name

| _ -> unsupported "indirect calls not supported in ONNX backend"

in

let (op_type, expected_arity) = match recognise_op fn_name with

| Some pair -> pair

| None -> unsupported ("unknown ONNX op (no name match): " ^ fn_name)

in

let actual = List.length args in

if actual <> expected_arity then

unsupported

(Printf.sprintf "%s expects %d args, got %d" op_type expected_arity actual);

let arg_names = List.map (lower_expr st) args in

let out_name = fresh_node_name st op_type in

let node = {

Onnx_proto.n_input = arg_names;

n_output = [out_name];

n_name = out_name;

n_op_type = op_type;

n_domain = "";

} in

st.nodes <- node :: st.nodes;

out_name

| ExprLet { el_pat; el_value; el_body; el_mut = _; el_quantity = _; el_ty = _ } ->

(* Ignore the binder for the SSA-style lowering: each call already

produces a uniquely-named output. We propagate the *output of the

RHS* to the body, threading variables through scope by their

user-visible name. *)

let rhs_name = lower_expr st el_value in

let var = match el_pat with

| PatVar id -> id.name

| _ -> unsupported "destructuring let not supported in ONNX backend"

in

(* Rename the latest node's output to the bound name so the user's

identifier survives into the graph. Only safe when the RHS produced

a fresh node (not a passthrough variable). *)

(match st.nodes with

| latest :: rest when latest.n_output = [rhs_name] ->

let renamed = { latest with

Onnx_proto.n_output = [var];

n_name = var;

} in

st.nodes <- renamed :: rest

| _ ->

(* RHS was a variable — emit an Identity node so the alias is

real in the graph. *)

let node = {

Onnx_proto.n_input = [rhs_name];

n_output = [var];

n_name = var;

n_op_type = "Identity";

n_domain = "";

} in

st.nodes <- node :: st.nodes);

(match el_body with

| Some body -> lower_expr st body

| None -> unsupported "let without body cannot produce graph output")

| ExprBlock blk ->

List.fold_left (fun _last s ->

match s with

| StmtLet { sl_pat = PatVar id; sl_value; _ } ->

let rhs_name = lower_expr st sl_value in

(match st.nodes with

| latest :: rest when latest.n_output = [rhs_name] ->

st.nodes <- { latest with

Onnx_proto.n_output = [id.name]; n_name = id.name } :: rest

| _ ->

st.nodes <- {

Onnx_proto.n_input = [rhs_name]; n_output = [id.name];

n_name = id.name; n_op_type = "Identity"; n_domain = "";

} :: st.nodes);

id.name

| StmtExpr e -> lower_expr st e

| _ -> unsupported "only let-bindings and trailing expressions allowed in ONNX block"

) "" blk.blk_stmts

|> ignore;

(match blk.blk_expr with

| Some e -> lower_expr st e

| None -> unsupported "block must end with an expression naming the output")

| ExprLit _ ->

unsupported "literal not supported as ONNX value (need Constant op + initializer)"

| _ ->

unsupported "expression form not supported in ONNX kernel"

(* ============================================================================

Driver

============================================================================ *)

(* ONNX's canonical entry name is [graph]; fall back to main / first-fn via

Kernel_sublang's shared finder. *)

let onnx_pick_entry (program : program) : fn_decl =

pick_entry ~names:["graph"; "main"; "kernel"] program

let validate_entry (fd : fn_decl) : unit =

List.iter (fun (p : param) ->

if not (is_array_float p.p_ty) then

unsupported

(Printf.sprintf "parameter %s must be Array[Float]" p.p_name.name)

) fd.fd_params;

match fd.fd_ret_ty with

| Some t when is_array_float t -> ()

| None -> unsupported "graph function must declare a return type"

| _ -> unsupported "graph function must return Array[Float]"

(** Produce a ValueInfoProto for a parameter or result name. We use a single

dynamic dimension named [N] so consumers can pass any-length tensors. *)

let value_info_for (name : string) : Onnx_proto.value_info = {

vi_name = name;

vi_type = Onnx_proto.TensorType {

elem_type = 1; (* FLOAT *)

shape = [Onnx_proto.DimParam "N"];

};

}

let generate (program : program) (_symbols : Symbol.t) : string =

let entry = onnx_pick_entry program in

validate_entry entry;

let st = { nodes = []; next_id = 0 } in

let output_name = match entry.fd_body with

| FnExpr e -> lower_expr st e

| FnBlock b -> lower_expr st (ExprBlock b)

in

let inputs = List.map (fun (p : param) -> value_info_for p.p_name.name)

entry.fd_params in

let outputs = [value_info_for output_name] in

let graph = {

Onnx_proto.g_node = List.rev st.nodes;

g_name = entry.fd_name.name;

g_input = inputs;

g_output = outputs;

} in

let model = {

Onnx_proto.m_ir_version = 7; (* ONNX 1.10+ *)

m_producer_name = "affinescript";

m_producer_version = Version.value;

m_opset_import = [{ op_domain = ""; op_version = 13 }];

m_graph = graph;

} in

Onnx_proto.serialize_model model

let codegen_onnx (program : program) (symbols : Symbol.t) : (string, string) result =

try Ok (generate program symbols)

with

| Unsupported msg -> Error ("ONNX backend: " ^ msg)

| Failure msg -> Error ("ONNX codegen error: " ^ msg)

| e -> Error ("ONNX codegen error: " ^ Printexc.to_string e)