|
10 | 10 | #include <gflags/gflags.h> |
11 | 11 | #include <sstream> |
12 | 12 | #include <vector> |
13 | | - |
14 | 13 | #include <executorch/examples/models/llama/runner/runner.h> |
15 | 14 |
|
| 15 | +#ifdef ET_EVENT_TRACER_ENABLED |
| 16 | +#include <executorch/devtools/etdump/etdump_flatcc.h> |
| 17 | +#endif |
| 18 | + |
16 | 19 | #if defined(ET_USE_THREADPOOL) |
17 | 20 | #include <executorch/extension/threadpool/cpuinfo_utils.h> |
18 | 21 | #include <executorch/extension/threadpool/threadpool.h> |
@@ -64,6 +67,11 @@ DEFINE_int32( |
64 | 67 |
|
65 | 68 | DEFINE_bool(warmup, false, "Whether to run a warmup run."); |
66 | 69 |
|
| 70 | +DEFINE_string( |
| 71 | + etdump_path, |
| 72 | + "etdump.in", |
| 73 | + "If an etdump path is provided, generate an ETDump file at the specified path for profiling purposes."); |
| 74 | + |
67 | 75 | // Helper function to parse comma-separated string lists |
68 | 76 | std::vector<std::string> parseStringList(const std::string& input) { |
69 | 77 | std::vector<std::string> result; |
@@ -117,9 +125,26 @@ int32_t main(int32_t argc, char** argv) { |
117 | 125 | ->_unsafe_reset_threadpool(num_performant_cores); |
118 | 126 | } |
119 | 127 | #endif |
| 128 | + |
| 129 | +#ifdef ET_EVENT_TRACER_ENABLED |
| 130 | + // Create ETDumpGen and get raw pointer reference for later access |
| 131 | + auto etdump_gen_ptr = std::make_unique<executorch::etdump::ETDumpGen>(); |
| 132 | + executorch::etdump::ETDumpGen* etdump_gen = etdump_gen_ptr.get(); |
| 133 | +#endif |
| 134 | + |
120 | 135 | // create llama runner |
121 | 136 | std::unique_ptr<::executorch::extension::llm::TextLLMRunner> runner = |
122 | | - example::create_llama_runner(model_path, tokenizer_path, data_paths); |
| 137 | + example::create_llama_runner( |
| 138 | + model_path, |
| 139 | + tokenizer_path, |
| 140 | + data_paths, |
| 141 | + temperature, |
| 142 | +#ifdef ET_EVENT_TRACER_ENABLED |
| 143 | + std::move(etdump_gen_ptr) |
| 144 | +#else |
| 145 | + nullptr |
| 146 | +#endif |
| 147 | + ); |
123 | 148 |
|
124 | 149 | if (runner == nullptr) { |
125 | 150 | ET_LOG(Error, "Failed to create llama runner"); |
@@ -157,5 +182,22 @@ int32_t main(int32_t argc, char** argv) { |
157 | 182 | return 1; |
158 | 183 | } |
159 | 184 |
|
| 185 | +#ifdef ET_EVENT_TRACER_ENABLED |
| 186 | + if (etdump_gen != nullptr) { |
| 187 | + executorch::etdump::ETDumpResult result = etdump_gen->get_etdump_data(); |
| 188 | + if (result.buf != nullptr && result.size > 0) { |
| 189 | + FILE* f = fopen(FLAGS_etdump_path.c_str(), "w+"); |
| 190 | + if (f == nullptr) { |
| 191 | + ET_LOG(Error, "Failed to open etdump file at path: %s", FLAGS_etdump_path.c_str()); |
| 192 | + } else { |
| 193 | + fwrite((uint8_t*)result.buf, 1, result.size, f); |
| 194 | + fclose(f); |
| 195 | + ET_LOG(Info, "ETDump file written to: %s", FLAGS_etdump_path.c_str()); |
| 196 | + } |
| 197 | + free(result.buf); |
| 198 | + } |
| 199 | + } |
| 200 | +#endif |
| 201 | + |
160 | 202 | return 0; |
161 | 203 | } |
0 commit comments