// Copyright 2024 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

module webnn.mojom;

import "gpu/ipc/common/mailbox.mojom";
import "gpu/ipc/common/sync_token.mojom";
import "services/webnn/public/mojom/features.mojom";
import "services/webnn/public/mojom/webnn_compiler_context.mojom";
import "services/webnn/public/mojom/webnn_tensor.mojom";
import "services/webnn/public/mojom/webnn_error.mojom";
import "services/webnn/public/mojom/webnn_graph_builder.mojom";
import "third_party/blink/public/mojom/tokens/tokens.mojom";

// Represents a successful call to `WebNNContext::CreateTensor()`.
struct CreateTensorSuccess {
  pending_associated_remote<WebNNTensor> tensor_remote;
  // tensor_handle is a generated token used as a handle to identify the tensor
  // from the renderer. The token is only valid for the lifetime
  // of the tensor and is used by context operations in the service using the
  // tensor corresponding to this handle.
  blink.mojom.WebNNTensorToken tensor_handle;
};

// Represents the return value of `WebNNContext::CreateTensor()`. Let it be
// `success` if the tensor was successfully created and `error` otherwise.
union CreateTensorResult {
  CreateTensorSuccess success;
  Error error;
};

// Represents the `MLContext` object in the WebIDL definition that is a global
// state of neural network compute workload and execution processes. This
// interface runs in the GPU process and is called from the renderer process.
[RuntimeFeature=webnn.mojom.features.kWebMachineLearningNeuralNetwork]
interface WebNNContext {
  // Creates a connection to an `MLGraphBuilder` in the WebNN service.
  CreateGraphBuilder(pending_receiver<WebNNGraphBuilder> receiver);

  // Called by the renderer process to create `WebNNTensor` message pipe for
  // creating platform specific tensors, the WebNN tensor will be validated and
  // created. This method guarantees memory allocation on the device.
  CreateTensor(TensorInfo tensor_info)
    => (CreateTensorResult result);

  // Similar to CreateTensor() above, except that it creates the tensor from
  // an existing shared image. The tensor must have WebGPUInterop usage.
  // The mailbox must reference a valid shared image, and the provided SyncToken
  // must be verified and obtained from the SharedImageInterface.
  CreateTensorFromMailbox(TensorInfo tensor_info, gpu.mojom.Mailbox mailbox,
   gpu.mojom.SyncToken fence) => (CreateTensorResult result);

  // Execute the compiled platform graph identified by `graph`. The graph is
  // identified by a token rather than via the WebNNGraph interface to preserve
  // message ordering with ReadTensor/WriteTensor: both Dispatch and
  // ReadTensor/WriteTensor are on the WebNNContext pipe, so the service
  // processes them in order. If Dispatch were on a separate (non-associated)
  // WebNNGraph pipe, the service could receive Dispatch before a preceding
  // WriteTensor, executing inference with stale input data.
  Dispatch(blink.mojom.WebNNGraphToken graph,
           map<string, blink.mojom.WebNNTensorToken> named_inputs,
           map<string, blink.mojom.WebNNTensorToken> named_outputs);

  // Explicitly destroy a graph, removing it from the context. This method is
  // on the context pipe (rather than relying on WebNNGraph pipe disconnect)
  // to ensure ordering with Dispatch/ReadTensor/WriteTensor: preceding
  // operations on the context pipe are guaranteed to be processed before the
  // graph is destroyed.
  DestroyGraph(blink.mojom.WebNNGraphToken graph_handle);

  // Requests a new CompilerContext from the Compiler process. The renderer
  // creates the pipe, binds the remote immediately, and sends the
  // receiver here. On failure, the pipe disconnects (triggering the
  // renderer's disconnect handler).
  RequestCompilerContext(
      pending_receiver<WebNNCompilerContext> compiler_context_receiver);
};