Add runtime control to Android ETDump profiling

extension/android/CMakeLists.txt

@@ -91,11 +91,35 @@ list(
   fbjni
 )
 
+# =============================================================================
+# Android ETDump Profiling with Runtime Control
+# =============================================================================
 if(EXECUTORCH_ANDROID_PROFILING)
-  list(APPEND link_libraries etdump flatccrt)
+  # Add ETDump JNI implementation
+  target_sources(executorch_jni PRIVATE jni/jni_etdump.cpp)
+
+  # Add SDK include directories for ETDump headers
+  target_include_directories(
+    executorch_jni
+    PRIVATE
+      ${EXECUTORCH_ROOT}/sdk
+      ${EXECUTORCH_ROOT}/third-party/flatcc/include
+  )
+
+  # Add compile definitions
   target_compile_definitions(
-    executorch_jni PUBLIC EXECUTORCH_ANDROID_PROFILING=1
+    executorch_jni
+    PUBLIC
+      ET_EVENT_TRACER_ENABLED
+      EXECUTORCH_ANDROID_PROFILING=1
   )
+
+  # Link ETDump libraries
+  list(APPEND link_libraries etdump flatccrt)
+
+  message(STATUS "Android ETDump profiling with runtime control: ENABLED")
+else()
+  message(STATUS "Android ETDump profiling: DISABLED")
 endif()
 
 if(TARGET optimized_native_cpu_ops_lib)

extension/android/executorch_android/src/main/java/org/pytorch/executorch/ETDump.java

@@ -0,0 +1,95 @@
+/*
+ * 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.
+ */
+
+package org.pytorch.executorch;
+
+import android.util.Log;
+
+/**
+ * ETDump provides runtime control for ExecuTorch profiling.
+ *
+ * <p>Enable profiling before loading models to capture execution traces. Use
+ * Module.writeETDumpToPath() to write profiling data to custom locations.
+ *
+ * <p>Example usage:
+ *
+ * <pre>{@code
+ * // Enable profiling
+ * ETDump.enableProfiling();
+ *
+ * // Load and run model
+ * Module module = Module.load("model.pte");
+ * module.forward(inputs);
+ *
+ * // Write profiling data to custom path (no root access needed!)
+ * module.writeETDumpToPath(getCacheDir() + "/profile.etdump");
+ *
+ * // Disable profiling
+ * ETDump.disableProfiling();
+ * }</pre>
+ */
+public class ETDump {
+  private static final String TAG = "ExecuTorch-ETDump";
+
+  static {
+    try {
+      System.loadLibrary("executorch");
+      nativeInit();
+    } catch (UnsatisfiedLinkError e) {
+      Log.e(TAG, "Failed to load executorch library", e);
+      throw e;
+    }
+  }
+
+  /** Initialize the ETDump subsystem. Called automatically when the class is loaded. */
+  private static native void nativeInit();
+
+  /**
+   * Enable profiling for subsequently loaded models.
+   *
+   * <p>Modules loaded after calling this method will capture profiling data. This has zero runtime
+   * overhead until a model is actually loaded.
+   *
+   * @return true if profiling was successfully enabled
+   */
+  public static boolean enableProfiling() {
+    boolean result = nativeEnableProfiling();
+    if (result) {
+      Log.i(TAG, "Profiling enabled");
+    } else {
+      Log.e(TAG, "Failed to enable profiling");
+    }
+    return result;
+  }
+
+  /**
+   * Disable profiling.
+   *
+   * <p>Modules loaded after calling this method will not capture profiling data.
+   */
+  public static void disableProfiling() {
+    nativeDisableProfiling();
+    Log.i(TAG, "Profiling disabled");
+  }
+
+  /**
+   * Check if profiling is currently enabled.
+   *
+   * @return true if profiling is enabled
+   */
+  public static boolean isProfilingEnabled() {
+    return nativeIsProfilingEnabled();
+  }
+
+  // Native methods
+  private static native boolean nativeEnableProfiling();
+
+  private static native void nativeDisableProfiling();
+
+  private static native boolean nativeIsProfilingEnabled();
+}

extension/android/executorch_android/src/main/java/org/pytorch/executorch/Module.java

@@ -257,4 +257,71 @@ public void destroy() {
               + " released.");
     }
   }
