Skip to content

Commit

Permalink
[WebNN EP] Add cache for MLContexts in the WebNNBackend (microsof…
Browse files Browse the repository at this point in the history
…t#22510)

### Description
This change adds a cache of `MLContext`s keyed by their options to the
`WebNNBackend`. This makes is so that multiple `InferenceSession`s
create with the same options will share the same context.

### Motivation and Context
Since `MLTensor`s are tied `MLContext`s, developer can't easily share
tensors between `InferenceSession` (outside of manually an `MLContext`
and specifying the `context` options). This leads strange behaviors such
as,
```js
const sessionsA = ort.InferenceSession.create(urlA, {
  executionProviders: ["webnn"],
  preferredOutputLocation: "ml-buffer",
});
const sessionsB = ort.InferenceSession.create(urlB, {
  executionProviders: ["webnn"],
});
const temp = await sessionA.run({/* arguments */});
const result = await sessionB.run({"input":temp["output"]}); // ERROR: Failed to execute 'dispatch' on 'MLContext': Invalid inputs: The context of MLGraph doesn't match the context of the MLTensor with name "input".
```
We encountered this behavior when updating the transformers.js version
in the developer preview demos. microsoft/webnn-developer-preview#46
  • Loading branch information
egalli authored and ankitm3k committed Dec 11, 2024
1 parent 54bbb2a commit 229f8f7
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 20 deletions.
17 changes: 0 additions & 17 deletions js/web/lib/wasm/wasm-types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -232,23 +232,6 @@ export declare namespace JSEP {
* @returns
*/
jsepCreateMLContext(optionsOrGpuDevice?: MLContextOptions | GPUDevice): Promise<MLContext>;

/**
* [exported from pre-jsep.js] Register a WebNN Constant operand from external data.
* @param externalFilePath - specify the external file path.
* @param dataOffset - specify the external data offset.
* @param dataLength - specify the external data length.
* @param builder - specify the MLGraphBuilder used for constructing the Constant.
* @param desc - specify the MLOperandDescriptor of the Constant.
* @returns the WebNN Constant operand for the specified external data.
*/
jsepRegisterMLConstant(
externalFilePath: string,
dataOffset: number,
dataLength: number,
builder: MLGraphBuilder,
desc: MLOperandDescriptor,
): MLOperand;
}
}

Expand Down
8 changes: 5 additions & 3 deletions onnxruntime/wasm/pre-jsep.js
Original file line number Diff line number Diff line change
Expand Up @@ -237,11 +237,13 @@ Module['jsepInit'] = (name, params) => {
}
Module['jsepRegisterMLTensor'] = (tensor, dataType, shape) => {
return backend['registerMLTensor'](tensor, dataType, shape);
}

};
Module['jsepCreateMLContext'] = (optionsOrGpuDevice) => {
return backend['createMLContext'](optionsOrGpuDevice);
};
Module.jsepRegisterMLConstant = (externalFilePath, dataOffset, dataLength, builder, desc) => {
return backend['registerMLConstant'](
externalFilePath, dataOffset, dataLength, builder, desc, Module.MountedFiles);
}
};
}
};

0 comments on commit 229f8f7

Please sign in to comment.