Avoid copying output from GPU to CPU

Author: larryliu0820

backends/cuda/runtime/cuda_backend.cpp

@@ -35,17 +35,24 @@ using executorch::runtime::ArrayRef;
 using executorch::runtime::Backend;
 using executorch::runtime::BackendExecutionContext;
 using executorch::runtime::BackendInitContext;
+using executorch::runtime::BackendOption;
+using executorch::runtime::BackendOptionContext;
 using executorch::runtime::CompileSpec;
 using executorch::runtime::DelegateHandle;
 using executorch::runtime::Error;
 using executorch::runtime::EValue;
 using executorch::runtime::FreeableBuffer;
+using executorch::runtime::kMaxOptionValueLength;
 using executorch::runtime::MemoryAllocator;
 using executorch::runtime::NamedDataMap;
 using executorch::runtime::Result;
 using executorch::runtime::Span;
 using executorch::runtime::etensor::Tensor;
 
+namespace {
+constexpr char kSkipCopyOutputToCpuForMethodOption[] = "skip_copy_output_to_cpu_for_method";
+}
+
 class ET_EXPERIMENTAL CudaBackend final
     : public ::executorch::runtime::BackendInterface {
  private:
@@ -91,6 +98,36 @@ class ET_EXPERIMENTAL CudaBackend final
     return 1;
   }
 
+  Error set_option(
+      ET_UNUSED BackendOptionContext& context,
+      const executorch::runtime::Span<BackendOption>& backend_options)
+      override {
+    for (const auto& option : backend_options) {
+      if (std::strcmp(option.key, kSkipCopyOutputToCpuForMethodOption) == 0) {
+        if (auto* val = std::get_if<bool>(&option.value)) {
+          copy_gpu_outputs_to_cpu_.store(*val, std::memory_order_relaxed);
+        } else {
+          ET_LOG(
+              Error, "Option %s must be a bool.", kSkipCopyOutputToCpuForMethodOption);
+          return Error::InvalidArgument;
+        }
+      }
+    }
+    return Error::Ok;
+  }
+
+  Error get_option(
+      ET_UNUSED BackendOptionContext& context,
+      executorch::runtime::Span<BackendOption>& backend_options) override {
+    for (auto& option : backend_options) {
+      if (std::strcmp(option.key, kCopyGpuOutputsToCpuOption) == 0) {
+        option.value = static_cast<bool>(
+            copy_gpu_outputs_to_cpu_.load(std::memory_order_relaxed));
+      }
+    }
+    return Error::Ok;
+  }
+
   // Once per loaded binary blob
   Result<DelegateHandle*> init(
       BackendInitContext& context,
@@ -303,18 +340,27 @@ class ET_EXPERIMENTAL CudaBackend final
         "AOTInductorModelContainerRun failed with error code %d",
         error);
 
-    // Copy GPU output results back to CPU output tensors
-    for (int i = 0; i < n_outputs; i++) {
-      auto cpu_output_tensor = &(args[i + n_inputs]->toTensor());
-      // For DYNAMIC_BOUND tensors we try to resize
-      ET_CHECK_OK_OR_RETURN_ERROR(
-          resize_tensor(*cpu_output_tensor, gpu_outputs[i]->sizes()),
-          "Error resizing tensor at output index %d",
-          i);
-      ET_CHECK_OK_OR_RETURN_ERROR(
-          aoti_torch_copy_(cpu_output_tensor, gpu_outputs[i], 0),
-          "Failed to copy GPU output %d back to CPU",
-          i);
+    const bool copy_outputs =
+        copy_gpu_outputs_to_cpu_.load(std::memory_order_relaxed);
+
+    if (copy_outputs) {
+      // Copy GPU output results back to CPU output tensors
+      for (int i = 0; i < n_outputs; i++) {
+        auto cpu_output_tensor = &(args[i + n_inputs]->toTensor());
+        // For DYNAMIC_BOUND tensors we try to resize
+        ET_CHECK_OK_OR_RETURN_ERROR(
+            resize_tensor(*cpu_output_tensor, gpu_outputs[i]->sizes()),
+            "Error resizing tensor at output index %d",
+            i);
+        ET_CHECK_OK_OR_RETURN_ERROR(
+            aoti_torch_copy_(cpu_output_tensor, gpu_outputs[i], 0),
+            "Failed to copy GPU output %d back to CPU",
+            i);
+      }
+    } else {
+      for (int i = 0; i < n_outputs; i++) {
+        args[i + n_inputs]->toTensor() = *gpu_outputs[i];
+      }
     }
 
     return Error::Ok;
@@ -365,6 +411,9 @@ class ET_EXPERIMENTAL CudaBackend final
     delete handle;
     clear_all_tensors();
   }
+
+ private:
+  std::atomic<bool> copy_gpu_outputs_to_cpu_{true};
 };
 
 } // namespace executorch::backends::cuda

extension/asr/runner/runner.cpp

@@ -107,7 +107,21 @@ Error AsrRunner::load() {
 
   ET_CHECK_OK_OR_RETURN_ERROR(module_->load_method(kDecoderMethodName));
   decoder_method_loaded_ = true;
-
+#ifdef CUDA_AVAILABLE
+  executorch::runtime::BackendOptions<1> backend_options;
+  // For decoder still copy output from GPU to CPU for sampling.
+  // TODO: change this to use a CUDA kernel to sample and then skip copying decoder output
+  ET_CHECK_OK_OR_RETURN_ERROR(
+      backend_options.set_option("skip_copy_output_to_cpu_for_method", kEncoderMethodName));
+  const auto opt_err =
+      executorch::runtime::set_option("CudaBackend", backend_options.view());
+  if (opt_err != ::executorch::runtime::Error::Ok) {
+    ET_LOG(
+        Warning,
+        "Failed to set CUDA backend options: %d",
+        static_cast<int>(opt_err));
+  }
+#endif
   ET_CHECK_OK_OR_RETURN_ERROR(load_tokenizer());
   auto eos_ids = get_eos_ids(tokenizer_.get(), module_.get());
   if (!eos_ids.empty()) {