xref: /aosp_15_r20/external/XNNPACK/tools/dump-jit-output.py (revision 4bdc94577ba0e567308109d787f7fec7b531ce36) 
1#!/usr/bin/env python
2# Copyright 2022 Google LLC
3#
4# This source code is licensed under the BSD-style license found in the
5# LICENSE file in the root directory of this source tree.
6
7# Dumps JIT codegen given specific JIT parameters.
8# Usage is similar to xngen:
9#   dump-jit-output.py <path to JIT cc file> --max_mr=6 [--clamp_min]
10# E.g.
11#   ./tools/dump-jit-output.py \
12#     src/f32-gemm/upto6x8-aarch64-neonfma-cortex-a75.cc
13#     --max_mr=6
14#
15# The parameters prefetch, clamp_min, clamp_max defaults to True if not
16# specified on the command line.
17
18import argparse
19import codecs
20import re
21import sys
22import xngen
23from itertools import chain
24
25
26parser = argparse.ArgumentParser(description='Dump output of JIT')
27parser.add_argument("input", metavar="FILE", nargs=1,
28          help="Input file")
29parser.add_argument("--prefetch", action="store_true")
30parser.add_argument("--clamp_min", action="store_true")
31parser.add_argument("--clamp_max", action="store_true")
32parser.add_argument("--max_mr", type=int, required=True)
33parser.add_argument("-o", "--output",
34          help='Output file')
35parser.set_defaults(defines=list())
36
37
38def preprocess(input_text):
39  input_lines = input_text.splitlines()
40  in_function = False
41  output = []
42  for i, line in enumerate(input_lines):
43    if line.startswith('void Generator::generate'):
44      in_function = True
45    if not in_function:
46      continue
47    if line == '}':
48      in_function = False
49      output.append(line)
50      continue
51    if line.strip() == '}':
52      continue
53
54    o = re.sub(r'(if|else)( +\(.*\)) +{', r'$\1\2:', line)
55    o = re.sub(r'&&', 'and', o)
56    o = re.sub(r'\|\|', 'or', o)
57    output.append(o)
58  return output
59
60
61def call_xngen(text, options):
62  input_globals = {
63      'prefetch': options.prefetch,
64      'clamp_min': options.clamp_min,
65      'clamp_max': options.clamp_max,
66      'max_mr': options.max_mr,
67  }
68  return xngen.preprocess(text, input_globals, "codegen")
69
70
71def main(args):
72  options = parser.parse_args(args)
73  input_text = codecs.open(options.input[0], "r", encoding="utf-8").read()
74  output = preprocess(input_text)
75  result = call_xngen("\n".join(output), options)
76  if (options.output):
77    with codecs.open(options.output, "w", encoding="utf-8") as output_file:
78      output_file.write(result)
79  else:
80    print(result)
81
82
83if __name__ == "__main__":
84  main(sys.argv[1:])
85