Metal backend: compute init/execute times (pytorch#16639)

This pull request introduces detailed performance tracking and reporting
for the Metal backend, focusing on timing statistics for both
initialization and execution phases. It adds infrastructure to collect,
reset, and print timing data, and integrates statistics output into the
Parakeet model example. Additionally, a preprocessor macro is defined to
signal Metal backend availability.

**Metal Backend Performance Statistics:**

* Added timing statistics collection for `init()` and `execute()` calls
in the Metal backend, including total time, call count, and per-method
breakdowns. Accessor and reset functions are provided in `stats.h` and
implemented in `metal_backend.cpp`.
* Introduced a new `stats.cpp` file with a function to print all
collected Metal backend statistics, including per-method breakdowns for
both initialization and execution.

**Build and Integration Improvements:**

* Added `runtime/stats.cpp` to the Metal backend build sources and
defined the `ET_BUILD_METAL` preprocessor macro to indicate Metal
backend support.

**Example Model Enhancements:**

* Updated the Parakeet example (`main.cpp`) to show Metal backend timing
stats after model execution if Metal is enabled.

These changes provide better visibility into Metal backend performance
and make it easier to profile and optimize model execution on Apple
devices.

Author: manuelcandales

CMakePresets.json

@@ -167,7 +167,8 @@
         "llm-release"
       ],
       "cacheVariables": {
-        "EXECUTORCH_BUILD_METAL": "ON"
+        "EXECUTORCH_BUILD_METAL": "ON",
+        "EXECUTORCH_METAL_COLLECT_STATS": "OFF"
       },
       "condition": {
         "lhs": "${hostSystemName}",
@@ -216,6 +217,23 @@
         "rhs": "Darwin"
       }
     },
+    {
+      "name": "llm-metal-stats",
+      "displayName": "LLM Metal build with stats collection and logging",
+      "inherits": [
+        "llm-release-metal"
+      ],
+      "cacheVariables": {
+        "EXECUTORCH_METAL_COLLECT_STATS": "ON",
+        "EXECUTORCH_ENABLE_LOGGING": "ON",
+        "EXECUTORCH_LOG_LEVEL": "Info"
+      },
+      "condition": {
+        "lhs": "${hostSystemName}",
+        "type": "equals",
+        "rhs": "Darwin"
+      }
+    },
     {
       "name": "profiling",
       "displayName": "Build ExecuTorch with Profiling Enabled",
@@ -320,6 +338,15 @@
         "install"
       ],
       "jobs": 0
+    },
+    {
+      "name": "llm-metal-stats-install",
+      "displayName": "Build and install LLM extension artifacts with Metal stats",
+      "configurePreset": "llm-metal-stats",
+      "targets": [
+        "install"
+      ],
+      "jobs": 0
     }
   ],
   "workflowPresets": [
@@ -406,6 +433,20 @@
           "name": "llm-debug-metal-install"
         }
       ]
+    },
+    {
+      "name": "llm-metal-stats",
+      "displayName": "Configure, build and install ExecuTorch LLM extension with Metal stats and logging",
+      "steps": [
+        {
+          "type": "configure",
+          "name": "llm-metal-stats"
+        },
+        {
+          "type": "build",
+          "name": "llm-metal-stats-install"
+        }
+      ]
     }
   ]
 }

backends/apple/metal/CMakeLists.txt

