No public description

PiperOrigin-RevId: 572309070
google-ai-edge · Oct 10, 2023 · df13788 · df13788
1 parent 0d5f35d
commit df13788
Show file tree

Hide file tree

Showing 9 changed files with 83 additions and 43 deletions.
diff --git a/mediapipe/calculators/image/BUILD b/mediapipe/calculators/image/BUILD
@@ -786,6 +786,7 @@ cc_library(
             ":affine_transformation_runner_gl",
             "//mediapipe/gpu:gl_calculator_helper",
             "//mediapipe/gpu:gpu_buffer",
+            "//mediapipe/gpu:gpu_service",
         ],
     }) + select({
         "//mediapipe/framework/port:disable_opencv": [],

diff --git a/mediapipe/calculators/image/image_clone_calculator.cc b/mediapipe/calculators/image/image_clone_calculator.cc
@@ -64,17 +64,15 @@ class ImageCloneCalculator : public Node {
           "GPU processing is disabled in build flags");
     }
 #else
-    MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
+    MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(
+        cc, /*requesst_gpu_as_optional=*/true));
 #endif  // MEDIAPIPE_DISABLE_GPU
     return absl::OkStatus();
   }
 
   absl::Status Open(CalculatorContext* cc) override {
     const auto& options = cc->Options<mediapipe::ImageCloneCalculatorOptions>();
     output_on_gpu_ = options.output_on_gpu();
-#if !MEDIAPIPE_DISABLE_GPU
-    MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
-#endif  // !MEDIAPIPE_DISABLE_GPU
     return absl::OkStatus();
   }
 
