|
5 | 5 | # LICENSE file in the root directory of this source tree. |
6 | 6 |
|
7 | 7 | import json |
| 8 | + |
| 9 | +import logging |
8 | 10 | import os |
9 | 11 | import tempfile |
10 | 12 |
|
11 | 13 | from dataclasses import dataclass, fields, is_dataclass |
12 | | -from typing import ClassVar, Literal |
| 14 | +from typing import ClassVar, Literal, Optional |
13 | 15 |
|
14 | 16 | import pkg_resources |
15 | 17 | from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import XNNGraph |
16 | 18 | from executorch.exir._serialize._dataclass import _DataclassEncoder |
17 | 19 |
|
18 | 20 | from executorch.exir._serialize._flatbuffer import _flatc_compile |
19 | 21 |
|
| 22 | +logger = logging.getLogger(__name__) |
| 23 | +logger.setLevel(logging.WARNING) |
| 24 | + |
20 | 25 | # Byte order of numbers written to program headers. Always little-endian |
21 | 26 | # regardless of the host system, since all commonly-used modern CPUs are little |
22 | 27 | # endian. |
@@ -273,19 +278,42 @@ def _pad_to(data: bytes, length: int) -> bytes: |
273 | 278 | return data |
274 | 279 |
|
275 | 280 |
|
276 | | -def pretty_print_xnngraph(xnnpack_graph_json: str): |
| 281 | +def pretty_print_xnngraph(xnnpack_graph_json: str, filename: Optional[str] = None): |
277 | 282 | """ |
278 | | - Pretty print the XNNGraph |
| 283 | + Pretty print the XNNGraph, optionally writing to a file if filename is provided |
279 | 284 | """ |
280 | | - from pprint import pprint |
| 285 | + from pprint import pformat |
281 | 286 |
|
282 | 287 | d = json.loads(xnnpack_graph_json) |
283 | | - pprint(d) |
| 288 | + pstr = pformat(d, indent=2, compact=True).replace("'", '"') |
| 289 | + if filename: |
| 290 | + with open(filename, "w") as f: |
| 291 | + if filename.endswith(".json"): |
| 292 | + pstr = pstr.replace("None", "null") |
| 293 | + f.write(pstr) |
| 294 | + else: # dump to stdout |
| 295 | + print("XNNGraph:") |
| 296 | + print(pstr) |
| 297 | + print("End of XNNGraph") |
| 298 | + |
| 299 | + |
| 300 | +# TODO: Replace this with an actual delegate id |
| 301 | +_delegate_instance_id = 0 |
284 | 302 |
|
285 | 303 |
|
286 | 304 | def convert_to_flatbuffer(xnnpack_graph: XNNGraph) -> bytes: |
| 305 | + global _delegate_instance_id |
287 | 306 | sanity_check_xnngraph_dataclass(xnnpack_graph) |
288 | 307 | xnnpack_graph_json = json.dumps(xnnpack_graph, cls=_DataclassEncoder) |
| 308 | + |
| 309 | + # Log the XNNGraph if debugging |
| 310 | + if logger.getEffectiveLevel() == logging.DEBUG: |
| 311 | + filename: str = f"./xnnpack_delegate_graph_{_delegate_instance_id}.json" |
| 312 | + logger.debug(f"Writing XNNGraph to {filename}") |
| 313 | + pretty_print_xnngraph(xnnpack_graph_json, filename) |
| 314 | + |
| 315 | + _delegate_instance_id += 1 |
| 316 | + |
289 | 317 | with tempfile.TemporaryDirectory() as d: |
290 | 318 | schema_path = os.path.join(d, "schema.fbs") |
291 | 319 | with open(schema_path, "wb") as schema_file: |
|
0 commit comments