@@ -36,6 +36,7 @@ find_package_torch()
 
 set(_aoti_metal_sources
     runtime/metal_backend.cpp
+    runtime/stats.cpp
     runtime/shims/memory.cpp
     runtime/shims/et_metal.mm
     runtime/shims/et_metal_ops.mm
@@ -68,6 +69,17 @@ target_link_libraries(
 
 target_compile_options(metal_backend PUBLIC -fexceptions -frtti -fPIC)
 
+# Define C++ preprocessor macro for Metal backend availability
+target_compile_definitions(metal_backend PUBLIC ET_BUILD_METAL)
+
+# Define macro for stats collection if enabled
+if(EXECUTORCH_METAL_COLLECT_STATS)
+  target_compile_definitions(
+    metal_backend PUBLIC EXECUTORCH_METAL_COLLECT_STATS
+  )
+  message(STATUS "Metal backend stats collection enabled")
+endif()
+
 target_link_options(metal_backend PUBLIC -Wl,-export_dynamic)
 
 # Find PyTorch's OpenMP library specifically for libtorch-less AOTI

backends/apple/metal/runtime/metal_backend.cpp

@@ -12,10 +12,12 @@
 #include <executorch/runtime/core/evalue.h>
 #include <executorch/runtime/core/exec_aten/util/tensor_util.h>
 #include <unistd.h>
+#include <chrono>
 #include <cstdio>
 
 #include <filesystem>
 #include <fstream>
+#include <mutex>
 #include <string>
 #include <vector>
 
@@ -29,9 +31,96 @@
 #include <executorch/backends/apple/metal/runtime/shims/shim_mps.h>
 #include <executorch/backends/apple/metal/runtime/shims/tensor_attribute.h>
 #include <executorch/backends/apple/metal/runtime/shims/utils.h>
+#include <executorch/backends/apple/metal/runtime/stats.h>
 
 namespace executorch::backends::metal {
 
+#ifdef EXECUTORCH_METAL_COLLECT_STATS
+
+// Per-method timing statistics
+struct MethodStats {
+  double total_ms = 0.0;
+  int64_t call_count = 0;
+};
+
+// Singleton struct containing all timing statistics and mutex
+struct StatsData {
+  std::mutex mutex;
+  double execute_total_ms = 0.0;
+  int64_t execute_call_count = 0;
+  double init_total_ms = 0.0;
+  int64_t init_call_count = 0;
+  std::unordered_map<std::string, MethodStats> method_stats;
+  std::unordered_map<std::string, MethodStats> init_method_stats;
+};
+
+// Thread-safe singleton accessor using C++11 magic statics
+static StatsData& get_stats_data() {
+  static StatsData instance;
+  return instance;
+}
+
+// Accessor functions for execute timing statistics
+double get_metal_backend_execute_total_ms() {
+  auto& stats = get_stats_data();
+  std::lock_guard<std::mutex> lock(stats.mutex);
+  return stats.execute_total_ms;
+}
+
+int64_t get_metal_backend_execute_call_count() {
+  auto& stats = get_stats_data();
+  std::lock_guard<std::mutex> lock(stats.mutex);
+  return stats.execute_call_count;
+}
+
+// Accessor functions for init timing statistics
+double get_metal_backend_init_total_ms() {
+  auto& stats = get_stats_data();
+  std::lock_guard<std::mutex> lock(stats.mutex);
+  return stats.init_total_ms;
+}
+
+int64_t get_metal_backend_init_call_count() {
+  auto& stats = get_stats_data();
+  std::lock_guard<std::mutex> lock(stats.mutex);
+  return stats.init_call_count;
+}
+
+void reset_metal_backend_stats() {
+  auto& stats = get_stats_data();
+  std::lock_guard<std::mutex> lock(stats.mutex);
+  stats.execute_total_ms = 0.0;
+  stats.execute_call_count = 0;
+  stats.init_total_ms = 0.0;
+  stats.init_call_count = 0;
+  stats.method_stats.clear();
+  stats.init_method_stats.clear();
+}
+
+std::unordered_map<std::string, std::pair<double, int64_t>>
+get_metal_backend_per_method_stats() {
+  auto& stats = get_stats_data();
+  std::lock_guard<std::mutex> lock(stats.mutex);
+  std::unordered_map<std::string, std::pair<double, int64_t>> result;
+  for (const auto& entry : stats.method_stats) {
+    result[entry.first] = {entry.second.total_ms, entry.second.call_count};
+  }
+  return result;
+}
+
+std::unordered_map<std::string, std::pair<double, int64_t>>
+get_metal_backend_init_per_method_stats() {
+  auto& stats = get_stats_data();
+  std::lock_guard<std::mutex> lock(stats.mutex);
+  std::unordered_map<std::string, std::pair<double, int64_t>> result;
+  for (const auto& entry : stats.init_method_stats) {
+    result[entry.first] = {entry.second.total_ms, entry.second.call_count};
+  }
+  return result;
+}
+
+#endif // EXECUTORCH_METAL_COLLECT_STATS
+
 #define LOAD_SYMBOL(handle, member, name, so_handle)                        \
   do {                                                                      \
     handle->member = reinterpret_cast<name##Func>(dlsym(so_handle, #name)); \
@@ -137,6 +226,9 @@ class ET_EXPERIMENTAL MetalBackend final
       FreeableBuffer* processed, // This will be a empty buffer
       ArrayRef<CompileSpec> compile_specs // This will be my empty list
   ) const override {
+#ifdef EXECUTORCH_METAL_COLLECT_STATS
+    auto init_start = std::chrono::high_resolution_clock::now();
+#endif
     ET_LOG(Info, "MetalBackend::init - Starting initialization");
 
     std::string method_name;
@@ -261,6 +353,29 @@ class ET_EXPERIMENTAL MetalBackend final
     }
 
     ET_LOG(Info, "MetalBackend::init - Initialization completed successfully");
+
+#ifdef EXECUTORCH_METAL_COLLECT_STATS
+    // Accumulate init timing statistics
+    auto init_end = std::chrono::high_resolution_clock::now();
+    double elapsed_ms =
+        std::chrono::duration<double, std::milli>(init_end - init_start)
+            .count();
+
+    {
+      auto& stats_data = get_stats_data();
+      std::lock_guard<std::mutex> lock(stats_data.mutex);
+      stats_data.init_total_ms += elapsed_ms;
+      stats_data.init_call_count++;
+
+      // Track per-method init timing
+      if (!method_name.empty()) {
+        auto& method_stats = stats_data.init_method_stats[method_name];
+        method_stats.total_ms += elapsed_ms;
+        method_stats.call_count++;
+      }
+    }
+#endif
+
     return (DelegateHandle*)handle; // Return the handle post-processing
   }
 
@@ -269,6 +384,9 @@ class ET_EXPERIMENTAL MetalBackend final
       BackendExecutionContext& context,
       DelegateHandle* handle_,
       Span<EValue*> args) const override {
+#ifdef EXECUTORCH_METAL_COLLECT_STATS
+    auto execute_start = std::chrono::high_resolution_clock::now();
+#endif
     ET_LOG(Debug, "MetalBackend execute");
 
     AOTIDelegateHandle* handle = (AOTIDelegateHandle*)handle_;
@@ -514,6 +632,29 @@ class ET_EXPERIMENTAL MetalBackend final
 
     ET_LOG(Debug, "MetalBackend execution completed successfully");
 
+#ifdef EXECUTORCH_METAL_COLLECT_STATS
+    // Accumulate timing statistics
+    auto execute_end = std::chrono::high_resolution_clock::now();
+    double elapsed_ms =
+        std::chrono::duration<double, std::milli>(execute_end - execute_start)
+            .count();
+
+    {
+      auto& stats_data = get_stats_data();
+      std::lock_guard<std::mutex> lock(stats_data.mutex);
+      stats_data.execute_total_ms += elapsed_ms;
+      stats_data.execute_call_count++;
+
+      // Track per-method timing
+      const char* method_name = context.get_method_name();
+      if (method_name != nullptr) {
+        auto& method_stats = stats_data.method_stats[method_name];
+        method_stats.total_ms += elapsed_ms;
+        method_stats.call_count++;
+      }
+    }
+#endif
+
     return Error::Ok;
   }
 

backends/apple/metal/runtime/stats.cpp

@@ -0,0 +1,86 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/apple/metal/runtime/stats.h>
+
+#ifdef EXECUTORCH_METAL_COLLECT_STATS
+
+#include <iostream>
+
+namespace executorch {
+namespace backends {
+namespace metal {
+
+void print_metal_backend_stats() {
+  std::cout << "\n--- Metal Backend Performance Statistics ---" << std::endl;
+
+  // Init stats
+  double metal_init_total_ms = get_metal_backend_init_total_ms();
+  int64_t metal_init_call_count = get_metal_backend_init_call_count();
+  std::cout << "Metal init() total: " << metal_init_total_ms << " ms ("
+            << metal_init_call_count << " calls)";
+  if (metal_init_call_count > 0) {
+    std::cout << " (avg: " << metal_init_total_ms / metal_init_call_count
+              << " ms/call)";
+  }
+  std::cout << std::endl;
+
+  // Per-method init breakdown
+  auto init_per_method_stats = get_metal_backend_init_per_method_stats();
+  if (!init_per_method_stats.empty()) {
+    std::cout << "  Per-method init breakdown:" << std::endl;
+    for (const auto& entry : init_per_method_stats) {
+      const std::string& method_name = entry.first;
+      double method_total_ms = entry.second.first;
+      int64_t method_call_count = entry.second.second;
+      std::cout << "    " << method_name << ": " << method_total_ms << " ms ("
+                << method_call_count << " calls)";
+      if (method_call_count > 0) {
+        std::cout << " (avg: " << method_total_ms / method_call_count
+                  << " ms/call)";
+      }
+      std::cout << std::endl;
+    }
+  }
+
+  // Execute stats
+  double metal_total_ms = get_metal_backend_execute_total_ms();
+  int64_t metal_call_count = get_metal_backend_execute_call_count();
+  std::cout << "\nMetal execute() total: " << metal_total_ms << " ms ("
+            << metal_call_count << " calls)";
+  if (metal_call_count > 0) {
+    std::cout << " (avg: " << metal_total_ms / metal_call_count << " ms/call)";
+  }
+  std::cout << std::endl;
+
+  // Per-method execute breakdown
+  auto per_method_stats = get_metal_backend_per_method_stats();
+  if (!per_method_stats.empty()) {
+    std::cout << "  Per-method execute breakdown:" << std::endl;
+    for (const auto& entry : per_method_stats) {
+      const std::string& method_name = entry.first;
+      double method_total_ms = entry.second.first;
+      int64_t method_call_count = entry.second.second;
+      std::cout << "    " << method_name << ": " << method_total_ms << " ms ("
+                << method_call_count << " calls)";
+      if (method_call_count > 0) {
+        std::cout << " (avg: " << method_total_ms / method_call_count
+                  << " ms/call)";
+      }
+      std::cout << std::endl;
+    }
+  }
+
+  std::cout << "--------------------------------------------\n" << std::endl;
+}
+
+} // namespace metal
+} // namespace backends
+} // namespace executorch
+
+#endif // EXECUTORCH_METAL_COLLECT_STATS