xref: /aosp_15_r20/external/XNNPACK/tools/check-assembly-and-jit-microkernels-in-sync.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
7from pathlib import Path
8import subprocess
9
10
11# Check if `line` can be found in file.
12def line_in_file(line, filename):
13  with open(filename) as o:
14    return any(line in l for l in o)
15
16
17# Finds all the JIT generated microkernel cc files.
18# We could use the JIT_AARCH32_SRCS and JIT_AARCH64_SRCS in BUILD too.
19def find_jit_files():
20  paths = list(Path('src').glob('**/*.cc'))
21  return [p for p in paths if line_in_file('Converted from', p)]
22
23
24# Check that file exists.
25def check_assembly_file_exists(filename):
26  p = Path('.')
27  assert((p / filename).exists())
28
29
30# JIT microkernel can be converted from a assembly file or a templated assembly file, gets the right "root" file.
31def get_assembly_file_or_template(converted_from):
32  file = converted_from
33  with open(converted_from) as f:
34    for line in f.readlines():
35      if 'Template: ' in line:
36        file = line.split()[2]
37        break
38  if file.startswith('/'):
39    file = file[len('/'):]
40  return file
41
42
43# Finds the assembly file associated with the JIT file to write the LINT marker to.
44def find_assembly_file(jit_file):
45  with jit_file.open() as f:
46    for line in f.readlines():
47      if 'Converted from' in line:
48        converted_from = line.split()[3]
49        check_assembly_file_exists(converted_from)
50        assembly_file = get_assembly_file_or_template(converted_from)
51        return assembly_file
52    return f'{jit_file} does not have converted from'
53
54
55# Writes LINT markers to `assembly_file`.
56def write_lint_markers(assembly_file, jit_file):
57  with open(assembly_file) as f:
58    lines = f.readlines()
59    if any('LINT.IfChange' in l for l in lines):
60      # Has lint marker, check that the JIT file is found.
61      if not any(jit_file.name in l for l in lines):
62        print(f'{jit_file.name} not found in {assembly_file}')
63        assert(False)
64      return
65
66  relative_jit_file = jit_file.name
67  # This sed script will write LINT markers that wraps the function definition in the assembly file.
68  sed_args = [
69      'sed',
70      '-i',
71      '-e',
72      '/# void xnn/i# LINT.IfChange',
73      '-e',
74      '/\/\/ void xnn/i# LINT.IfChange',
75      '-e',
76      f'/END_FUNCTION/a# LINT.ThenChange({relative_jit_file})',
77      assembly_file,
78  ]
79  subprocess.run(sed_args)
80
81
82def main():
83  for f in find_jit_files():
84    write_lint_markers(find_assembly_file(f), f)
85
86
87if __name__ == '__main__':
88  main()
89