Force xnnpack when CPU inference is enforced

Also renames the flag to MEDIAPIPE_ to abide to new naming scheme. PiperOrigin-RevId: 683942820
kp-forks · Oct 9, 2024 · c4f475e · c4f475e
1 parent a6637e4
commit c4f475e
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 5 deletions.
diff --git a/mediapipe/calculators/tensor/BUILD b/mediapipe/calculators/tensor/BUILD
@@ -430,7 +430,7 @@ cc_library_with_tflite(
     hdrs = ["inference_calculator.h"],
     local_defines = select({
         ":force_cpu_inference": ["MEDIAPIPE_FORCE_CPU_INFERENCE=1"],
-        "//conditions:default": [],
+        "//conditions:default": ["MEDIAPIPE_FORCE_CPU_INFERENCE=0"],
     }),
     tflite_deps = [
         ":inference_runner",
@@ -782,6 +782,10 @@ cc_library(
     srcs = [
         "inference_calculator_cpu.cc",
     ],
+    local_defines = select({
+        ":force_cpu_inference": ["MEDIAPIPE_FORCE_CPU_INFERENCE=1"],
+        "//conditions:default": ["MEDIAPIPE_FORCE_CPU_INFERENCE=0"],
+    }),
     deps = [
         ":inference_calculator_interface",
         ":inference_calculator_utils",

diff --git a/mediapipe/calculators/tensor/inference_calculator.cc b/mediapipe/calculators/tensor/inference_calculator.cc
@@ -51,7 +51,7 @@ class InferenceCalculatorSelectorImpl
             subgraph_node);
     std::vector<absl::string_view> impls;
 
-#if !defined(MEDIAPIPE_FORCE_CPU_INFERENCE) || !MEDIAPIPE_FORCE_CPU_INFERENCE
+#if !MEDIAPIPE_FORCE_CPU_INFERENCE
 
     const bool should_use_gpu =
         !options.has_delegate() ||  // Use GPU delegate if not specified
@@ -73,8 +73,7 @@ class InferenceCalculatorSelectorImpl
         impls.emplace_back("GlAdvanced");
       }
     }
-#endif  // !defined(MEDIAPIPE_FORCE_CPU_INFERENCE) ||
-        // !MEDIAPIPE_FORCE_CPU_INFERENCE
+#endif  // !MEDIAPIPE_FORCE_CPU_INFERENCE
     impls.emplace_back("Cpu");
     impls.emplace_back("Xnnpack");
     std::vector<std::string> missing_impls;

diff --git a/mediapipe/calculators/tensor/inference_calculator_cpu.cc b/mediapipe/calculators/tensor/inference_calculator_cpu.cc
@@ -146,7 +146,7 @@ InferenceCalculatorCpuImpl::MaybeCreateDelegate(CalculatorContext* cc) {
   }
 #endif  // MEDIAPIPE_ANDROID
 
-#if defined(__EMSCRIPTEN__)
+#if defined(__EMSCRIPTEN__) || MEDIAPIPE_FORCE_CPU_INFERENCE
   const bool use_xnnpack = true;
 #else
   const bool use_xnnpack = opts_has_delegate && opts_delegate.has_xnnpack();