Added ETDump to Wasm

Author: Conarnar

.github/workflows/pull.yml

@@ -801,6 +801,8 @@ jobs:
       id-token: write
       contents: read
     strategy:
+      matrix:
+        enable-etdump: ['', '--enable-etdump']
       fail-fast: false
     with:
       runner: linux.2xlarge
@@ -820,7 +822,7 @@ jobs:
         source .ci/scripts/setup-emscripten.sh
 
         # Test selective build
-        bash scripts/build_wasm_tests.sh
+        bash scripts/build_wasm_tests.sh {{ matrix.enable-etdump }}
 
         # Install Jest
         cd cmake-out-wasm/extension/wasm/test

extension/wasm/CMakeLists.txt

@@ -44,6 +44,13 @@ list(
 
 add_library(executorch_wasm OBJECT wasm_bindings.cpp)
 
+if(EXECUTORCH_ENABLE_EVENT_TRACER)
+  list(APPEND link_libraries etdump)
+  target_compile_definitions(
+    executorch_wasm PUBLIC EXECUTORCH_ENABLE_EVENT_TRACER
+  )
+endif()
+
 target_compile_options(executorch_wasm PUBLIC ${_common_compile_options})
 target_include_directories(
   executorch_wasm PUBLIC ${_common_include_directories}

extension/wasm/README.md

@@ -82,6 +82,7 @@ The output `my_project.js` should contain both the emitted JS code and the conte
 - `getMethods()`: Returns the list of methods in the model.
 - `loadMethod(methodName)`: Load a method from the model.
 - `getMethodMetadata(methodName)`: Get the metadata of a method.
+- `etdump()`: If enabled, flushes the etdump buffer and return the results.
 - `execute(methodName, inputs)`: Execute a method with the given inputs.
 - `forward(inputs)`: Execute the forward method with the given inputs.
 - `delete()`: Delete the model from memory.
@@ -118,6 +119,10 @@ The output `my_project.js` should contain both the emitted JS code and the conte
 - `name`: The name of the tensor.
 - These are value types and do not need to be manually deleted.
 
+### ETDumpResult
+- `buffer`: The buffer containing the ETDump data.
+- `delete()`: Delete the ETDumpResult from memory.
+
 ### ScalarType
 - Only `Float` and `Long` are currently supported.
 - `value`: The int constant value of the enum.

extension/wasm/test/CMakeLists.txt

@@ -40,6 +40,27 @@ add_custom_target(
   DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/package.json
 )
 
+if(EXECUTORCH_ENABLE_EVENT_TRACER)
+  set(ETDUMP_UNIT_TESTS ${CMAKE_CURRENT_SOURCE_DIR}/unittests_etdump.js)
+else()
+  set(ETDUMP_UNIT_TESTS
+      ${CMAKE_CURRENT_SOURCE_DIR}/unittests_etdump_disabled.js
+  )
+endif()
+
+add_custom_command(
+  OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/unittests_full.js
+  COMMAND cat ${CMAKE_CURRENT_SOURCE_DIR}/unittests.js ${ETDUMP_UNIT_TESTS} >
+          ${CMAKE_CURRENT_BINARY_DIR}/unittests_full.js
+  DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/unittests.js ${ETDUMP_UNIT_TESTS}
+  COMMENT "Copying unittests_full.js to build output directory"
+)
+
+add_custom_target(
+  executorch_wasm_unittests
+  DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/unittests_full.js
+)
+
 add_executable(executorch_wasm_tests)
 target_link_libraries(executorch_wasm_tests PRIVATE executorch_wasm)
 target_link_options(
@@ -48,7 +69,7 @@ target_link_options(
   --embed-file
   "${MODELS_DIR}@/"
   --pre-js
-  ${CMAKE_CURRENT_SOURCE_DIR}/unittests.js
+  ${CMAKE_CURRENT_BINARY_DIR}/unittests_full.js
   -sASSERTIONS=2
 )
 set_target_properties(
@@ -57,9 +78,9 @@ set_target_properties(
 set_property(
   TARGET executorch_wasm_tests
   APPEND
-  PROPERTY LINK_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/unittests.js
+  PROPERTY LINK_DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/unittests_full.js
 )
 add_dependencies(
-  executorch_wasm_tests executorch_wasm_test_models
+  executorch_wasm_tests executorch_wasm_unittests executorch_wasm_test_models
   executorch_wasm_test_package_json
 )

extension/wasm/test/unittests_etdump.js

@@ -0,0 +1,24 @@
+/*
+ * 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.
+ */
+
+describe("ETDump", () => {
+    test("etdump enabled", () => {
+        const module = et.Module.load("add_mul.pte");
+        const inputs = [et.Tensor.ones([2, 2]), et.Tensor.ones([2, 2]), et.Tensor.ones([2, 2])];
+        const output = module.forward(inputs);
+
+        inputs.forEach((input) => input.delete());
+        output.forEach((output) => output.delete());
+        const etdump = module.etdump();
+        const buffer = etdump.buffer;
+        expect(buffer).toBeInstanceOf(Uint8Array);
+        expect(buffer.length).toBeGreaterThan(0);
+        etdump.delete();
+        module.delete();
+    });
+});

extension/wasm/test/unittests_etdump_disabled.js

@@ -0,0 +1,15 @@
+/*
+ * 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.
+ */
+
+describe("ETDump", () => {
+    test("etdump disabled", () => {
+        const module = et.Module.load("add_mul.pte");
+        expect(() => module.etdump()).toThrow();
+        module.delete();
+    });
+});

extension/wasm/wasm_bindings.cpp

@@ -14,6 +14,10 @@
 #include <executorch/runtime/core/exec_aten/util/tensor_util.h>
 #include <numeric>
 
+#ifdef EXECUTORCH_ENABLE_EVENT_TRACER
+#include <executorch/devtools/etdump/etdump_flatcc.h>
+#endif
+
 #define THROW_JS_ERROR(errorType, message, ...)                           \
   ({                                                                      \
     char msg_buf[256];                                                    \
@@ -51,10 +55,15 @@ using executorch::aten::Tensor;
 using ::executorch::extension::BufferDataLoader;
 using ::executorch::runtime::Error;
 using ::executorch::runtime::EValue;
+using ::executorch::runtime::EventTracer;
 using ::executorch::runtime::Result;
 using ::executorch::runtime::Tag;
 using ::executorch::runtime::TensorInfo;
 
+#ifdef EXECUTORCH_ENABLE_EVENT_TRACER
+using executorch::etdump::ETDumpGen;
+#endif
+
 namespace executorch {
 namespace extension {
 namespace wasm {
@@ -495,6 +504,35 @@ struct ET_EXPERIMENTAL JsMethodMeta {
   }
 };
 
+/**
+ * EXPERIMENTAL: Wrapper around ETDumpResult for JavaScript.
+ */
+#ifdef EXECUTORCH_ENABLE_EVENT_TRACER
+class ET_EXPERIMENTAL JsETDumpResult final {
+ public:
+  JsETDumpResult() = delete;
+  JsETDumpResult(const JsETDumpResult&) = delete;
+  JsETDumpResult& operator=(const JsETDumpResult&) = delete;
+  JsETDumpResult(JsETDumpResult&&) = default;
+  JsETDumpResult& operator=(JsETDumpResult&&) = default;
+
+  explicit JsETDumpResult(uint8_t* buffer, size_t size)
+      : buffer_(buffer), size_(size) {}
+
+  ~JsETDumpResult() {
+    free(buffer_);
+  }
+
+  val get_buffer() const {
+    return val(typed_memory_view(size_, buffer_));
+  }
+
+ private:
+  uint8_t* buffer_;
+  size_t size_;
+};
+#endif
+
 /**
  * EXPERIMENTAL: Wrapper around extension/Module for JavaScript.
  */
@@ -518,17 +556,32 @@ class ET_EXPERIMENTAL JsModule final {
     val memory_view = val(typed_memory_view(length, buffer.data()));
     memory_view.call<void>("set", data);
     auto loader = std::make_unique<BufferDataLoader>(buffer.data(), length);
+
+#ifdef EXECUTORCH_ENABLE_EVENT_TRACER
+    std::unique_ptr<EventTracer> etdump_gen = std::make_unique<ETDumpGen>();
+#else
+    std::unique_ptr<EventTracer> etdump_gen = nullptr;
+#endif
     return std::make_unique<JsModule>(
-        std::move(buffer), std::make_unique<Module>(std::move(loader)));
+        std::move(buffer),
+        std::make_unique<Module>(
+            std::move(loader), nullptr, nullptr, std::move(etdump_gen)));
   }
 
   static std::unique_ptr<JsModule> load(val data) {
     if (data.isNull() || data.isUndefined()) {
       THROW_JS_ERROR(TypeError, "Data cannot be null or undefined");
     }
     if (data.isString()) {
-      return std::make_unique<JsModule>(
-          std::make_unique<Module>(data.as<std::string>()));
+#ifdef EXECUTORCH_ENABLE_EVENT_TRACER
+      std::unique_ptr<EventTracer> etdump_gen = std::make_unique<ETDumpGen>();
+#else
+      std::unique_ptr<EventTracer> etdump_gen = nullptr;
+#endif
+      return std::make_unique<JsModule>(std::make_unique<Module>(
+          data.as<std::string>(),
+          Module::LoadMode::File,
+          std::move(etdump_gen)));
     } else if (data.instanceof (val::global("Uint8Array"))) {
       return load_from_uint8_array(data);
     } else if (data.instanceof (val::global("ArrayBuffer"))) {
@@ -569,6 +622,18 @@ class ET_EXPERIMENTAL JsModule final {
     return JsMethodMeta::from_method_meta(res.get());
   }
 
+#ifdef EXECUTORCH_ENABLE_EVENT_TRACER
+  std::unique_ptr<JsETDumpResult> etdump() {
+    ETDumpGen* etdump_gen = dynamic_cast<ETDumpGen*>(module_->event_tracer());
+    if (etdump_gen == nullptr) {
+      return nullptr;
+    }
+    auto etdump_data = etdump_gen->get_etdump_data();
+    return std::make_unique<JsETDumpResult>(
+        static_cast<uint8_t*>(etdump_data.buf), etdump_data.size);
+  }
+#endif
+
   val_array<val> execute(const std::string& method, val js_inputs) {
     std::vector<EValue> inputs;
     if (js_inputs.isArray()) {
@@ -613,11 +678,19 @@ EMSCRIPTEN_BINDINGS(WasmBindings) {
 #define JS_DECLARE_TAG(NAME) .value(#NAME, Tag::NAME)
       EXECUTORCH_FORALL_TAGS(JS_DECLARE_TAG);
 
+#ifdef EXECUTORCH_ENABLE_EVENT_TRACER
+  class_<JsETDumpResult>("ETDumpResult")
+      .property("buffer", &JsETDumpResult::get_buffer);
+#endif
+
   class_<JsModule>("Module")
       .class_function("load", &JsModule::load)
       .function("getMethods", &JsModule::get_methods)
       .function("loadMethod", &JsModule::load_method)
       .function("getMethodMeta", &JsModule::get_method_meta)
+#ifdef EXECUTORCH_ENABLE_EVENT_TRACER
+      .function("etdump", &JsModule::etdump)
+#endif
       .function("execute", &JsModule::execute)
       .function("forward", &JsModule::forward);
   class_<JsTensor>("Tensor")

scripts/build_wasm_tests.sh

@@ -5,6 +5,15 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+echo "Args: $@"
+for arg in "$@"; do
+  if [ "$arg" == "--enable-etdump" ]; then
+    ETDUMP_OPTS="-DEXECUTORCH_ENABLE_EVENT_TRACER=ON \
+                 -DEXECUTORCH_BUILD_DEVTOOLS=ON \
+                 -DFLATCC_ALLOW_WERROR=OFF"
+  fi
+done
+
 CMAKE_OUT=cmake-out-wasm
 
 cd "$(dirname "${BASH_SOURCE[0]}")/../"
@@ -16,7 +25,7 @@ emcmake cmake . -DEXECUTORCH_BUILD_WASM=ON \
     -DEXECUTORCH_SELECT_OPS_LIST="aten::mm.out,aten::add.out" \
     -DEXECUTORCH_BUILD_TESTS=ON \
     -DCMAKE_BUILD_TYPE=Release \
-    -B"${CMAKE_OUT}"
+    ${ETDUMP_OPTS} -B"${CMAKE_OUT}"
 
 if [ "$(uname)" == "Darwin" ]; then
     CMAKE_JOBS=$(( $(sysctl -n hw.ncpu) - 1 ))