From 74d9b580f9ab713fe3cd3210ef71635798d5505e Mon Sep 17 00:00:00 2001 From: ptredak Date: Tue, 11 Sep 2018 14:04:32 -0700 Subject: [PATCH] Use pinned memory only when interfacing with MakeContiguous Signed-off-by: ptredak --- dali/pipeline/executor/executor.cc | 9 +++++++-- dali/pipeline/executor/pipelined_executor.cc | 9 +++++++++ dali/pipeline/executor/pipelined_executor.h | 1 + 3 files changed, 17 insertions(+), 2 deletions(-) diff --git a/dali/pipeline/executor/executor.cc b/dali/pipeline/executor/executor.cc index 8053808d44..d6e4078958 100644 --- a/dali/pipeline/executor/executor.cc +++ b/dali/pipeline/executor/executor.cc @@ -458,8 +458,13 @@ void Executor::SetupDataForGraph(WorkspaceBlob *wsb) { HostWorkspace &src_ws = wsb->cpu_op_data[parent_idx]; auto input = src_ws.SharedCPUOutput(input_src_idx); - for (auto t : input) { - t->set_pinned(true); + // Use pinned memory only when it is useful + if (node.spec.name() == "MakeContiguous" && + node.spec.NumOutput() == 1 && + node.spec.OutputDevice(0) == "gpu") { + for (auto t : input) { + t->set_pinned(true); + } } ws.AddInput(input); } diff --git a/dali/pipeline/executor/pipelined_executor.cc b/dali/pipeline/executor/pipelined_executor.cc index 8a0dfb908f..d4c27e9eaa 100644 --- a/dali/pipeline/executor/pipelined_executor.cc +++ b/dali/pipeline/executor/pipelined_executor.cc @@ -223,6 +223,15 @@ void PipelinedExecutor::SetStageOutputsForIter( int input_idx = info.con_and_idx[j].second; wsb->mixed_op_data[mixed_op_id].SetInput( input_idx, tvp.Get(queue_idx)); + const OpNode &node = graph_->mixed_node(mixed_op_id); + // Use pinned memory only when it is useful + if (node.spec.name() == "MakeContiguous" && + node.spec.NumOutput() == 1 && + node.spec.OutputDevice(0) == "gpu") { + for (auto& v : tvp.Get(queue_idx)) { + v->set_pinned(true); + } + } } else if (graph_->NodeType(node_id) == DALI_CPU) { int cpu_op_id = graph_->NodeIdx(node_id); int input_idx = info.con_and_idx[j].second; diff --git a/dali/pipeline/executor/pipelined_executor.h b/dali/pipeline/executor/pipelined_executor.h index 8bd22229a9..524b3b2f56 100644 --- a/dali/pipeline/executor/pipelined_executor.h +++ b/dali/pipeline/executor/pipelined_executor.h @@ -65,6 +65,7 @@ class DLL_PUBLIC PipelinedExecutor : public Executor { for (int j = 0; j < batch_size; ++j) { tvs_[i].push_back(std::make_shared>()); tvs_[i].back()->Resize({(Index)bytes_hint}); + tvs_[i].back()->set_pinned(false); } } }