cix_p1_tpu

defmodule CixP1.Graph do
  @moduledoc """
  A compiled model graph loaded onto the NPU.

  Load an offline-built graph binary (`.cix`, produced by the NOE Compiler or
  pulled from the CIX AI Model Hub) with `load/2`, then inspect its input and
  output tensors via the descriptor functions. Run inference through
  `CixP1.Job` or the `CixP1.run/2` convenience.
  """

  alias CixP1.{Context, Nif}

  @enforce_keys [:resource, :context, :path]
  defstruct [:resource, :context, :path]

  @type t :: %__MODULE__{resource: reference(), context: Context.t(), path: Path.t()}

  @typedoc "Tensor name and logical shape, read from the `.cix` model file."
  @type shape_info :: %{name: binary(), shape: [integer()]}

  @typedoc """
  A tensor descriptor as reported by libnoe.

  * `:id` — tensor id within the graph
  * `:size` — buffer size in bytes
  * `:scale` / `:zero_point` — quantization parameters (`real = (q - zero_point) * scale`)
  * `:data_type` — element type atom (`:u8`, `:s8`, `:s16`, `:s32`, `:f16`, `:f32`, …)
    or the raw integer for types without an atom mapping
  """
  @type descriptor :: %{
          id: non_neg_integer(),
          size: non_neg_integer(),
          scale: float(),
          zero_point: integer(),
          data_type: atom() | non_neg_integer()
        }

  @doc """
  Loads a compiled graph binary from `path` into `context`.
  """
  @spec load(Context.t(), Path.t()) :: {:ok, t()} | {:error, String.t()}
  def load(%Context{resource: ctx} = context, path) when is_binary(path) do
    case Nif.graph_load(ctx, path) do
      {:ok, resource} ->
        {:ok, %__MODULE__{resource: resource, context: context, path: path}}

      {:error, reason} ->
        {:error, to_string(reason)}
    end
  end

  @doc "Number of input tensors the graph expects."
  @spec input_count(t()) :: {:ok, non_neg_integer()} | {:error, String.t()}
  def input_count(%__MODULE__{resource: g}), do: count(g, Nif.tensor_type_input())

  @doc "Number of output tensors the graph produces."
  @spec output_count(t()) :: {:ok, non_neg_integer()} | {:error, String.t()}
  def output_count(%__MODULE__{resource: g}), do: count(g, Nif.tensor_type_output())

  @doc "Descriptor for input tensor `index` (0-based)."
  @spec input_descriptor(t(), non_neg_integer()) ::
          {:ok, descriptor()} | {:error, String.t()}
  def input_descriptor(%__MODULE__{resource: g}, index),
    do: descriptor(g, Nif.tensor_type_input(), index)

  @doc "Descriptor for output tensor `index` (0-based)."
  @spec output_descriptor(t(), non_neg_integer()) ::
          {:ok, descriptor()} | {:error, String.t()}
  def output_descriptor(%__MODULE__{resource: g}, index),
    do: descriptor(g, Nif.tensor_type_output(), index)

  @doc """
  Name and logical shape of input tensor `index` (0-based).

  Descriptors carry only element type and flat byte size; the shape is read
  from the `.cix` model file on disk, so the file must still be present at the
  path the graph was loaded from.
  """
  @spec input_shape(t(), non_neg_integer()) ::
          {:ok, shape_info()} | {:error, String.t()}
  def input_shape(%__MODULE__{} = graph, index),
    do: shape(graph, &input_count/1, Nif.tensor_type_input(), index)

  @doc "Name and logical shape of output tensor `index` (0-based). See `input_shape/2`."
  @spec output_shape(t(), non_neg_integer()) ::
          {:ok, shape_info()} | {:error, String.t()}
  def output_shape(%__MODULE__{} = graph, index),
    do: shape(graph, &output_count/1, Nif.tensor_type_output(), index)

  @doc "All input descriptors, in order."
  @spec input_descriptors(t()) :: {:ok, [descriptor()]} | {:error, String.t()}
  def input_descriptors(%__MODULE__{} = graph),
    do: descriptors(graph, &input_count/1, &input_descriptor/2)

  @doc "All output descriptors, in order."
  @spec output_descriptors(t()) :: {:ok, [descriptor()]} | {:error, String.t()}
  def output_descriptors(%__MODULE__{} = graph),
    do: descriptors(graph, &output_count/1, &output_descriptor/2)

  defp count(g, type) do
    case Nif.graph_tensor_count(g, type) do
      {:ok, cnt} -> {:ok, cnt}
      {:error, reason} -> {:error, to_string(reason)}
    end
  end

  # The shape NIF is path-based (libnoe reads the .cix directly), so it can't
  # cheaply bounds-check the index in C. Validate it against the loaded graph's
  # tensor count here so a wild index never reaches libnoe's file parser.
  defp shape(%__MODULE__{path: path} = graph, count_fun, type, index)
       when is_integer(index) and index >= 0 do
    with {:ok, n} <- count_fun.(graph) do
      if index < n do
        case Nif.graph_tensor_shape(path, type, index) do
          {:ok, info} -> {:ok, info}
          {:error, reason} -> {:error, to_string(reason)}
        end
      else
        {:error, "tensor_index_out_of_range"}
      end
    end
  end

  defp descriptor(g, type, index) do
    case Nif.graph_tensor_descriptor(g, type, index) do
      {:ok, desc} -> {:ok, desc}
      {:error, reason} -> {:error, to_string(reason)}
    end
  end

  defp descriptors(graph, count_fun, desc_fun) do
    with {:ok, n} <- count_fun.(graph) do
      Enum.reduce_while(0..(n - 1)//1, {:ok, []}, fn i, {:ok, acc} ->
        case desc_fun.(graph, i) do
          {:ok, desc} -> {:cont, {:ok, [desc | acc]}}
          {:error, _} = err -> {:halt, err}
        end
      end)
      |> case do
        {:ok, acc} -> {:ok, Enum.reverse(acc)}
        err -> err
      end
    end
  end
end