saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[webgpu]: add a simple GetCapability implementation (#17643)

Browse source 
Most of the function body was copied from CUDA EP.
This commit is contained in:
Changming Sun 2023-10-06 10:52:17 -07:00 committed by GitHub
parent cb9408e89c
commit 735df7e2a8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 42 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

43
onnxruntime/core/providers/js/js_execution_provider.cc
View file

@ -5,6 +5,7 @@
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
#include "js_execution_provider.h"
@ -13,9 +14,11 @@
#endif
#include "core/graph/function_utils.h"
#include "core/graph/indexed_sub_graph.h"
#include "core/framework/compute_capability.h"
#include "core/framework/data_transfer_manager.h"
#include "core/framework/kernel_registry.h"
#include "core/framework/fallback_cpu_capability.h"
#include "core/providers/shared/node_unit/node_unit.h"
#include "allocator.h"
#include "data_transfer.h"
@ -645,7 +648,45 @@ std::vector<AllocatorPtr> JsExecutionProvider::CreatePreferredAllocators() {
std::vector<std::unique_ptr<ComputeCapability>> JsExecutionProvider::GetCapability(
const onnxruntime::GraphViewer& graph,
const IKernelLookup& kernel_lookup) const {
return IExecutionProvider::GetCapability(graph, kernel_lookup);
InlinedVector<NodeIndex> candidates;
// `tenative_candidates` is a subset of `candidates`.
InlinedVector<NodeIndex> tenative_candidates;
for (auto& node_index : graph.GetNodesInTopologicalOrder()) {
const auto* p_node = graph.GetNode(node_index);
if (p_node == nullptr)
continue;
const auto& node = *p_node;
if (!node.GetExecutionProviderType().empty()) {
// If the node was added by layout transformer, do not move it to CPU
if (node.GetExecutionProviderType() == kJsExecutionProvider) {
candidates.push_back(node.Index());
}
continue;
}
const KernelCreateInfo* webgpu_kernel_def = kernel_lookup.LookUpKernel(node);
// none of the provided registries has a webgpu kernel for this node
if (webgpu_kernel_def == nullptr) {
LOGS(*GetLogger(), INFO) << "webgpu kernel not found in registries for Op type: "
<< node.OpType() << " node name: " << node.Name();
continue;
}
candidates.push_back(node.Index());
tenative_candidates.push_back(node.Index());
}
auto cpu_nodes = GetCpuPreferredNodes(graph, kernel_lookup, tenative_candidates);
std::vector<std::unique_ptr<ComputeCapability>> result;
for (auto& node_index : candidates) {
if (cpu_nodes.count(node_index) > 0) {
continue;
}
auto sub_graph = std::make_unique<IndexedSubGraph>();
sub_graph->nodes.push_back(node_index);
result.emplace_back(std::make_unique<ComputeCapability>(std::move(sub_graph)));
}
return result;
}
std::shared_ptr<KernelRegistry> JsExecutionProvider::GetKernelRegistry() const {