nGraph: Avoid input and output data copies (#940)

This commit is contained in:
nivas-x86 2019-04-30 12:10:28 -07:00 committed by jywu-msft
parent 01cd7eaca8
commit 3b0dda0aca
3 changed files with 4 additions and 4 deletions

View file

@ -176,7 +176,7 @@ Status GraphPartitioner::Partition(Graph& graph, bool export_dll, FuncManager& f
//prepare the func kernel
KernelDefBuilder builder;
BuildFusedKernelDef(builder, *node);
if (node->GetExecutionProviderType() == onnxruntime::kTensorrtExecutionProvider) {
if (node->GetExecutionProviderType() == onnxruntime::kTensorrtExecutionProvider || node->GetExecutionProviderType() == onnxruntime::kNGraphExecutionProvider) {
builder.SetDefaultInputsMemoryType(OrtMemTypeCPUInput);
builder.SetDefaultOutputMemoryType(OrtMemTypeCPUOutput);
}

View file

@ -128,8 +128,8 @@ common::Status CopyOneInputAcrossDevices(const SessionState& session_state,
ORT_ENFORCE(p_input_provider);
}
//no copy for TRT
if (required_provider_type == onnxruntime::kTensorrtExecutionProvider) {
//no copy for TRT and nGraph
if (required_provider_type == onnxruntime::kTensorrtExecutionProvider || required_provider_type == onnxruntime::kNGraphExecutionProvider) {
new_mlvalue = orig_mlvalue;
break;
}

View file

@ -184,7 +184,7 @@ void TransformerMemcpyImpl::ProcessDefs(onnxruntime::Node& node, const KernelReg
} else {
// TODO: copy between devices? i.e. multiple GPUs
if (node.GetExecutionProviderType() != onnxruntime::kCpuExecutionProvider && node.GetExecutionProviderType() != onnxruntime::kTensorrtExecutionProvider &&
!node.GetExecutionProviderType().empty()) {
node.GetExecutionProviderType() != onnxruntime::kNGraphExecutionProvider && !node.GetExecutionProviderType().empty()) {
ORT_THROW("Execution type '", node.GetExecutionProviderType(), "' doesn't support memcpy ");
}