|
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