@@ -104,6 +102,10 @@ class ImageCloneCalculator : public Node {
 
     if (output_on_gpu_ && !input_on_gpu) {
 #if !MEDIAPIPE_DISABLE_GPU
+      if (!gpu_initialized_) {
+        MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
+        gpu_initialized_ = true;
+      }
       gpu_helper_.RunInGlContext([&output]() { output->ConvertToGpu(); });
 #endif  // !MEDIAPIPE_DISABLE_GPU
     } else if (!output_on_gpu_ && input_on_gpu) {
@@ -118,6 +120,7 @@ class ImageCloneCalculator : public Node {
   bool output_on_gpu_;
 #if !MEDIAPIPE_DISABLE_GPU
   mediapipe::GlCalculatorHelper gpu_helper_;
+  bool gpu_initialized_ = false;
 #endif  // !MEDIAPIPE_DISABLE_GPU
 };
 MEDIAPIPE_REGISTER_NODE(ImageCloneCalculator);

diff --git a/mediapipe/calculators/image/segmentation_smoothing_calculator.cc b/mediapipe/calculators/image/segmentation_smoothing_calculator.cc
@@ -117,7 +117,8 @@ absl::Status SegmentationSmoothingCalculator::GetContract(
   cc->Outputs().Tag(kOutputMaskTag).Set<Image>();
 
 #if !MEDIAPIPE_DISABLE_GPU
-  MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
+  MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(
+      cc, /*requesst_gpu_as_optional=*/true));
 #endif  // !MEDIAPIPE_DISABLE_GPU
 
   return absl::OkStatus();
@@ -130,10 +131,6 @@ absl::Status SegmentationSmoothingCalculator::Open(CalculatorContext* cc) {
       cc->Options<mediapipe::SegmentationSmoothingCalculatorOptions>();
   combine_with_previous_ratio_ = options.combine_with_previous_ratio();
 
-#if !MEDIAPIPE_DISABLE_GPU
-  MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
-#endif  //  !MEDIAPIPE_DISABLE_GPU
-
   return absl::OkStatus();
 }
 
@@ -154,6 +151,9 @@ absl::Status SegmentationSmoothingCalculator::Process(CalculatorContext* cc) {
 
   if (use_gpu) {
 #if !MEDIAPIPE_DISABLE_GPU
+    if (!gpu_initialized_) {
+      MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
+    }
     MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this, cc]() -> absl::Status {
       if (!gpu_initialized_) {
         MP_RETURN_IF_ERROR(GlSetup(cc));
@@ -178,10 +178,12 @@ absl::Status SegmentationSmoothingCalculator::Process(CalculatorContext* cc) {
 
 absl::Status SegmentationSmoothingCalculator::Close(CalculatorContext* cc) {
 #if !MEDIAPIPE_DISABLE_GPU
-  gpu_helper_.RunInGlContext([this] {
-    if (program_) glDeleteProgram(program_);
-    program_ = 0;
-  });
+  if (gpu_initialized_) {
+    gpu_helper_.RunInGlContext([this] {
+      if (program_) glDeleteProgram(program_);
+      program_ = 0;
+    });
+  }
 #endif  // !MEDIAPIPE_DISABLE_GPU
 
   return absl::OkStatus();

diff --git a/mediapipe/calculators/image/warp_affine_calculator.cc b/mediapipe/calculators/image/warp_affine_calculator.cc
@@ -36,6 +36,7 @@
 #if !MEDIAPIPE_DISABLE_GPU
 #include "mediapipe/gpu/gl_calculator_helper.h"
 #include "mediapipe/gpu/gpu_buffer.h"
+#include "mediapipe/gpu/gpu_service.h"
 #endif  // !MEDIAPIPE_DISABLE_GPU
 
 namespace mediapipe {
@@ -106,6 +107,7 @@ class WarpAffineRunnerHolder<mediapipe::GpuBuffer> {
         cc->Options<mediapipe::WarpAffineCalculatorOptions>().interpolation());
     return gl_helper_->Open(cc);
   }
+
   absl::StatusOr<RunnerType*> GetRunner() {
     if (!runner_) {
       MP_ASSIGN_OR_RETURN(
@@ -142,7 +144,10 @@ class WarpAffineRunnerHolder<mediapipe::Image> {
       MP_RETURN_IF_ERROR(cpu_holder_.Open(cc));
 #endif  // !MEDIAPIPE_DISABLE_OPENCV
 #if !MEDIAPIPE_DISABLE_GPU
-      MP_RETURN_IF_ERROR(gpu_holder_.Open(cc));
+      if (cc->Service(kGpuService).IsAvailable()) {
+        MP_RETURN_IF_ERROR(gpu_holder_.Open(cc));
+        gpu_holder_initialized_ = true;
+      }
 #endif  // !MEDIAPIPE_DISABLE_GPU
       return absl::OkStatus();
     }
@@ -152,6 +157,9 @@ class WarpAffineRunnerHolder<mediapipe::Image> {
         AffineTransformation::BorderMode border_mode) override {
       if (input.UsesGpu()) {
 #if !MEDIAPIPE_DISABLE_GPU
+        if (!gpu_holder_initialized_) {
+          return absl::UnavailableError("GPU support is not available");
+        }
         MP_ASSIGN_OR_RETURN(auto* runner, gpu_holder_.GetRunner());
         MP_ASSIGN_OR_RETURN(
             auto result,
@@ -183,6 +191,7 @@ class WarpAffineRunnerHolder<mediapipe::Image> {
 #endif  // !MEDIAPIPE_DISABLE_OPENCV
 #if !MEDIAPIPE_DISABLE_GPU
     WarpAffineRunnerHolder<mediapipe::GpuBuffer> gpu_holder_;
+    bool gpu_holder_initialized_ = false;
 #endif  // !MEDIAPIPE_DISABLE_GPU
   };
 
@@ -196,20 +205,24 @@ class WarpAffineCalculatorImpl : public mediapipe::api2::NodeImpl<InterfaceT> {
   static absl::Status UpdateContract(CalculatorContract* cc) {
     if constexpr (std::is_same_v<InterfaceT, WarpAffineCalculatorGpu> ||
                   std::is_same_v<InterfaceT, WarpAffineCalculator>) {
-      MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
+      MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(
+          cc, /*requesst_gpu_as_optional=*/true));
     }
     return absl::OkStatus();
   }
 #endif  // !MEDIAPIPE_DISABLE_GPU
-
-  absl::Status Open(CalculatorContext* cc) override { return holder_.Open(cc); }
-
   absl::Status Process(CalculatorContext* cc) override {
     if (InterfaceT::kInImage(cc).IsEmpty() ||
         InterfaceT::kMatrix(cc).IsEmpty() ||
         InterfaceT::kOutputSize(cc).IsEmpty()) {
       return absl::OkStatus();
     }
+
+    if (!holder_initialized_) {
+      MP_RETURN_IF_ERROR(holder_.Open(cc));
+      holder_initialized_ = true;
+    }
+
     const std::array<float, 16>& transform = *InterfaceT::kMatrix(cc);
     auto [out_width, out_height] = *InterfaceT::kOutputSize(cc);
     AffineTransformation::Size output_size;
@@ -230,6 +243,7 @@ class WarpAffineCalculatorImpl : public mediapipe::api2::NodeImpl<InterfaceT> {
  private:
   WarpAffineRunnerHolder<typename decltype(InterfaceT::kInImage)::PayloadT>
       holder_;
+  bool holder_initialized_ = false;
 };
 
 }  // namespace

diff --git a/mediapipe/calculators/tensor/tensors_to_detections_calculator.cc b/mediapipe/calculators/tensor/tensors_to_detections_calculator.cc
@@ -266,7 +266,8 @@ absl::Status TensorsToDetectionsCalculator::UpdateContract(
     CalculatorContract* cc) {
   if (CanUseGpu()) {
 #ifndef MEDIAPIPE_DISABLE_GL_COMPUTE
-    MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
+    MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(
+        cc, /*requesst_gpu_as_optional=*/true));
 #elif MEDIAPIPE_METAL_ENABLED
     MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
 #endif  // !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
@@ -280,7 +281,6 @@ absl::Status TensorsToDetectionsCalculator::Open(CalculatorContext* cc) {
 
   if (CanUseGpu()) {
 #ifndef MEDIAPIPE_DISABLE_GL_COMPUTE
-    MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
 #elif MEDIAPIPE_METAL_ENABLED
     gpu_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
     RET_CHECK(gpu_helper_);
@@ -676,13 +676,15 @@ absl::Status TensorsToDetectionsCalculator::ProcessGPU(
 
 absl::Status TensorsToDetectionsCalculator::Close(CalculatorContext* cc) {
 #ifndef MEDIAPIPE_DISABLE_GL_COMPUTE
-  gpu_helper_.RunInGlContext([this] {
-    decoded_boxes_buffer_ = nullptr;
-    scored_boxes_buffer_ = nullptr;
-    raw_anchors_buffer_ = nullptr;
-    glDeleteProgram(decode_program_);
-    glDeleteProgram(score_program_);
-  });
+  if (gpu_inited_) {
+    gpu_helper_.RunInGlContext([this] {
+      decoded_boxes_buffer_ = nullptr;
+      scored_boxes_buffer_ = nullptr;
+      raw_anchors_buffer_ = nullptr;
+      glDeleteProgram(decode_program_);
+      glDeleteProgram(score_program_);
+    });
+  }
 #elif MEDIAPIPE_METAL_ENABLED
   decoded_boxes_buffer_ = nullptr;
   scored_boxes_buffer_ = nullptr;
@@ -942,6 +944,7 @@ absl::Status TensorsToDetectionsCalculator::GpuInit(CalculatorContext* cc) {
       break;
   }
 #ifndef MEDIAPIPE_DISABLE_GL_COMPUTE
+  MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
   MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this, output_format_flag]()
                                                     -> absl::Status {
     // A shader to decode detection boxes.
@@ -1420,7 +1423,6 @@ kernel void scoreKernel(
           num_classes_, max_wg_size));
     }
   }
-
 #endif  // !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
 
   return absl::OkStatus();

diff --git a/mediapipe/calculators/tensor/tensors_to_segmentation_calculator.cc b/mediapipe/calculators/tensor/tensors_to_segmentation_calculator.cc
@@ -173,6 +173,7 @@ class TensorsToSegmentationCalculator : public CalculatorBase {
 #if !MEDIAPIPE_DISABLE_GPU
   mediapipe::GlCalculatorHelper gpu_helper_;
   GLuint upsample_program_;
+  bool gpu_initialized_ = false;
 #if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
   int cached_width_ = 0;
   int cached_height_ = 0;
@@ -206,7 +207,8 @@ absl::Status TensorsToSegmentationCalculator::GetContract(
 
   if (CanUseGpu()) {
 #if !MEDIAPIPE_DISABLE_GPU
-    MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
+    MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(
+        cc, /*requesst_gpu_as_optional=*/true));
 #if MEDIAPIPE_METAL_ENABLED
     MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
 #endif  // MEDIAPIPE_METAL_ENABLED
@@ -218,12 +220,9 @@ absl::Status TensorsToSegmentationCalculator::GetContract(
 
 absl::Status TensorsToSegmentationCalculator::Open(CalculatorContext* cc) {
   cc->SetOffset(TimestampDiff(0));
-  bool use_gpu = false;
 
   if (CanUseGpu()) {
 #if !MEDIAPIPE_DISABLE_GPU
-    use_gpu = true;
-    MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
 #if MEDIAPIPE_METAL_ENABLED
     metal_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
     RET_CHECK(metal_helper_);
@@ -233,14 +232,6 @@ absl::Status TensorsToSegmentationCalculator::Open(CalculatorContext* cc) {
 
   MP_RETURN_IF_ERROR(LoadOptions(cc));
 
-  if (use_gpu) {
-#if !MEDIAPIPE_DISABLE_GPU
-    MP_RETURN_IF_ERROR(InitGpu(cc));
-#else
-    RET_CHECK_FAIL() << "GPU processing disabled.";
-#endif  // !MEDIAPIPE_DISABLE_GPU
-  }
-
   return absl::OkStatus();
 }
 
@@ -285,6 +276,15 @@ absl::Status TensorsToSegmentationCalculator::Process(CalculatorContext* cc) {
   }
 
   if (use_gpu) {
+#if !MEDIAPIPE_DISABLE_GPU
+    if (!gpu_initialized_) {
+      MP_RETURN_IF_ERROR(InitGpu(cc));
+      gpu_initialized_ = true;
+    }
+#else
+    RET_CHECK_FAIL() << "GPU processing disabled.";
+#endif  // !MEDIAPIPE_DISABLE_GPU
+
 #if !MEDIAPIPE_DISABLE_GPU
     MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this, cc]() -> absl::Status {
       MP_RETURN_IF_ERROR(ProcessGpu(cc));
@@ -306,6 +306,10 @@ absl::Status TensorsToSegmentationCalculator::Process(CalculatorContext* cc) {
 
 absl::Status TensorsToSegmentationCalculator::Close(CalculatorContext* cc) {
 #if !MEDIAPIPE_DISABLE_GPU
+  if (!gpu_initialized_) {
+    return absl::OkStatus();
+  }
+
   gpu_helper_.RunInGlContext([this] {
     if (upsample_program_) glDeleteProgram(upsample_program_);
     upsample_program_ = 0;
@@ -635,6 +639,7 @@ absl::Status TensorsToSegmentationCalculator::LoadOptions(
 
 absl::Status TensorsToSegmentationCalculator::InitGpu(CalculatorContext* cc) {
 #if !MEDIAPIPE_DISABLE_GPU
+  MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
   MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]() -> absl::Status {
   // A shader to process a segmentation tensor into an output mask.
   // Currently uses 4 channels for output, and sets R+A channels as mask value.
@@ -899,6 +904,8 @@ void main() {
 
     return absl::OkStatus();
   }));
+
+  gpu_initialized_ = true;
 #endif  // !MEDIAPIPE_DISABLE_GPU
 
   return absl::OkStatus();

diff --git a/mediapipe/gpu/gl_calculator_helper.cc b/mediapipe/gpu/gl_calculator_helper.cc
@@ -56,8 +56,13 @@ void GlCalculatorHelper::InitializeForTest(GpuResources* gpu_resources) {
 }
 
 // static
-absl::Status GlCalculatorHelper::UpdateContract(CalculatorContract* cc) {
-  cc->UseService(kGpuService);
+absl::Status GlCalculatorHelper::UpdateContract(CalculatorContract* cc,
+                                                bool requesst_gpu_as_optional) {
+  if (requesst_gpu_as_optional) {
+    cc->UseService(kGpuService).Optional();
+  } else {
+    cc->UseService(kGpuService);
+  }
   // Allow the legacy side packet to be provided, too, for backwards
   // compatibility with existing graphs. It will just be ignored.
   auto& input_side_packets = cc->InputSidePackets();

diff --git a/mediapipe/gpu/gl_calculator_helper.h b/mediapipe/gpu/gl_calculator_helper.h
@@ -67,7 +67,8 @@ class GlCalculatorHelper {
 
   // This method can be called from GetContract to set up the needed GPU
   // resources.
-  static absl::Status UpdateContract(CalculatorContract* cc);
+  static absl::Status UpdateContract(CalculatorContract* cc,
+                                     bool requesst_gpu_as_optional = false);
 
   // This method can be called from FillExpectations to set the correct types
   // for the shared GL input side packet(s).

diff --git a/mediapipe/python/pybind/calculator_graph.cc b/mediapipe/python/pybind/calculator_graph.cc
@@ -112,6 +112,11 @@ void CalculatorGraphSubmodule(pybind11::module* module) {
                              "graph with a ValidatedGraphConfig object.");
         }
         auto calculator_graph = absl::make_unique<CalculatorGraph>();
+        // Disable default service initialization. This allows us to use
+        // the CPU versions of calculators that only optionally request
+        // kGpuService.
+        RaisePyErrorIfNotOk(
+            calculator_graph->DisallowServiceDefaultInitialization());
         RaisePyErrorIfNotOk(calculator_graph->Initialize(graph_config_proto));
         return calculator_graph.release();
       }),