+
+  /**
+   * Write ETDump profiling data to a specific path.
+   *
+   * <p>This allows writing profiling data to custom locations such as app cache directory or
+   * external storage. No root access required.
+   *
+   * <p>Example:
+   *
+   * <pre>{@code
+   * Module module = Module.load("model.pte");
+   * module.forward(inputs);
+   * module.writeETDumpToPath(getCacheDir() + "/profile.etdump");
+   * }</pre>
+   *
+   * @param outputPath The file path where ETDump data will be written. Must be a writable location.
+   * @return true if the data was successfully written, false otherwise
+   */
+  @DoNotStrip
+  public boolean writeETDumpToPath(String outputPath) {
+    if (outputPath == null || outputPath.isEmpty()) {
+      android.util.Log.e("ExecuTorch-Module", "Output path cannot be null or empty");
+      return false;
+    }
+
+    // Validate that parent directory exists
+    java.io.File file = new java.io.File(outputPath);
+    java.io.File parentDir = file.getParentFile();
+    if (parentDir != null && !parentDir.exists()) {
+      if (!parentDir.mkdirs()) {
+        android.util.Log.e("ExecuTorch-Module", "Failed to create parent directory: " + parentDir);
+        return false;
+      }
+    }
+
+    return nativeWriteETDumpToPath(outputPath);
+  }
+
+  /**
+   * Get ETDump profiling data as a byte array.
+   *
+   * <p>This is useful for custom handling of profiling data, such as uploading to a server or
+   * processing in-memory without writing to disk.
+   *
+   * <p>Example:
+   *
+   * <pre>{@code
+   * Module module = Module.load("model.pte");
+   * module.forward(inputs);
+   * byte[] profileData = module.getETDumpData();
+   * if (profileData != null) {
+   *   uploadToServer(profileData);
+   * }
+   * }</pre>
+   *
+   * @return byte array containing the ETDump data, or null if no data is available
+   */
+  @DoNotStrip
+  public byte[] getETDumpData() {
+    return nativeGetETDumpData();
+  }
+
+  @DoNotStrip
+  private native boolean nativeWriteETDumpToPath(String outputPath);
+
+  @DoNotStrip
+  private native byte[] nativeGetETDumpData();
 }

extension/android/jni/jni_etdump.cpp

@@ -0,0 +1,80 @@
+/*
+ * 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 <fbjni/fbjni.h>
+#include <jni.h>
+
+#ifdef EXECUTORCH_ANDROID_PROFILING
+
+#include "jni_etdump.h"
+
+namespace executorch {
+namespace extension {
+namespace jni {
+
+// ============================================================================
+// Global ETDump Manager
+// ============================================================================
+
+// Global ETDump manager (one per process)
+static std::unique_ptr<ETDumpManager> g_etdump_manager = nullptr;
+
+extern "C" ETDumpManager* getGlobalETDumpManager() {
+  return g_etdump_manager.get();
+}
+
+// ============================================================================
+// JNI Methods for ETDump Java class
+// ============================================================================
+
+// Initialize ETDump manager
+extern "C" JNIEXPORT void JNICALL
+Java_org_pytorch_executorch_ETDump_nativeInit(JNIEnv* env, jclass clazz) {
+  if (!g_etdump_manager) {
+    g_etdump_manager = std::make_unique<ETDumpManager>();
+  }
+}
+
+// Enable profiling
+extern "C" JNIEXPORT jboolean JNICALL
+Java_org_pytorch_executorch_ETDump_nativeEnableProfiling(
+    JNIEnv* env,
+    jclass clazz) {
+  if (!g_etdump_manager) {
+    g_etdump_manager = std::make_unique<ETDumpManager>();
+  }
+  g_etdump_manager->enableProfiling();
+  return JNI_TRUE;
+}
+
+// Disable profiling
+extern "C" JNIEXPORT void JNICALL
+Java_org_pytorch_executorch_ETDump_nativeDisableProfiling(
+    JNIEnv* env,
+    jclass clazz) {
+  if (g_etdump_manager) {
+    g_etdump_manager->disableProfiling();
+  }
+}
+
+// Check if profiling is enabled
+extern "C" JNIEXPORT jboolean JNICALL
+Java_org_pytorch_executorch_ETDump_nativeIsProfilingEnabled(
+    JNIEnv* env,
+    jclass clazz) {
+  if (!g_etdump_manager) {
+    return JNI_FALSE;
+  }
+  return g_etdump_manager->isProfilingEnabled() ? JNI_TRUE : JNI_FALSE;
+}
+
+} // namespace jni
+} // namespace extension
+} // namespace executorch
+
+#endif // EXECUTORCH_ANDROID_PROFILING

extension/android/jni/jni_etdump.h

@@ -0,0 +1,63 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#ifdef EXECUTORCH_ANDROID_PROFILING
+#include <executorch/devtools/etdump/etdump_flatcc.h>
+
+namespace executorch {
+namespace extension {
+namespace jni {
+
+/**
+ * ETDumpManager manages the global profiling enabled state.
+ * This is a singleton that controls whether newly loaded modules
+ * should capture profiling data.
+ */
+class ETDumpManager {
+ public:
+  ETDumpManager() : profiling_enabled_(false) {}
+
+  /**
+   * Enable profiling for subsequently loaded modules.
+   */
+  void enableProfiling() {
+    profiling_enabled_ = true;
+  }
+
+  /**
+   * Disable profiling.
+   */
+  void disableProfiling() {
+    profiling_enabled_ = false;
+  }
+
+  /**
+   * Check if profiling is currently enabled.
+   * @return true if profiling is enabled
+   */
+  bool isProfilingEnabled() const {
+    return profiling_enabled_;
+  }
+
+ private:
+  bool profiling_enabled_;
+};
+
+/**
+ * Get the global ETDump manager instance.
+ * @return Pointer to the global ETDumpManager, or nullptr if not initialized
+ */
+extern "C" ETDumpManager* getGlobalETDumpManager();
+
+} // namespace jni
+} // namespace extension
+} // namespace executorch
+
+#endif // EXECUTORCH_ANDROID_PROFILING