Allow Python to be build on Mac with GPU support

PiperOrigin-RevId: 574625520
google-ai-edge · Oct 18, 2023 · 364048d · 364048d
1 parent 4f29ffc
commit 364048d
Show file tree

Hide file tree

Showing 8 changed files with 80 additions and 19 deletions.
diff --git a/mediapipe/calculators/core/BUILD b/mediapipe/calculators/core/BUILD
@@ -21,10 +21,10 @@ licenses(["notice"])
 package(default_visibility = ["//visibility:public"])
 
 selects.config_setting_group(
-    name = "ios_or_disable_gpu",
+    name = "apple_or_disable_gpu",
     match_any = [
         "//mediapipe/gpu:disable_gpu",
-        "//mediapipe:ios",
+        "//mediapipe:apple",
     ],
 )
 
@@ -299,7 +299,7 @@ cc_library(
         "//mediapipe/util:render_data_cc_proto",
         "@org_tensorflow//tensorflow/lite:framework",
     ] + select({
-        ":ios_or_disable_gpu": [],
+        ":apple_or_disable_gpu": [],
         "//conditions:default": [
             "@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_buffer",
         ],
@@ -913,7 +913,7 @@ cc_library(
         "@org_tensorflow//tensorflow/lite:framework",
         "@org_tensorflow//tensorflow/lite/kernels:builtin_ops",
     ] + select({
-        ":ios_or_disable_gpu": [],
+        ":apple_or_disable_gpu": [],
         "//conditions:default": [
             "@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_buffer",
         ],

diff --git a/mediapipe/calculators/tensor/BUILD b/mediapipe/calculators/tensor/BUILD
@@ -50,10 +50,18 @@ more_selects.config_setting_negation(
 )
 
 selects.config_setting_group(
-    name = "platform_ios_with_gpu",
+    name = "platform_apple_with_gpu",
     match_all = [
         ":not_disable_gpu",
-        "//mediapipe:ios",
+        "//mediapipe:apple",
+    ],
+)
+
+selects.config_setting_group(
+    name = "platform_apple_without_gpu",
+    match_all = [
+        ":disable_gpu",
+        "//mediapipe:apple",
     ],
 )
 
@@ -614,7 +622,7 @@ cc_library(
         ":inference_calculator_interface",
     ] + select({
         "//conditions:default": [":inference_calculator_gl_if_compute_shader_available"],
-        ":platform_ios_with_gpu": [":inference_calculator_metal"],
+        ":platform_apple_with_gpu": [":inference_calculator_metal"],
     }),
     alwayslink = 1,
 )
@@ -687,12 +695,13 @@ cc_library(
             "//mediapipe/gpu:gl_calculator_helper",
             "//mediapipe/gpu:gpu_buffer",
         ],
-        "//mediapipe:ios": [
+        ":platform_apple_with_gpu": [
             "//mediapipe/gpu:MPPMetalHelper",
             "//mediapipe/gpu:MPPMetalUtil",
             "//mediapipe/objc:mediapipe_framework_ios",
         ],
-        "//mediapipe:macos": [],
+        # This setting is needed to allow bazel to build all targets on Mac with GPU disabled
+        ":platform_apple_without_gpu": [],
         "//conditions:default": [
             "//mediapipe/gpu:gl_calculator_helper",
             "//mediapipe/gpu:gl_simple_shaders",
@@ -777,11 +786,12 @@ cc_library(
     name = "tensors_to_detections_calculator_gpu_deps",
     visibility = ["//visibility:private"],
     deps = select({
-        "//mediapipe:ios": [
+        ":platform_apple_with_gpu": [
             "//mediapipe/gpu:MPPMetalHelper",
             "//mediapipe/gpu:MPPMetalUtil",
         ],
-        "//mediapipe:macos": [],
+        # This setting is needed to allow bazel to build all targets on Mac with GPU disabled
+        ":platform_apple_without_gpu": [],
         "//conditions:default": [
             "//mediapipe/gpu:gl_calculator_helper",
         ],
@@ -1428,7 +1438,7 @@ cc_library(
         ],
     }) + selects.with_or({
         ":gpu_inference_disabled": [],
-        "//mediapipe:ios": [
+        ":platform_apple_with_gpu": [
             "//mediapipe/gpu:MPPMetalUtil",
             "//mediapipe/gpu:MPPMetalHelper",
             "//third_party/apple_frameworks:MetalKit",

diff --git a/mediapipe/framework/formats/BUILD b/mediapipe/framework/formats/BUILD
@@ -456,7 +456,7 @@ cc_library(
         "tensor.h",
         "//mediapipe/framework/formats/tensor:internal.h",
     ] + select({
-        "//mediapipe:ios": ["tensor_mtl_buffer_view.h"],
+        "//mediapipe:apple": ["tensor_mtl_buffer_view.h"],
         "//conditions:default": [],
     }),
     copts = select({
@@ -477,7 +477,7 @@ cc_library(
         "//conditions:default": [],
     }),
     linkopts = select({
-        "//mediapipe:ios": [
+        "//mediapipe:apple": [
             "-framework CoreVideo",
             "-framework MetalKit",
         ],

diff --git a/mediapipe/gpu/BUILD b/mediapipe/gpu/BUILD
@@ -96,7 +96,7 @@ GL_BASE_LINK_OPTS_OSS = GL_BASE_LINK_OPTS + select({
         "-lEGL",
     ],
     "//mediapipe:android": [],
-    "//mediapipe:ios": [],
+    ":platform_apple_with_gpu": [],
     ":disable_gpu": [],
 })
 
@@ -246,7 +246,7 @@ cc_library(
         "//mediapipe/framework/formats:image_frame",
     ] + select({
         "//conditions:default": [],
-        ":platform_ios_with_gpu": [
+        ":platform_apple_with_gpu": [
             ":gl_texture_util",
             ":gpu_buffer_storage_cv_pixel_buffer",
         ],
@@ -287,6 +287,14 @@ selects.config_setting_group(
     ],
 )
 
+selects.config_setting_group(
+    name = "platform_apple_with_gpu",
+    match_all = [
+        ":not_disable_gpu",
+        "//mediapipe:apple",
+    ],
+)
+
 cc_library(
     name = "gpu_buffer",
     srcs = ["gpu_buffer.cc"],
@@ -315,6 +323,7 @@ cc_library(
         ":platform_macos_with_gpu": [
             ":gl_texture_buffer",
             ":gl_texture_view",
+            ":gpu_buffer_storage_cv_pixel_buffer",
             "//mediapipe/objc:CFHolder",
         ],
         ":disable_gpu": [],

diff --git a/mediapipe/gpu/gpu_buffer_storage_cv_pixel_buffer.cc b/mediapipe/gpu/gpu_buffer_storage_cv_pixel_buffer.cc
@@ -43,7 +43,7 @@ GlTextureView GpuBufferStorageCvPixelBuffer::GetTexture(
   cv_texture.adopt(cv_texture_temp);
   return GlTextureView(
       gl_context.get(), CVOpenGLTextureGetTarget(*cv_texture),
-      CVOpenGLTextureGetName(*cv_texture), width(), height(), *this, plane,
+      CVOpenGLTextureGetName(*cv_texture), width(), height(), plane,
       [cv_texture](mediapipe::GlTextureView&) { /* only retains cv_texture */ },
       done_writing);
 #else

diff --git a/mediapipe/tasks/cc/vision/face_stylizer/calculators/BUILD b/mediapipe/tasks/cc/vision/face_stylizer/calculators/BUILD
@@ -12,12 +12,43 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+load("@bazel_skylib//lib:selects.bzl", "selects")
+load("//mediapipe/framework:more_selects.bzl", "more_selects")
 load("//mediapipe/framework/port:build_config.bzl", "mediapipe_proto_library")
 
 licenses(["notice"])
 
 package(default_visibility = ["//mediapipe/tasks:internal"])
 
+config_setting(
+    name = "disable_gpu",
+    define_values = {
+        "MEDIAPIPE_DISABLE_GPU": "1",
+    },
+    visibility = ["//visibility:public"],
+)
+
+more_selects.config_setting_negation(
+    name = "not_disable_gpu",
+    negate = ":disable_gpu",
+)
+
+selects.config_setting_group(
+    name = "platform_apple_with_gpu",
+    match_all = [
+        ":not_disable_gpu",
+        "//mediapipe:apple",
+    ],
+)
+
+selects.config_setting_group(
+    name = "platform_apple_without_gpu",
+    match_all = [
+        ":disable_gpu",
+        "//mediapipe:apple",
+    ],
+)
+
 mediapipe_proto_library(
     name = "tensors_to_image_calculator_proto",
     srcs = ["tensors_to_image_calculator.proto"],
@@ -91,13 +122,14 @@ cc_library(
             "@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_texture",
             "@org_tensorflow//tensorflow/lite/delegates/gpu/gl/converters:util",
         ],
-        "//mediapipe:ios": [
+        ":platform_apple_with_gpu": [
             "//mediapipe/gpu:MPPMetalHelper",
             "//mediapipe/gpu:MPPMetalUtil",
             "//mediapipe/gpu:gl_calculator_helper",
             "//mediapipe/gpu:gpu_buffer",
         ],
-        "//mediapipe:macos": [],
+        # This setting is needed to allow bazel to build all targets on Mac with GPU disabled
+        ":platform_apple_without_gpu": [],
         "//conditions:default": [
             "//mediapipe/gpu:gl_calculator_helper",
             "//mediapipe/gpu:gl_quad_renderer",

diff --git a/mediapipe/util/tflite/config.h b/mediapipe/util/tflite/config.h
@@ -32,6 +32,10 @@
 #define MEDIAPIPE_TFLITE_METAL_INFERENCE 0
 #endif  // MEDIAPIPE_IOS
 
+#if TARGET_OS_OSX && !MEDIAPIPE_DISABLE_GPU
+#define MEDIAPIPE_TFLITE_METAL_INFERENCE 1
+#endif  // TARGET_OS_OSX && !MEDIAPIPE_DISABLE_GPU
+
 #define MEDIAPIPE_TFLITE_GPU_SUPPORTED \
   ((MEDIAPIPE_TFLITE_GL_INFERENCE) || (MEDIAPIPE_TFLITE_METAL_INFERENCE))
 

diff --git a/setup.py b/setup.py
@@ -40,12 +40,18 @@
 MP_ROOT_INIT_PY = os.path.join(MP_ROOT_PATH, '__init__.py')
 
 GPU_OPTIONS_DISBALED = ['--define=MEDIAPIPE_DISABLE_GPU=1']
+
 GPU_OPTIONS_ENBALED = [
     '--copt=-DTFLITE_GPU_EXTRA_GLES_DEPS',
     '--copt=-DMEDIAPIPE_OMIT_EGL_WINDOW_BIT',
     '--copt=-DMESA_EGL_NO_X11_HEADERS',
     '--copt=-DEGL_NO_X11',
 ]
+if IS_MAC:
+  GPU_OPTIONS_ENBALED.append(
+      '--copt=-DMEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER'
+  )
+
 GPU_OPTIONS = GPU_OPTIONS_DISBALED if MP_DISABLE_GPU else GPU_OPTIONS_ENBALED