File size: 2,553 Bytes
8b7c501
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
#!/usr/bin/env python
# Copyright 2022 Google LLC
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

from pathlib import Path
import subprocess


# Check if `line` can be found in file.
def line_in_file(line, filename):
  with open(filename) as o:
    return any(line in l for l in o)


# Finds all the JIT generated microkernel cc files.
# We could use the JIT_AARCH32_SRCS and JIT_AARCH64_SRCS in BUILD too.
def find_jit_files():
  paths = list(Path('src').glob('**/*.cc'))
  return [p for p in paths if line_in_file('Converted from', p)]


# Check that file exists.
def check_assembly_file_exists(filename):
  p = Path('.')
  assert((p / filename).exists())


# JIT microkernel can be converted from a assembly file or a templated assembly file, gets the right "root" file.
def get_assembly_file_or_template(converted_from):
  file = converted_from
  with open(converted_from) as f:
    for line in f.readlines():
      if 'Template: ' in line:
        file = line.split()[2]
        break
  if file.startswith('/'):
    file = file[len('/'):]
  return file


# Finds the assembly file associated with the JIT file to write the LINT marker to.
def find_assembly_file(jit_file):
  with jit_file.open() as f:
    for line in f.readlines():
      if 'Converted from' in line:
        converted_from = line.split()[3]
        check_assembly_file_exists(converted_from)
        assembly_file = get_assembly_file_or_template(converted_from)
        return assembly_file
    return f'{jit_file} does not have converted from'


# Writes LINT markers to `assembly_file`.
def write_lint_markers(assembly_file, jit_file):
  with open(assembly_file) as f:
    lines = f.readlines()
    if any('LINT.IfChange' in l for l in lines):
      # Has lint marker, check that the JIT file is found.
      if not any(jit_file.name in l for l in lines):
        print(f'{jit_file.name} not found in {assembly_file}')
        assert(False)
      return

  relative_jit_file = jit_file.name
  # This sed script will write LINT markers that wraps the function definition in the assembly file.
  sed_args = [
      'sed',
      '-i',
      '-e',
      '/# void xnn/i# LINT.IfChange',
      '-e',
      '/\/\/ void xnn/i# LINT.IfChange',
      '-e',
      f'/END_FUNCTION/a# LINT.ThenChange({relative_jit_file})',
      assembly_file,
  ]
  subprocess.run(sed_args)


def main():
  for f in find_jit_files():
    write_lint_markers(find_assembly_file(f), f)


if __name__ == '__main__':
  main()