Use pinned memory only when interfacing with MakeContiguous

Signed-off-by: ptredak <[email protected]>

Author: ptrendx

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);
     }

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;

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<Tensor<Backend>>());
           tvs_[i].back()->Resize({(Index)bytes_hint});
+          tvs_[i].back()->set_pinned(false);
         }
       }
     }