Merge branch 'master' into ios-llm-inference-refactor

google-ai-edge · Sep 6, 2024 · e1a2141 · e1a2141
2 parents a4c6a6a + 5dbd02e
commit e1a2141
Show file tree

Hide file tree

Showing 173 changed files with 4,373 additions and 820 deletions.
diff --git a/WORKSPACE b/WORKSPACE
@@ -104,20 +104,20 @@ http_archive(
 # KleidiAI is needed to get the best possible performance out of XNNPack
 http_archive(
     name = "KleidiAI",
-    sha256 = "e1a3a6a27dcae459e61c33f5eb235a7c809c3208b3b8a649f361a641269ebdc8",
-    strip_prefix = "kleidiai-8fda0bd9224cad4360c011a09bbb582c5ab7496a",
+    sha256 = "88233e427be6579560073267575f00f3b5fc370a31a43bbdd87a1810bd4bf1b6",
+    strip_prefix = "kleidiai-cddf991af5de49fd34949fa39690e4e906e04074",
     urls = [
-        "https://gitlab.arm.com/kleidi/kleidiai/-/archive/8fda0bd9224cad4360c011a09bbb582c5ab7496a/kleidiai-1976f8661e8d5aa7d4cdca0f3d2a915e5ecb4c53.zip",
+        "https://gitlab.arm.com/kleidi/kleidiai/-/archive/cddf991af5de49fd34949fa39690e4e906e04074/kleidiai-cddf991af5de49fd34949fa39690e4e906e04074.zip",
     ],
 )
 
 # XNNPACK on 2024-07-16
 http_archive(
     name = "XNNPACK",
     # `curl -L <url> | shasum -a 256`
-    sha256 = "b5a9406dac47e63980db281c52788f15a881519e95ed5140d927750c06828a2b",
-    strip_prefix = "XNNPACK-203651a9d4bc6b9c33e943c842d8bada139bb71d",
-    url = "https://github.com/google/XNNPACK/archive/203651a9d4bc6b9c33e943c842d8bada139bb71d.zip",
+    sha256 = "0e5d5c16686beff813e3946b26ca412f28acaf611228d20728ffb6479264fe19",
+    strip_prefix = "XNNPACK-9ddeb74f9f6866174d61888947e4aa9ffe963b1b",
+    url = "https://github.com/google/XNNPACK/archive/9ddeb74f9f6866174d61888947e4aa9ffe963b1b.zip",
 )
 
 # TODO: This is an are indirect depedency. We should factor it out.
@@ -576,9 +576,6 @@ http_archive(
         # Works around Bazel issue with objc_library.
         # See https://github.com/bazelbuild/bazel/issues/19912
         "@//third_party:org_tensorflow_objc_build_fixes.diff",
-        # Restores scores for text pipelines, which return different results
-        # with subgraph reshaping
-        "@//third_party:org_tensorflow_disable_subgraph_reshaping.diff"
     ],
     patch_args = [
         "-p1",

diff --git a/mediapipe/BUILD b/mediapipe/BUILD
@@ -196,6 +196,15 @@ config_setting_and_platform(
     visibility = ["//visibility:public"],
 )
 
+# Emscripten (WebAssembly)
+config_setting_and_platform(
+    name = "emscripten",
+    constraint_values = [
+        "@platforms//cpu:wasm32",
+    ],
+    visibility = ["//visibility:public"],
+)
+
 exports_files(
     ["provisioning_profile.mobileprovision"],
     visibility = ["//visibility:public"],

diff --git a/mediapipe/calculators/audio/BUILD b/mediapipe/calculators/audio/BUILD
@@ -347,10 +347,10 @@ cc_test(
         "//mediapipe/framework/formats:matrix",
         "//mediapipe/framework/formats:time_series_header_cc_proto",
         "//mediapipe/framework/port:gtest_main",
-        "//mediapipe/framework/port:integral_types",
-        "//mediapipe/framework/port:status",
         "//mediapipe/util:time_series_test_util",
+        "//mediapipe/util:time_series_util",
         "@com_google_absl//absl/log:absl_log",
+        "@com_google_absl//absl/status",
         "@com_google_audio_tools//audio/dsp:window_functions",
         "@eigen_archive//:eigen3",
     ],

diff --git a/mediapipe/calculators/audio/time_series_framer_calculator_test.cc b/mediapipe/calculators/audio/time_series_framer_calculator_test.cc
@@ -15,12 +15,13 @@
 #include <math.h>
 
 #include <cstdint>
-#include <memory>
+#include <new>
 #include <string>
 #include <vector>
 
 #include "Eigen/Core"
 #include "absl/log/absl_log.h"
+#include "absl/status/status.h"
 #include "audio/dsp/window_functions.h"
 #include "mediapipe/calculators/audio/time_series_framer_calculator.pb.h"
 #include "mediapipe/framework/calculator_framework.h"
@@ -29,8 +30,8 @@
 #include "mediapipe/framework/formats/time_series_header.pb.h"
 #include "mediapipe/framework/port/gmock.h"
 #include "mediapipe/framework/port/gtest.h"
-#include "mediapipe/framework/port/status.h"
 #include "mediapipe/util/time_series_test_util.h"
+#include "mediapipe/util/time_series_util.h"
 
 namespace mediapipe {
 namespace {
@@ -376,7 +377,7 @@ class TimeSeriesFramerCalculatorWindowingSanityTest
     MP_ASSERT_OK(RunGraph());
     ASSERT_EQ(1, output().packets.size());
     ASSERT_NEAR(expected_average * FrameDurationSamples(),
-                output().packets[0].Get<Matrix>().sum(), 1e-5);
+                output().packets[0].Get<Matrix>().sum(), 1e-3);
   }
 };
 
@@ -459,7 +460,7 @@ class TimeSeriesFramerCalculatorTimestampingTest
                kGapBetweenPacketsInSeconds;
   }
 
-  // Returns the timestamp inseconds based on cumulative timestamping.
+  // Returns the timestamp in seconds based on cumulative timestamping.
   double GetExpectedCumulativeTimestamp(int sample_index) {
     return kInitialTimestampOffsetMicroseconds * 1.0e-6 +
            sample_index / FrameDurationSamples() * FrameDurationSamples() /

diff --git a/mediapipe/calculators/image/BUILD b/mediapipe/calculators/image/BUILD
@@ -751,6 +751,7 @@ cc_library(
         "//mediapipe/gpu:gl_simple_shaders",
         "//mediapipe/gpu:gpu_buffer",
         "//mediapipe/gpu:gpu_origin_cc_proto",
+        "//mediapipe/gpu:gpu_origin_utils",
         "//mediapipe/gpu:shader_util",
         "@com_google_absl//absl/log:absl_log",
         "@com_google_absl//absl/memory",

diff --git a/mediapipe/calculators/image/affine_transformation_runner_gl.cc b/mediapipe/calculators/image/affine_transformation_runner_gl.cc
@@ -31,6 +31,7 @@
 #include "mediapipe/gpu/gl_simple_shaders.h"
 #include "mediapipe/gpu/gpu_buffer.h"
 #include "mediapipe/gpu/gpu_origin.pb.h"
+#include "mediapipe/gpu/gpu_origin_utils.h"
 #include "mediapipe/gpu/shader_util.h"
 
 namespace mediapipe {
@@ -43,24 +44,6 @@ using mediapipe::GlTexture;
 using mediapipe::GpuBuffer;
 using mediapipe::GpuOrigin;
 
-bool IsMatrixVerticalFlipNeeded(GpuOrigin::Mode gpu_origin) {
-  switch (gpu_origin) {
-    case GpuOrigin::DEFAULT:
-    case GpuOrigin::CONVENTIONAL:
-#ifdef __APPLE__
-      return false;
-#else
-      return true;
-#endif  //  __APPLE__
-    case GpuOrigin::TOP_LEFT:
-      return false;
-    default:
-      ABSL_LOG(ERROR) << "Incorrect GpuOrigin: "
-                      << static_cast<int>(gpu_origin);
-      return true;
-  }
-}
-
 #ifdef __APPLE__
 #define GL_CLAMP_TO_BORDER_MAY_BE_SUPPORTED 0
 #else
@@ -335,7 +318,9 @@ class GlTextureWarpAffineRunner
 
     // uniforms
     Eigen::Matrix<float, 4, 4, Eigen::RowMajor> eigen_mat(matrix.data());
-    if (IsMatrixVerticalFlipNeeded(gpu_origin_)) {
+    MP_ASSIGN_OR_RETURN(bool is_matrix_vertical_flip_needed,
+                        mediapipe::IsGpuOriginAtBottom(gpu_origin_));
+    if (is_matrix_vertical_flip_needed) {
       // @matrix describes affine transformation in terms of TOP LEFT origin, so
       // in some cases/on some platforms an extra flipping should be done before
       // and after.

diff --git a/mediapipe/calculators/tensor/BUILD b/mediapipe/calculators/tensor/BUILD
@@ -620,13 +620,12 @@ cc_library(
     deps = [
         ":inference_calculator_interface",
         ":inference_io_mapper",
+        ":inference_on_disk_cache_helper",
         ":tensor_span",
         "//mediapipe/framework:calculator_framework",
         "//mediapipe/framework:mediapipe_profiling",
         "//mediapipe/framework/api2:packet",
-        "//mediapipe/framework/deps:file_path",
         "//mediapipe/framework/formats:tensor",
-        "//mediapipe/framework/port:file_helpers",
         "//mediapipe/framework/port:ret_check",
         "//mediapipe/framework/port:status",
         "//mediapipe/gpu:gl_calculator_helper",
@@ -819,6 +818,7 @@ cc_library(
         "@org_tensorflow//tensorflow/lite/c:common",
         "@org_tensorflow//tensorflow/lite/core:framework",
     ] + select({
+        "//mediapipe:emscripten": [],
         "//conditions:default": [
             "//mediapipe/util:cpu_util",
         ],
@@ -845,6 +845,27 @@ cc_test(
     ],
 )
 
+cc_library(
+    name = "inference_on_disk_cache_helper",
+    srcs = ["inference_on_disk_cache_helper.cc"],
+    hdrs = ["inference_on_disk_cache_helper.h"],
+    tags = ["nomac"],
+    deps = [
+        ":inference_calculator_cc_proto",
+        ":inference_calculator_interface",
+        ":inference_io_mapper",
+        "//mediapipe/framework/deps:file_path",
+        "//mediapipe/framework/port:file_helpers",
+        "//mediapipe/framework/port:status",
+        "//mediapipe/util/tflite:tflite_gpu_runner",
+        "//mediapipe/util/tflite:tflite_model_loader",
+        "@com_google_absl//absl/log:absl_log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+    ],
+)
+
 cc_library(
     name = "inference_calculator_xnnpack",
     srcs = [
@@ -931,6 +952,7 @@ cc_library(
         "//mediapipe/framework/port:status",
         "//mediapipe/framework/port:statusor",
         "//mediapipe/gpu:gpu_origin_cc_proto",
+        "//mediapipe/gpu:gpu_origin_utils",
         "//mediapipe/util:resource_util",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:absl_check",
@@ -1369,6 +1391,9 @@ cc_library(
         "//mediapipe:ios": [
             "//mediapipe/util/android/file/base",
         ],
+        "//mediapipe:emscripten": [
+            "//mediapipe/util/android/file/base",
+        ],
         "//conditions:default": [
             "//mediapipe/framework/port:file_helpers",
         ],
@@ -1475,6 +1500,8 @@ cc_library(
         "//mediapipe/framework/port:status",
         "//mediapipe/framework/port:statusor",
         "//mediapipe/gpu:gpu_origin_cc_proto",
+        "//mediapipe/gpu:gpu_origin_utils",
+        "//mediapipe/gpu/webgpu:webgpu_check",
         "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/log:absl_log",
     ] + select({
@@ -1502,6 +1529,12 @@ cc_library(
             "//mediapipe/gpu:gpu_buffer",
             "//mediapipe/gpu:gpu_service",
         ],
+        "//mediapipe:emscripten": [
+            ":image_to_tensor_converter_gl_texture",
+            "//mediapipe/gpu:gl_calculator_helper",
+            "//mediapipe/gpu:gpu_buffer",
+            "//mediapipe/gpu:gpu_service",
+        ],
         "//mediapipe:apple": [
             ":image_to_tensor_converter_metal",
             "//mediapipe/gpu:gl_calculator_helper",
@@ -1515,6 +1548,12 @@ cc_library(
             "//mediapipe/gpu:gpu_buffer",
             "//mediapipe/gpu:gpu_service",
         ],
+    }) + select({
+        "//mediapipe/gpu/webgpu:use_webgpu": [
+            "//mediapipe/gpu/webgpu:image_to_tensor_converter_webgpu_texture",
+            "//mediapipe/gpu/webgpu:webgpu_texture_buffer",
+        ],
+        "//conditions:default": [],
     }),
 )
 
@@ -1653,6 +1692,7 @@ cc_library(
     hdrs = ["image_to_tensor_converter_gl_buffer.h"],
     visibility = ["//visibility:private"],
     deps = ["//mediapipe/framework:port"] + selects.with_or({
+        "//mediapipe:emscripten": [],
         "//mediapipe:apple": [],
         "//conditions:default": [
             ":image_to_tensor_converter",
@@ -1874,6 +1914,10 @@ cc_library(
             ":tensors_to_segmentation_converter_gl_buffer",
             "//mediapipe/gpu:gl_calculator_helper",
         ],
+        "//mediapipe:emscripten": [
+            ":tensors_to_segmentation_converter_gl_texture",
+            "//mediapipe/gpu:gl_calculator_helper",
+        ],
         ":platform_apple_with_gpu": [
             ":tensors_to_segmentation_converter_metal",
             "//mediapipe/gpu:MPPMetalHelper",
@@ -1973,6 +2017,7 @@ cc_library(
             "//mediapipe/gpu:gpu_buffer",
             "//mediapipe/gpu:gpu_buffer_format",
             "//mediapipe/gpu:gpu_origin_cc_proto",
+            "//mediapipe/gpu:gpu_origin_utils",
             "//mediapipe/gpu:shader_util",
             "@com_google_absl//absl/status",
             "@com_google_absl//absl/status:statusor",
@@ -1987,6 +2032,7 @@ cc_library(
     srcs = ["tensors_to_segmentation_converter_gl_buffer.cc"],
     hdrs = ["tensors_to_segmentation_converter_gl_buffer.h"],
     deps = ["//mediapipe/framework:port"] + selects.with_or({
+        "//mediapipe:emscripten": [],
         "//mediapipe:apple": [],
         "//conditions:default": [
             ":tensors_to_segmentation_calculator_cc_proto",
@@ -2004,6 +2050,7 @@ cc_library(
             "//mediapipe/framework/port:ret_check",
             "//mediapipe/framework/port:status",
             "//mediapipe/gpu:gpu_origin_cc_proto",
+            "//mediapipe/gpu:gpu_origin_utils",
             "//mediapipe/gpu:gl_base_hdr",
             "//mediapipe/gpu:gl_calculator_helper",
             "//mediapipe/gpu:gl_simple_shaders",
@@ -2053,6 +2100,7 @@ cc_library(
             "//mediapipe/gpu:gl_calculator_helper",
             "//mediapipe/gpu:gl_simple_shaders",
             "//mediapipe/gpu:gpu_origin_cc_proto",
+            "//mediapipe/gpu:gpu_origin_utils",
             "//mediapipe/gpu:shader_util",
             "//mediapipe/util:resource_util",
             "@org_tensorflow//tensorflow/lite:framework",
@@ -2083,7 +2131,7 @@ cc_test(
 
 cc_library(
     name = "tensors_to_segmentation_calculator_test_utils",
-    testonly = 1,
+    testonly = True,
     srcs = ["tensors_to_segmentation_calculator_test_utils.cc"],
     hdrs = ["tensors_to_segmentation_calculator_test_utils.h"],
     deps = [

diff --git a/mediapipe/calculators/tensor/image_to_tensor_calculator.cc b/mediapipe/calculators/tensor/image_to_tensor_calculator.cc
@@ -37,6 +37,8 @@
 #include "mediapipe/framework/port/status.h"
 #include "mediapipe/framework/port/statusor.h"
 #include "mediapipe/gpu/gpu_origin.pb.h"
+#include "mediapipe/gpu/gpu_origin_utils.h"
+#include "mediapipe/gpu/webgpu/webgpu_check.h"
 
 #if !MEDIAPIPE_DISABLE_OPENCV
 #include "mediapipe/calculators/tensor/image_to_tensor_converter_opencv.h"
@@ -58,6 +60,11 @@
 #include "mediapipe/calculators/tensor/image_to_tensor_converter_gl_texture.h"
 #include "mediapipe/gpu/gl_calculator_helper.h"
 #include "mediapipe/gpu/gpu_service.h"
+#if MEDIAPIPE_USE_WEBGPU
+#include "mediapipe/gpu/webgpu/image_to_tensor_converter_webgpu_texture.h"
+#include "mediapipe/gpu/webgpu/webgpu_service.h"
+#include "mediapipe/gpu/webgpu/webgpu_texture_buffer.h"
+#endif  // MEDIAPIPE_USE_WEBGPU
 #endif  // MEDIAPIPE_METAL_ENABLED
 #endif  // !MEDIAPIPE_DISABLE_GPU
 
@@ -161,6 +168,9 @@ class ImageToTensorCalculator : public Node {
     MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
 #else
     cc->UseService(kGpuService).Optional();
+#if MEDIAPIPE_USE_WEBGPU
+    cc->UseService(kWebGpuService).Optional();
+#endif  // MEDIAPIPE_USE_WEBGPU
 #endif  // MEDIAPIPE_METAL_ENABLED
 #endif  // MEDIAPIPE_DISABLE_GPU
 
@@ -269,16 +279,26 @@ class ImageToTensorCalculator : public Node {
             gpu_converter_,
             CreateMetalConverter(cc, GetBorderMode(options_.border_mode())));
 #elif MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
+        MP_ASSIGN_OR_RETURN(bool input_starts_at_bottom,
+                            IsGpuOriginAtBottom(options_.gpu_origin()));
         MP_ASSIGN_OR_RETURN(gpu_converter_,
                             CreateImageToGlBufferTensorConverter(
-                                cc, DoesGpuInputStartAtBottom(options_),
+                                cc, input_starts_at_bottom,
                                 GetBorderMode(options_.border_mode())));
 #else
+        if (IsWebGpuAvailable()) {
+#if MEDIAPIPE_USE_WEBGPU
+          MP_ASSIGN_OR_RETURN(gpu_converter_,
+                              CreateImageToWebGpuTextureTensorConverter(cc));
+#endif  // MEDIAPIPE_USE_WEBGPU
+        }
 #if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
         if (!gpu_converter_) {
+          MP_ASSIGN_OR_RETURN(bool input_starts_at_bottom,
+                              IsGpuOriginAtBottom(options_.gpu_origin()));
           MP_ASSIGN_OR_RETURN(gpu_converter_,
                               CreateImageToGlTextureTensorConverter(
-                                  cc, DoesGpuInputStartAtBottom(options_),
+                                  cc, input_starts_at_bottom,
                                   GetBorderMode(options_.border_mode())));
         }
 #endif  // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30

diff --git a/mediapipe/calculators/tensor/image_to_tensor_utils.h b/mediapipe/calculators/tensor/image_to_tensor_utils.h
@@ -201,12 +201,6 @@ OutputTensorParams GetOutputTensorParams(const T& options) {
   return params;
 }
 
-// Returns whether the GPU input format starts at the bottom.
-template <typename T>
-bool DoesGpuInputStartAtBottom(const T& options) {
-  return options.gpu_origin() != mediapipe::GpuOrigin_Mode_TOP_LEFT;
-}
-
 // Converts the BorderMode proto into struct.
 BorderMode GetBorderMode(
     const mediapipe::ImageToTensorCalculatorOptions::BorderMode& mode);