Enable GPU->CPU transfers

dali/c_api/c_api.cc

@@ -246,13 +246,33 @@ daliCreatePipeline2(daliPipelineHandle *pipe_handle, const char *serialized_pipe
                     int async_execution, int separated_execution, int prefetch_queue_depth,
                     int cpu_prefetch_queue_depth, int gpu_prefetch_queue_depth,
                     int enable_memory_stats) {
-  bool se = separated_execution != 0;
-  bool pe = pipelined_execution != 0;
-  bool ae = async_execution != 0;
+  dali_exec_flags_t flags = {};
+  if (async_execution)
+    flags = flags | DALI_EXEC_IS_ASYNC;
+  if (pipelined_execution)
+    flags = flags | DALI_EXEC_IS_PIPELINED;
+  if (separated_execution)
+    flags = flags | DALI_EXEC_IS_SEPARATED;
+  daliCreatePipeline3(pipe_handle, serialized_pipeline, length,
+                      max_batch_size, num_threads, device_id, flags,
+                      prefetch_queue_depth, cpu_prefetch_queue_depth, gpu_prefetch_queue_depth,
+                      enable_memory_stats);
+}
 
-  auto pipeline =
-          std::make_unique<dali::Pipeline>(std::string(serialized_pipeline, length), max_batch_size,
-                                           num_threads, device_id, pe, prefetch_queue_depth, ae);
+DLL_PUBLIC void
+daliCreatePipeline3(daliPipelineHandle *pipe_handle, const char *serialized_pipeline, int length,
+                    int max_batch_size, int num_threads, int device_id,
+                    dali_exec_flags_t exec_flags, int prefetch_queue_depth,
+                    int cpu_prefetch_queue_depth, int gpu_prefetch_queue_depth,
+                    int enable_memory_stats) {
+  bool se = exec_flags & DALI_EXEC_IS_SEPARATED;
+  bool pe = exec_flags & DALI_EXEC_IS_PIPELINED;
+  bool ae = exec_flags & DALI_EXEC_IS_ASYNC;
+  bool de = exec_flags & DALI_EXEC_IS_DYNAMIC;
+
+  auto pipeline = std::make_unique<dali::Pipeline>(
+    std::string(serialized_pipeline, length), max_batch_size, num_threads, device_id,
+    pe, prefetch_queue_depth, ae, de);
   pipeline->SetExecutionTypes(pe, se, ae);
   if (se) {
     pipeline->SetQueueSizes(cpu_prefetch_queue_depth, gpu_prefetch_queue_depth);
@@ -263,7 +283,6 @@ daliCreatePipeline2(daliPipelineHandle *pipe_handle, const char *serialized_pipe
   *pipe_handle = WrapPipeline(std::move(pipeline)).release();
 }
 
-
 void daliDeserializeDefault(daliPipelineHandle *pipe_handle, const char *serialized_pipeline,
                             int length) {
   auto pipeline = std::make_unique<dali::Pipeline>(std::string(serialized_pipeline, length));

dali/pipeline/executor/executor_factory.cc

@@ -58,40 +58,44 @@ bool ForceExec2() {
 }  // namespace
 
 template <typename... T>
-std::unique_ptr<ExecutorBase> GetExecutorImpl(bool pipelined, bool separated, bool async,
-                                              T&&... args) {
-  if (async && separated && pipelined) {
+std::unique_ptr<ExecutorBase> GetExecutorImpl(
+      bool pipelined, bool separated, bool async, bool dynamic, T&&... args) {
+  // Go over possible combinations and throw otherwise.
+  if (async && separated && pipelined && !dynamic) {
     return std::make_unique<AsyncSeparatedPipelinedExecutor>(std::forward<T>(args)...);
   } else if (async && !separated && pipelined) {
-    if (ForceExec2()) {
-      std::cerr << "\n!!! Forced use of Executor 2.0 !!!" << std::endl;
+    bool force_exec2 = ForceExec2();
+    if (dynamic || force_exec2) {
+      if (force_exec2)
+        std::cerr << "\n!!! Forced use of Executor 2.0 !!!" << std::endl;
       return std::make_unique<exec2::Executor2>(MakeExec2Config(std::forward<T>(args)...));
     } else {
      return std::make_unique<AsyncPipelinedExecutor>(std::forward<T>(args)...);
     }
-  } else if (!async && separated && pipelined) {
+  } else if (!async && separated && pipelined && !dynamic) {
     return std::make_unique<SeparatedPipelinedExecutor>(std::forward<T>(args)...);
-  } else if (!async && !separated && pipelined) {
+  } else if (!async && !separated && pipelined && !dynamic) {
     return std::make_unique<PipelinedExecutor>(std::forward<T>(args)...);
-  } else if (!async && !separated && !pipelined) {
+  } else if (!async && !separated && !pipelined && !dynamic) {
     return std::make_unique<SimpleExecutor>(std::forward<T>(args)...);
   }
   std::stringstream error;
   error << std::boolalpha;
   error << "No supported executor selected for pipelined = " << pipelined
-        << ", separated = " << separated << ", async = " << async << std::endl;
+        << ", separated = " << separated << ", async = " << async
+        << ", dynamic = " << dynamic << std::endl;
   DALI_FAIL(error.str());
 }
 
 
-std::unique_ptr<ExecutorBase> GetExecutor(bool pipelined, bool separated, bool async,
+std::unique_ptr<ExecutorBase> GetExecutor(bool pipelined, bool separated, bool async, bool dynamic,
                                           int batch_size, int num_thread, int device_id,
                                           size_t bytes_per_sample_hint, bool set_affinity,
                                           int max_num_stream,
                                           int default_cuda_stream_priority,
                                           QueueSizes prefetch_queue_depth) {
   return GetExecutorImpl(
-    pipelined, separated, async,
+    pipelined, separated, async, dynamic,
     batch_size, num_thread, device_id, bytes_per_sample_hint, set_affinity,
     max_num_stream, default_cuda_stream_priority, prefetch_queue_depth);
 }

dali/pipeline/executor/executor_factory.h

@@ -23,7 +23,7 @@
 namespace dali {
 
 DLL_PUBLIC
-std::unique_ptr<ExecutorBase> GetExecutor(bool pipelined, bool separated, bool async,
+std::unique_ptr<ExecutorBase> GetExecutor(bool pipelined, bool separated, bool async, bool dynamic,
                                           int batch_size, int num_thread, int device_id,
                                           size_t bytes_per_sample_hint, bool set_affinity = false,
                                           int max_num_stream = -1,

dali/pipeline/executor/lowered_graph.cc

@@ -97,8 +97,18 @@ void OpGraph::Lower(const graph::OpGraph &definition) {
   if (!op_nodes_.empty() || !tensor_nodes_.empty())
     throw std::logic_error("The target graph must be empty");
   for (auto &node : definition.OpNodes()) {
-    auto &lowered_op = AddOp(node.spec, node.instance_name);
-    lowered_op.definition = &node;
+    try {
+      auto &lowered_op = AddOp(node.spec, node.instance_name);
+      lowered_op.definition = &node;
+    } catch (...) {
+      PropagateError({
+          std::current_exception(),
+          make_string(
+              "Critical error when building pipeline:\n",
+              GetErrorContextMessage(node.spec)),
+          "\nCurrent pipeline object is no longer valid."
+        });
+    }
   }
   for (auto &t : tensor_nodes_) {
     t.definition = definition.GetData(t.name);
@@ -131,22 +141,27 @@ OpNode &OpGraph::AddOp(const OpSpec &spec, const std::string &op_name) {
   // Validate the op specification
   CheckOpConstraints(spec);
 
+  const char *gpu2cpu_error =
+    "This pipeline doesn't support transition from GPU to CPU.\n"
+    "To enable GPU->CPU transitions, use the experimental \"dynamic\" executor.\n"
+    "Specify experimental_exec_dynamic=True in your Pipeline constructor or @pipeline_def.";
+
   string device = spec.GetArgument<string>("device");
   auto op_type = ParseOpType(device);
   // TODO(klecki): refactor this out
   switch (op_type) {
     case OpType::CPU: {
       // Enforce graph constraints
-      DALI_ENFORCE(AllInputsCPU(spec), "CPU ops cannot receive GPU input data.");
-      DALI_ENFORCE(AllOutputsCPU(spec), "CPU ops can only produce CPU output data.");
+      DALI_ENFORCE(AllInputsCPU(spec), gpu2cpu_error);
+      DALI_ENFORCE(AllOutputsCPU(spec), "CPU operators can only produce CPU output data.");
       break;
     }
     case OpType::GPU: {
       break;
     }
     case OpType::MIXED: {
       // Enforce graph constraints
-      DALI_ENFORCE(AllInputsCPU(spec), "Mixed ops cannot receive GPU input data.");
+      DALI_ENFORCE(AllInputsCPU(spec), gpu2cpu_error);
       break;
     }
     default:

dali/pipeline/operator/builtin/copy.cc

@@ -17,11 +17,47 @@
 namespace dali {
 
 template <>
-void Copy<CPUBackend>::RunCopies(Workspace &ws) {
-  scatter_gather_.Run(ws.GetThreadPool(), true);
+void Copy<CPUBackend>::RunImpl(Workspace &ws) {
+  if (ws.InputIsType<CPUBackend>(0)) {
+    auto &input = ws.Input<CPUBackend>(0);
+    auto &output = ws.Output<CPUBackend>(0);
+
+    int batch_size = input.num_samples();
+    output.SetLayout(input.GetLayout());
+    auto shapes = input.shape();
+
+    auto &thread_pool = ws.GetThreadPool();
+    for (int sample_id = 0; sample_id < batch_size; ++sample_id) {
+      thread_pool.AddWork(
+          [sample_id, &input, &output](int tid) {
+            output.CopySample(sample_id, input, sample_id, AccessOrder::host());
+          },
+          shapes.tensor_size(sample_id));
+    }
+    thread_pool.RunAll();
+  } else {
+    auto &input = ws.Input<GPUBackend>(0);
+    auto &output = ws.Output<CPUBackend>(0);
+    output.Copy(input, ws.output_order());
+  }
+}
+
+template <>
+void Copy<GPUBackend>::RunImpl(Workspace &ws) {
+  if (ws.InputIsType<CPUBackend>(0)) {
+    auto &input = ws.Input<CPUBackend>(0);
+    auto &output = ws.Output<GPUBackend>(0);
+    output.Copy(input, ws.output_order());
+  } else {
+    auto &input = ws.Input<GPUBackend>(0);
+    auto &output = ws.Output<GPUBackend>(0);
+    output.Copy(input, ws.output_order());
+  }
 }
 
 DALI_REGISTER_OPERATOR(Copy, Copy<CPUBackend>, CPU);
+DALI_REGISTER_OPERATOR(Copy, Copy<GPUBackend>, GPU);
+
 
 DALI_SCHEMA(Copy)
   .DocStr("Creates a copy of the input tensor.")

dali/pipeline/operator/builtin/copy.h

@@ -28,10 +28,8 @@ namespace dali {
 template <typename Backend>
 class Copy : public StatelessOperator<Backend> {
  public:
-  inline explicit Copy(const OpSpec &spec) :
-    StatelessOperator<Backend>(spec), scatter_gather_(kMaxSizePerBlock) {}
-
-  inline ~Copy() override = default;
+  explicit Copy(const OpSpec &spec) :
+    StatelessOperator<Backend>(spec) {}
 
   DISABLE_COPY_MOVE_ASSIGN(Copy);
 
@@ -42,37 +40,15 @@ class Copy : public StatelessOperator<Backend> {
 
   bool SetupImpl(std::vector<OutputDesc> &output_desc, const Workspace &ws) override {
     output_desc.resize(1);
-    const auto &input = ws.Input<Backend>(0);
-    output_desc[0].type = input.type();
-    output_desc[0].shape = input.shape();
+    output_desc[0].type = ws.GetInputDataType(0);
+    output_desc[0].shape = ws.GetInputShape(0);
     return true;
   }
 
-  void RunImpl(Workspace &ws) override {
-    auto &input = ws.Input<Backend>(0);
-    auto data_type_size = input.type_info().size();
-    auto &output = ws.Output<Backend>(0);
-    output.SetLayout(input.GetLayout());
-    for (int i = 0; i < input.num_samples(); i++) {
-      auto tensor_shape = input.tensor_shape(i);
-      auto tensor_size = volume(tensor_shape);
-      scatter_gather_.AddCopy(output.raw_mutable_tensor(i), input.raw_tensor(i),
-                              tensor_size * data_type_size);
-    }
-    RunCopies(ws);
-  }
-
-  void RunCopies(Workspace &ws);
-
-  std::conditional_t<
-      std::is_same<Backend, CPUBackend>::value,
-      kernels::ScatterGatherCPU,
-      kernels::ScatterGatherGPU> scatter_gather_;
-  // 256 kB per block for GPU
-  static constexpr size_t kMaxSizePerBlock =
-      std::is_same<Backend, CPUBackend>::value ? kernels::ScatterGatherCPU::kAnyBlockSize : 1 << 18;
+  void RunImpl(Workspace &ws) override;
 };
 
+
 }  // namespace dali
 
 #endif  // DALI_PIPELINE_OPERATOR_BUILTIN_COPY_H_

dali/pipeline/operator/checkpointing/checkpoint_test.cc

@@ -36,7 +36,7 @@ void BuildFromLegacyGraph(Checkpoint &checkpoint, const OpGraph &graph) {
 }
 
 auto GetSimpleExecutor() {
-  return GetExecutor(false, false, false, 1, 1, CPU_ONLY_DEVICE_ID, 0);
+  return GetExecutor(false, false, false, false, 1, 1, CPU_ONLY_DEVICE_ID, 0);
 }
 
 }  // namespace