| #!/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. |
| |
| import argparse |
| import codecs |
| import io |
| import os |
| import re |
| import sys |
| |
| parser = argparse.ArgumentParser( |
| description='Utility for re-generating microkernel lists' |
| ) |
| parser.add_argument( |
| '-o', |
| '--output', |
| required=False, |
| help='Directory for generated files.', |
| ) |
| parser.add_argument( |
| '-r', |
| '--root_dir', |
| required=False, |
| help=( |
| 'Start search for files at this directory instead of the script' |
| ' directory.' |
| ), |
| ) |
| |
| TOOLS_DIR = os.path.dirname(os.path.abspath(__file__)) |
| |
| _ISA_LIST = frozenset({ |
| 'armsimd32', |
| 'avx', |
| 'avx2', |
| 'avx512f', |
| 'avx512skx', |
| 'avx512vbmi', |
| 'avx512vnni', |
| 'avx512vnnigfni', |
| 'avx512amx', |
| 'avx512fp16', |
| 'avxvnni', |
| 'avxvnniint8', |
| 'avx256skx', |
| 'avx256vnni', |
| 'avx256vnnigfni', |
| 'f16c', |
| 'fma3', |
| 'fp16arith', |
| 'hexagon', |
| 'hvx', |
| 'neon', |
| 'neonbf16', |
| 'neondot', |
| 'neondotfp16arith', |
| 'neoni8mm', |
| 'neonfma', |
| 'neonfp16', |
| 'neonfp16arith', |
| 'neonsme', |
| 'neonsme2', |
| 'neonv8', |
| 'rvv', |
| 'rvvfp16arith', |
| 'scalar', |
| 'sse', |
| 'sse2', |
| 'sse2fma', |
| 'sse41', |
| 'ssse3', |
| 'wasmblendvps', |
| 'wasmrelaxedsimd', |
| 'wasmpshufb', |
| 'wasmsdot', |
| 'wasmusdot', |
| 'wasmsimd', |
| }) |
| |
| _ISA_MAP = { |
| 'wasmblendvps': 'wasmrelaxedsimd', |
| 'wasmpshufb': 'wasmrelaxedsimd', |
| 'wasmsdot': 'wasmrelaxedsimd', |
| 'wasmusdot': 'wasmrelaxedsimd', |
| } |
| |
| _ARCH_LIST = frozenset({ |
| 'aarch32', |
| 'aarch64', |
| 'wasm32', |
| 'wasmsimd32', |
| 'wasmrelaxedsimd32', |
| 'amd64', |
| }) |
| |
| _MICROKERNEL_NAME_REGEX = re.compile( |
| r'\bxnn_(?:[a-z0-9]+(?:_[a-z0-9]+)*)_ukernel(?:_[a-z0-9]+)*__(?:[a-z0-9]+(?:_[a-z0-9]+)*)\b' |
| ) |
| |
| _VERIFICATION_IGNORE_SUBDIRS = { |
| os.path.join('src', 'qs8-requantization'), |
| os.path.join('src', 'qu8-requantization'), |
| os.path.join('src', 'reference'), |
| os.path.join('src', 'xnnpack', 'simd'), |
| } |
| |
| |
| def overwrite_if_changed(filepath, content): |
| if isinstance(content, io.IOBase): |
| content.seek(0) |
| content = content.read() |
| txt_changed = True |
| if os.path.exists(filepath): |
| with codecs.open(filepath, 'r', encoding='utf-8') as output_file: |
| txt_changed = output_file.read() != content |
| if txt_changed: |
| with codecs.open(filepath, 'w', encoding='utf-8') as output_file: |
| output_file.write(content) |
| |
| |
| def human_sort_key(text): |
| return [ |
| token.zfill(10) if token.isdigit() else token.lower() |
| for token in re.split(r'(\d+|\W)', text) |
| if token |
| ] |
| |
| |
| def make_variable_name(prefix, key, suffix): |
| return '_'.join(token for token in [prefix, key.upper(), suffix] if token) |
| |
| |
| def write_grouped_microkernels_bzl(file, key, microkernels, prefix, suffix): |
| if key not in microkernels: |
| return [] |
| variable_name = make_variable_name(prefix, key, suffix) |
| file.write(f'\n{variable_name} = [\n') |
| for microkernel in sorted(microkernels[key], key=human_sort_key): |
| file.write(f' "{microkernel}",\n') |
| file.write(']\n') |
| return [variable_name] |
| |
| |
| def write_grouped_microkernels_cmake(file, key, microkernels, prefix, suffix): |
| if key not in microkernels: |
| return [] |
| variable_name = make_variable_name(prefix, key, suffix) |
| file.write(f'\nSET({variable_name}') |
| for microkernel in sorted(microkernels[key], key=human_sort_key): |
| file.write(f'\n {microkernel}') |
| file.write(')\n') |
| return [variable_name] |
| |
| |
| def main(args): |
| options = parser.parse_args(args) |
| root_dir = ( |
| options.root_dir |
| if options.root_dir |
| else os.path.normpath(os.path.join(TOOLS_DIR, '..')) |
| ) |
| src_dir = os.path.join(root_dir, 'src') |
| configs_dir = os.path.join(src_dir, 'configs') |
| dst_dir = options.output if options.output else root_dir |
| bzl_gen_dir = os.path.join(dst_dir, 'gen') |
| if not os.path.exists(bzl_gen_dir): |
| os.makedirs(bzl_gen_dir, exist_ok=True) |
| cmake_gen_dir = os.path.join(dst_dir, 'cmake', 'gen') |
| if not os.path.exists(cmake_gen_dir): |
| os.makedirs(cmake_gen_dir, exist_ok=True) |
| ignore_roots = { |
| src_dir, |
| os.path.join(src_dir, 'configs'), |
| os.path.join(src_dir, 'enums'), |
| os.path.join(src_dir, 'operators'), |
| os.path.join(src_dir, 'subgraph'), |
| os.path.join(src_dir, 'tables'), |
| os.path.join(src_dir, 'xnnpack'), |
| } |
| c_microkernels_per_isa = {isa: [] for isa in _ISA_LIST if isa not in _ISA_MAP} |
| c_microkernels_per_isa['neon_aarch64'] = list() |
| c_microkernels_per_isa['neondot_aarch64'] = list() |
| c_microkernels_per_isa['neonfma_aarch64'] = list() |
| c_microkernels_per_isa['neonfp16arith_aarch64'] = list() |
| c_microkernels_per_isa['neonsme'] = list() |
| c_microkernels_per_isa['neonsme2'] = list() |
| asm_microkernels_per_arch = {arch: [] for arch in _ARCH_LIST} |
| microkernel_name_to_filename = dict() |
| for root, _, files in os.walk(src_dir, topdown=False): |
| if root in ignore_roots: |
| continue |
| |
| for name in files: |
| if name.endswith('.in'): |
| continue |
| if name.endswith('.h'): |
| continue |
| if name.endswith('.inc'): |
| continue |
| basename, ext = os.path.splitext(name) |
| if ext == '.sollya': |
| continue |
| |
| subdir = os.path.relpath(root, root_dir) |
| filepath = os.path.join(subdir, name) |
| if 'pipertmp' in filepath: |
| continue |
| |
| # Skip files created by repository tools. |
| if ( |
| name.startswith('._') |
| or filepath.endswith('.swp') |
| or filepath.endswith('.orig') |
| ): |
| continue |
| |
| # Build microkernel name -> microkernel filepath mapping |
| with open(os.path.join(root_dir, filepath), 'r', encoding='utf-8') as f: |
| content = f.read() |
| microkernels = re.findall(_MICROKERNEL_NAME_REGEX, content) |
| if not microkernels: |
| if subdir in _VERIFICATION_IGNORE_SUBDIRS: |
| microkernel_name_to_filename[ |
| os.path.splitext(os.path.basename(filepath))[0] |
| ] = filepath |
| else: |
| print('No microkernel found in %s' % filepath) |
| else: |
| microkernels = sorted(set(microkernels)) |
| if len(microkernels) > 1: |
| if ext == '.cc': |
| # In generators it is normal to have multiple implementations in |
| # the same file. We need to filter out non-generators though. |
| microkernels = sorted( |
| filter( |
| lambda fn: fn.startswith('xnn_generate_'), microkernels |
| ) |
| ) |
| if ( |
| not microkernels |
| and subdir not in _VERIFICATION_IGNORE_SUBDIRS |
| ): |
| print('No microkernel generator found in %s' % filepath) |
| |
| for microkernel in microkernels: |
| if microkernel in microkernel_name_to_filename: |
| print( |
| 'Duplicate microkernel generator definition: %s and %s' |
| % (microkernel_name_to_filename[microkernel], filepath) |
| ) |
| else: |
| microkernel_name_to_filename[microkernel] = filepath |
| else: |
| # Extract the individual function names and their offsets. |
| matches = list(re.finditer(_MICROKERNEL_NAME_REGEX, content)) |
| |
| # Write them to temporary files. |
| for k, match in enumerate(matches): |
| microkernel = match.group(0) |
| if microkernel in microkernel_name_to_filename: |
| print( |
| 'Duplicate microkernel definition: %s and %s (%ith' |
| ' function)' |
| % ( |
| microkernel_name_to_filename[microkernel], |
| filepath, |
| k, |
| ) |
| ) |
| microkernel_name_to_filename[microkernel] = filepath |
| |
| else: |
| microkernel = microkernels[0] |
| if microkernel in microkernel_name_to_filename: |
| print( |
| 'Duplicate microkernel definition: %s and %s' |
| % (microkernel_name_to_filename[microkernel], filepath) |
| ) |
| else: |
| microkernel_name_to_filename[microkernel] = filepath |
| |
| if ext == '.c': |
| arch = None |
| for component in basename.split('-'): |
| if component in _ARCH_LIST: |
| arch = component |
| if component in _ISA_LIST: |
| isa = _ISA_MAP.get(component, component) |
| key = isa if arch is None else f'{isa}_{arch}' |
| c_microkernels_per_isa[key].append(filepath) |
| break |
| else: |
| print('Unknown ISA for C microkernel %s' % filepath) |
| elif ext == '.S': |
| for component in basename.split('-'): |
| if component in _ARCH_LIST: |
| asm_microkernels_per_arch[component].append(filepath) |
| break |
| else: |
| print('Unknown architecture for assembly microkernel %s' % filepath) |
| |
| # Collect filenames of production microkernels as a set |
| prod_microkernels = set() |
| for configs_filepath in os.listdir(configs_dir): |
| if not configs_filepath.endswith('-config.c'): |
| continue |
| with open( |
| os.path.join(configs_dir, configs_filepath), 'r', encoding='utf-8' |
| ) as config_file: |
| content = config_file.read() |
| microkernels = re.findall(_MICROKERNEL_NAME_REGEX, content) |
| prod_microkernels.update(microkernels) |
| prod_microkernels = set( |
| map(microkernel_name_to_filename.get, prod_microkernels) |
| ) |
| |
| # Split the microkernels into prod/test. |
| prod_c_microkernels_per_isa = dict() |
| non_prod_c_microkernels_per_isa = dict() |
| for arch, microkernels in c_microkernels_per_isa.items(): |
| prod_c_microkernels_per_isa[arch] = [ |
| v for v in microkernels if v in prod_microkernels |
| ] |
| non_prod_c_microkernels_per_isa[arch] = [ |
| v for v in microkernels if v not in prod_microkernels |
| ] |
| prod_asm_microkernels_per_arch = dict() |
| non_prod_asm_microkernels_per_arch = dict() |
| for arch, microkernels in asm_microkernels_per_arch.items(): |
| prod_asm_microkernels_per_arch[arch] = [ |
| v for v in microkernels if v in prod_microkernels |
| ] |
| non_prod_asm_microkernels_per_arch[arch] = [ |
| v for v in microkernels if v not in prod_microkernels |
| ] |
| |
| with io.StringIO() as microkernels_bzl: |
| microkernels_bzl.write('''\ |
| # |
| # Microkernel filenames lists. |
| # |
| # Auto-generated file. Do not edit! |
| # Generator: tools/update-microkernels.py |
| # |
| |
| ''') |
| prod_c_vars_per_arch = dict() |
| non_prod_c_vars_per_arch = dict() |
| prod_asm_vars_per_arch = dict() |
| non_prod_asm_vars_per_arch = dict() |
| keys = set(c_microkernels_per_isa.keys()).union( |
| asm_microkernels_per_arch.keys() |
| ) |
| keys = sorted(keys, key=lambda key: key + '_microkernels.bzl') |
| exports = ['\n'] |
| for key in keys: |
| arch_microkernels_bzl_filename = key + '_microkernels.bzl' |
| with io.StringIO() as arch_microkernels_bzl: |
| arch_microkernels_bzl.write(f'''\ |
| # |
| # Microkernel filenames lists for {key}. |
| # |
| # Auto-generated file. Do not edit! |
| # Generator: tools/update-microkernels.py |
| # |
| ''') |
| prod_c_vars = write_grouped_microkernels_bzl( |
| arch_microkernels_bzl, |
| key, |
| prod_c_microkernels_per_isa, |
| 'PROD', |
| 'MICROKERNEL_SRCS', |
| ) |
| if prod_c_vars: |
| prod_c_vars_per_arch[key] = prod_c_vars |
| non_prod_c_vars = write_grouped_microkernels_bzl( |
| arch_microkernels_bzl, |
| key, |
| non_prod_c_microkernels_per_isa, |
| 'NON_PROD', |
| 'MICROKERNEL_SRCS', |
| ) |
| if non_prod_c_vars: |
| non_prod_c_vars_per_arch[key] = non_prod_c_vars |
| c_vars = prod_c_vars + non_prod_c_vars |
| if c_vars: |
| all_srcs = make_variable_name('ALL', key, 'MICROKERNEL_SRCS') |
| arch_microkernels_bzl.write(f'\n{all_srcs} = {" + ".join(c_vars)}\n') |
| c_vars.append(all_srcs) |
| prod_asm_vars = write_grouped_microkernels_bzl( |
| arch_microkernels_bzl, |
| key, |
| prod_asm_microkernels_per_arch, |
| 'PROD', |
| 'ASM_MICROKERNEL_SRCS', |
| ) |
| if prod_asm_vars: |
| prod_asm_vars_per_arch[key] = prod_asm_vars |
| non_prod_asm_vars = write_grouped_microkernels_bzl( |
| arch_microkernels_bzl, |
| key, |
| non_prod_asm_microkernels_per_arch, |
| 'NON_PROD', |
| 'ASM_MICROKERNEL_SRCS', |
| ) |
| if non_prod_asm_vars: |
| non_prod_asm_vars_per_arch[key] = non_prod_asm_vars |
| asm_vars = prod_asm_vars + non_prod_asm_vars |
| if asm_vars: |
| all_srcs = make_variable_name('', key, 'ASM_MICROKERNEL_SRCS') |
| arch_microkernels_bzl.write( |
| f'\n{all_srcs} = {" + ".join(asm_vars)}\n' |
| ) |
| asm_vars.append(all_srcs) |
| all_vars = c_vars + asm_vars |
| imports = ', '.join(f'_{var} = "{var}"' for var in sorted(all_vars)) |
| microkernels_bzl.write( |
| f'load("{arch_microkernels_bzl_filename}", {imports})\n' |
| ) |
| for var in all_vars: |
| exports.append(f'{var} = _{var}\n') |
| |
| overwrite_if_changed( |
| os.path.join(bzl_gen_dir, arch_microkernels_bzl_filename), |
| arch_microkernels_bzl, |
| ) |
| |
| # Generate dictionaries of microkernel lists per arch. |
| microkernels_bzl.write(''.join(sorted(exports))) |
| microkernels_bzl.write('\nPROD_C_SRCS_FOR_ARCH = {\n') |
| for key, var_list in prod_c_vars_per_arch.items(): |
| microkernels_bzl.write(f' "{key}": {" + ".join(var_list)},\n') |
| microkernels_bzl.write('}\n') |
| microkernels_bzl.write('\nNON_PROD_C_SRCS_FOR_ARCH = {\n') |
| for key, var_list in non_prod_c_vars_per_arch.items(): |
| microkernels_bzl.write(f' "{key}": {" + ".join(var_list)},\n') |
| microkernels_bzl.write('}\n') |
| microkernels_bzl.write('\nPROD_ASM_SRCS_FOR_ARCH = {\n') |
| for key, var_list in prod_asm_vars_per_arch.items(): |
| microkernels_bzl.write(f' "{key}": {" + ".join(var_list)},\n') |
| microkernels_bzl.write('}\n') |
| microkernels_bzl.write('\nNON_PROD_ASM_SRCS_FOR_ARCH = {\n') |
| for key, var_list in non_prod_asm_vars_per_arch.items(): |
| microkernels_bzl.write(f' "{key}": {" + ".join(var_list)},\n') |
| microkernels_bzl.write('}\n') |
| microkernels_bzl.write(""" |
| def prod_c_srcs_for_arch(arch): |
| return PROD_C_SRCS_FOR_ARCH.get(arch, []) |
| |
| def non_prod_c_srcs_for_arch(arch): |
| return NON_PROD_C_SRCS_FOR_ARCH.get(arch, []) |
| |
| def all_c_srcs_for_arch(arch): |
| return prod_c_srcs_for_arch(arch) + non_prod_c_srcs_for_arch(arch) |
| |
| def prod_asm_srcs_for_arch(arch): |
| return PROD_ASM_SRCS_FOR_ARCH.get(arch, []) |
| |
| def non_prod_asm_srcs_for_arch(arch): |
| return NON_PROD_ASM_SRCS_FOR_ARCH.get(arch, []) |
| |
| def all_asm_srcs_for_arch(arch): |
| return prod_asm_srcs_for_arch(arch) + non_prod_asm_srcs_for_arch(arch) |
| |
| def prod_srcs_for_arch(arch): |
| return prod_c_srcs_for_arch(arch) + prod_asm_srcs_for_arch(arch) |
| |
| def non_prod_srcs_for_arch(arch): |
| return non_prod_c_srcs_for_arch(arch) + non_prod_asm_srcs_for_arch(arch) |
| |
| def all_srcs_for_arch(arch): |
| return all_c_srcs_for_arch(arch) + all_asm_srcs_for_arch(arch) |
| """) |
| |
| overwrite_if_changed( |
| os.path.join(bzl_gen_dir, 'microkernels.bzl'), microkernels_bzl |
| ) |
| |
| with io.StringIO() as microkernels_cmake: |
| microkernels_cmake.write("""\ |
| # 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. |
| # |
| # Description: microkernel filename lists. |
| # |
| # Auto-generated file. Do not edit! |
| # Generator: tools/update-microkernels.py |
| |
| """) |
| keys = sorted( |
| set(c_microkernels_per_isa.keys()).union( |
| asm_microkernels_per_arch.keys() |
| ) |
| ) |
| for key in keys: |
| arch_microkernels_cmake_filename = key + '_microkernels.cmake' |
| with io.StringIO() as arch_microkernels_cmake: |
| arch_microkernels_cmake.write(f"""\ |
| # 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. |
| # |
| # Description: microkernel filename lists for {key} |
| # |
| # Auto-generated file. Do not edit! |
| # Generator: tools/update-microkernels.py |
| |
| """) |
| c_vars = write_grouped_microkernels_cmake( |
| arch_microkernels_cmake, |
| key, |
| prod_c_microkernels_per_isa, |
| 'PROD', |
| 'MICROKERNEL_SRCS', |
| ) |
| c_vars += write_grouped_microkernels_cmake( |
| arch_microkernels_cmake, |
| key, |
| non_prod_c_microkernels_per_isa, |
| 'NON_PROD', |
| 'MICROKERNEL_SRCS', |
| ) |
| asm_vars = write_grouped_microkernels_cmake( |
| arch_microkernels_cmake, |
| key, |
| prod_asm_microkernels_per_arch, |
| 'PROD', |
| 'ASM_MICROKERNEL_SRCS', |
| ) |
| asm_vars += write_grouped_microkernels_cmake( |
| arch_microkernels_cmake, |
| key, |
| non_prod_asm_microkernels_per_arch, |
| 'NON_PROD', |
| 'ASM_MICROKERNEL_SRCS', |
| ) |
| if c_vars: |
| c_vars = [f'${{{v}}}' for v in c_vars] |
| arch_microkernels_cmake.write( |
| f'\nSET({make_variable_name("ALL", key, "MICROKERNEL_SRCS")} {" + ".join(c_vars)})\n' |
| ) |
| if asm_vars: |
| asm_vars = [f'${{{v}}}' for v in asm_vars] |
| arch_microkernels_cmake.write( |
| f'\nSET({make_variable_name("ALL", key, "ASM_MICROKERNEL_SRCS")} {" + ".join(asm_vars)})\n' |
| ) |
| microkernels_cmake.write( |
| f'INCLUDE(cmake/gen/{arch_microkernels_cmake_filename})\n' |
| ) |
| overwrite_if_changed( |
| os.path.join(cmake_gen_dir, arch_microkernels_cmake_filename), |
| arch_microkernels_cmake, |
| ) |
| |
| overwrite_if_changed( |
| os.path.join(cmake_gen_dir, 'microkernels.cmake'), microkernels_cmake |
| ) |
| |
| |
| if __name__ == '__main__': |
| main(sys.argv[1:]) |
