Support loading models by name (#38)

Co-authored-by: Matthew Tamayo-Rios <[email protected]>

README.md

@@ -123,7 +123,7 @@ detailed examples.
 
 ### Detailed design discussion
 
-For the details of the API, see [wasi-nn.wit.md](wasi-nn.wit.md).
+For the details of the API, see [wasi-nn.wit](wit/wasi-nn.wit).
 
 <!--
 This section should mostly refer to the .wit.md file that specifies the API. This section is for

wasi-nn.witx

@@ -2,6 +2,8 @@
 ;; version for the official specification and documentation.
 
 (typename $buffer_size u32)
+(typename $status u32)
+
 (typename $nn_errno
   (enum (@witx tag u16)
     $success
@@ -10,6 +12,9 @@
     $missing_memory
     $busy
     $runtime_error
+    $unsupported_operation
+    $model_too_large
+    $model_not_found
   )
 )
 (typename $tensor_dimensions (list u32))
@@ -39,6 +44,7 @@
     $tensorflow
     $pytorch
     $tensorflowlite
+    $autodetect
   )
 )
 (typename $execution_target
@@ -58,6 +64,17 @@
     (param $target $execution_target)
     (result $error (expected $graph (error $nn_errno)))
   )
+  ;;; Load an opaque sequence of bytes to use for inference.
+  ;;;
+  ;;; This allows runtime implementations to support multiple graph encoding formats. For unsupported graph encodings,
+  ;;; return `errno::inval`.
+  (@interface func (export "load_by_name")
+     ;;; The name of the model to load from the model registry
+     (param $model_name string)
+
+     (result $error (expected $graph (error $nn_errno)))
+  )
+
   (@interface func (export "init_execution_context")
     (param $graph $graph)
     (result $error (expected $graph_execution_context (error $nn_errno)))

wit/wasi-nn.wit

@@ -0,0 +1,93 @@
+package wasi:nn
+
+world inference {
+    import tensor
+    import graph
+    import execution
+    import errors
+}
+
+interface tensor {
+    type tensor-dimensions = list<u32>
+    type tensor-data = list<u8>
+
+    enum tensor-type {
+        fp16,
+        fp32,
+        bf16,
+        up8,
+        ip32
+    }
+
+    record tensor {
+        // Describe the size of the tensor (e.g., 2x2x2x2 -> [2, 2, 2, 2]). To represent a tensor
+        // containing a single value, use `[1]` for the tensor dimensions.
+        dimensions: tensor-dimensions,
+
+        // Describe the type of element in the tensor (e.g., f32).
+        tensor-type: tensor-type,
+
+        // Contains the tensor data.
+        data: tensor-data,
+    }
+}
+
+interface graph {
+    use errors.{error}
+    type graph-builder = list<u8>
+    type graph-builder-array = list<graph-builder>
+    use tensor.{tensor}
+
+    type graph = u32
+
+    enum graph-encoding {
+        openvino,
+        onnx,
+        tensorflow,
+        pytorch,
+        tensorflowlite,
+        autodetect,
+    }
+
+    enum execution-target {
+        cpu,
+        gpu,
+        tpu
+    }
+
+    load: func(builder: graph-builder-array, encoding: graph-encoding, target: execution-target) -> result<graph, error>
+    load-named-model: func(name: string) -> result<graph, error>
+}
+
+interface execution {
+    use errors.{error}
+    use tensor.{tensor, tensor-data}
+    use graph.{graph}
+
+    type graph-execution-context = u32
+    init-execution-context: func(graph: graph) -> result<graph-execution-context, error>
+    set-input: func(ctx: graph-execution-context, index: u32, tensor: tensor) -> result<_, error>
+    set-input-by-name: func(ctx: graph-execution-context, name: string, tensor: tensor) -> result<_, error>
+    compute: func(ctx: graph-execution-context) -> result<_, error>
+    get-output: func(ctx: graph-execution-context, index: u32) -> result<list<tensor-data>, error>
+    eval: func(tensors: list<tensor>) -> result<list<tensor>, error>
+
+}
+
+interface errors {
+    enum error {
+        // Caller module passed an invalid argument.
+        invalid-argument,
+        // Invalid encoding.
+        invalid-encoding,
+        busy,
+        // Runtime Error.
+        runtime-error,
+        // Unsupported operation
+        unsupported-operation,
+        // Model too large
+        model-too-large,
+        // Model not found
+        model-not-found
+    }
+}