mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
d92d4133e7
> **__Note:__** XNNPACK Upgrade is too large in the range of **40k** files and **10m** Lines of code, Thus we break the update of the library into multiple parts. All Parts [1 - n] Must be landed together for it to work. ***This also means If there is a revert. Please revert the Entire Stack.*** This change is everything remaining requiring XNNPACK version to work. @allow-large-files Differential Revision: [D52099769](https://our.internmc.facebook.com/intern/diff/D52099769/) --- submodule (unblock merge to make ShipIt happy) Pull Request resolved: https://github.com/pytorch/pytorch/pull/115714 Approved by: https://github.com/digantdesai
197 lines
8.8 KiB
Python
197 lines
8.8 KiB
Python
#!/usr/bin/env python3
|
|
|
|
from __future__ import print_function
|
|
import collections
|
|
import os
|
|
import sys
|
|
import logging
|
|
|
|
BANNER = "Auto-generated by generate-wrappers.py script. Do not modify"
|
|
WRAPPER_SRC_NAMES = {
|
|
"PROD_SCALAR_MICROKERNEL_SRCS": None,
|
|
"PROD_FMA_MICROKERNEL_SRCS": "defined(__riscv) || defined(__riscv__)",
|
|
"PROD_ARMSIMD32_MICROKERNEL_SRCS": "defined(__arm__)",
|
|
"PROD_FP16ARITH_MICROKERNEL_SRCS": "defined(__arm__)",
|
|
"PROD_NEON_MICROKERNEL_SRCS": "defined(__arm__) || defined(__aarch64__)",
|
|
"PROD_NEONFP16_MICROKERNEL_SRCS": "defined(__arm__) || defined(__aarch64__)",
|
|
"PROD_NEONFMA_MICROKERNEL_SRCS": "defined(__arm__) || defined(__aarch64__)",
|
|
"PROD_NEON_AARCH64_MICROKERNEL_SRCS": "defined(__aarch64__)",
|
|
"PROD_NEONV8_MICROKERNEL_SRCS": "defined(__arm__) || defined(__aarch64__)",
|
|
"PROD_NEONFP16ARITH_MICROKERNEL_SRCS": "defined(__arm__) || defined(__aarch64__)",
|
|
"PROD_NEONFP16ARITH_AARCH64_MICROKERNEL_SRCS": "defined(__aarch64__)",
|
|
"PROD_NEONDOT_MICROKERNEL_SRCS": "defined(__arm__) || defined(__aarch64__)",
|
|
"PROD_NEONDOT_AARCH64_MICROKERNEL_SRCS": "defined(__aarch64__)",
|
|
"PROD_NEONDOTFP16ARITH_MICROKERNEL_SRCS": "defined(__arm__) || defined(__aarch64__)",
|
|
"PROD_NEONDOTFP16ARITH_AARCH64_MICROKERNEL_SRCS": "defined(__aarch64__)",
|
|
"PROD_NEONI8MM_MICROKERNEL_SRCS": "defined(__arm__) || defined(__aarch64__)",
|
|
"PROD_SSE_MICROKERNEL_SRCS": "defined(__i386__) || defined(__i686__) || defined(__x86_64__)",
|
|
"PROD_SSE2_MICROKERNEL_SRCS": "defined(__i386__) || defined(__i686__) || defined(__x86_64__)",
|
|
"PROD_SSSE3_MICROKERNEL_SRCS": "defined(__i386__) || defined(__i686__) || defined(__x86_64__)",
|
|
"PROD_SSE41_MICROKERNEL_SRCS": "defined(__i386__) || defined(__i686__) || defined(__x86_64__)",
|
|
"PROD_AVX_MICROKERNEL_SRCS": "defined(__i386__) || defined(__i686__) || defined(__x86_64__)",
|
|
"PROD_F16C_MICROKERNEL_SRCS": "defined(__i386__) || defined(__i686__) || defined(__x86_64__)",
|
|
"PROD_XOP_MICROKERNEL_SRCS": "defined(__i386__) || defined(__i686__) || defined(__x86_64__)",
|
|
"PROD_FMA3_MICROKERNEL_SRCS": "defined(__i386__) || defined(__i686__) || defined(__x86_64__)",
|
|
"PROD_AVX2_MICROKERNEL_SRCS": "defined(__i386__) || defined(__i686__) || defined(__x86_64__)",
|
|
"PROD_AVX512F_MICROKERNEL_SRCS": "defined(__i386__) || defined(__i686__) || defined(__x86_64__)",
|
|
"PROD_AVX512SKX_MICROKERNEL_SRCS": "defined(__i386__) || defined(__i686__) || defined(__x86_64__)",
|
|
"PROD_AVX512VBMI_MICROKERNEL_SRCS": "defined(__i386__) || defined(__i686__) || defined(__x86_64__)",
|
|
"PROD_AVX512VNNI_MICROKERNEL_SRCS": "defined(__i386__) || defined(__i686__) || defined(__x86_64__)",
|
|
"PROD_RVV_MICROKERNEL_SRCS": "defined(__riscv) || defined(__riscv__)",
|
|
"PROD_AVXVNNI_MICROKERNEL_SRCS": "defined(__i386__) || defined(__i686__) || defined(__x86_64__)",
|
|
"AARCH32_ASM_MICROKERNEL_SRCS": "defined(__arm__)",
|
|
"AARCH64_ASM_MICROKERNEL_SRCS": "defined(__aarch64__)",
|
|
|
|
# add non-prod microkernel sources here:
|
|
}
|
|
|
|
SRC_NAMES = set([
|
|
"OPERATOR_SRCS",
|
|
"SUBGRAPH_SRCS",
|
|
"LOGGING_SRCS",
|
|
"XNNPACK_SRCS",
|
|
"TABLE_SRCS",
|
|
"JIT_SRCS",
|
|
"PROD_SCALAR_MICROKERNEL_SRCS",
|
|
"PROD_FMA_MICROKERNEL_SRCS",
|
|
"PROD_ARMSIMD32_MICROKERNEL_SRCS",
|
|
"PROD_FP16ARITH_MICROKERNEL_SRCS",
|
|
"PROD_NEON_MICROKERNEL_SRCS",
|
|
"PROD_NEONFP16_MICROKERNEL_SRCS",
|
|
"PROD_NEONFMA_MICROKERNEL_SRCS",
|
|
"PROD_NEON_AARCH64_MICROKERNEL_SRCS",
|
|
"PROD_NEONV8_MICROKERNEL_SRCS",
|
|
"PROD_NEONFP16ARITH_MICROKERNEL_SRCS",
|
|
"PROD_NEONFP16ARITH_AARCH64_MICROKERNEL_SRCS",
|
|
"PROD_NEONDOT_MICROKERNEL_SRCS",
|
|
"PROD_NEONDOT_AARCH64_MICROKERNEL_SRCS",
|
|
"PROD_NEONDOTFP16ARITH_MICROKERNEL_SRCS",
|
|
"PROD_NEONDOTFP16ARITH_AARCH64_MICROKERNEL_SRCS",
|
|
"PROD_NEONI8MM_MICROKERNEL_SRCS",
|
|
"PROD_SSE_MICROKERNEL_SRCS",
|
|
"PROD_SSE2_MICROKERNEL_SRCS",
|
|
"PROD_SSSE3_MICROKERNEL_SRCS",
|
|
"PROD_SSE41_MICROKERNEL_SRCS",
|
|
"PROD_AVX_MICROKERNEL_SRCS",
|
|
"PROD_F16C_MICROKERNEL_SRCS",
|
|
"PROD_XOP_MICROKERNEL_SRCS",
|
|
"PROD_FMA3_MICROKERNEL_SRCS",
|
|
"PROD_AVX2_MICROKERNEL_SRCS",
|
|
"PROD_AVX512F_MICROKERNEL_SRCS",
|
|
"PROD_AVX512SKX_MICROKERNEL_SRCS",
|
|
"PROD_AVX512VBMI_MICROKERNEL_SRCS",
|
|
"PROD_AVX512VNNI_MICROKERNEL_SRCS",
|
|
"PROD_RVV_MICROKERNEL_SRCS",
|
|
"PROD_AVXVNNI_MICROKERNEL_SRCS",
|
|
"AARCH32_ASM_MICROKERNEL_SRCS",
|
|
"AARCH64_ASM_MICROKERNEL_SRCS",
|
|
|
|
# add non-prod microkernel sources here:
|
|
])
|
|
|
|
def handle_singleline_parse(line):
|
|
start_index = line.find("(")
|
|
end_index = line.find(")")
|
|
line = line[start_index+1:end_index]
|
|
key_val = line.split(" ")
|
|
return key_val[0], list(map(lambda x: x[4:], key_val[1:]))
|
|
|
|
def update_sources(xnnpack_path, cmakefile = "XNNPACK/CMakeLists.txt"):
|
|
sources = collections.defaultdict(list)
|
|
with open(os.path.join(xnnpack_path, cmakefile)) as cmake:
|
|
lines = cmake.readlines()
|
|
i = 0
|
|
while i < len(lines):
|
|
line = lines[i]
|
|
|
|
if lines[i].startswith("SET") and "src/" in lines[i]:
|
|
name, val = handle_singleline_parse(line)
|
|
sources[name].extend(val)
|
|
i+=1
|
|
continue
|
|
|
|
if line.startswith("SET") and line.split('(')[1].strip(' \t\n\r') in set(WRAPPER_SRC_NAMES.keys()) | set(SRC_NAMES):
|
|
name = line.split('(')[1].strip(' \t\n\r')
|
|
i += 1
|
|
while i < len(lines) and len(lines[i]) > 0 and ')' not in lines[i]:
|
|
# remove "src/" at the beginning, remove whitespaces and newline
|
|
value = lines[i].strip(' \t\n\r')
|
|
sources[name].append(value[4:])
|
|
i += 1
|
|
if i < len(lines) and len(lines[i]) > 4:
|
|
# remove "src/" at the beginning, possibly ')' at the end
|
|
value = lines[i].strip(' \t\n\r)')
|
|
sources[name].append(value[4:])
|
|
else:
|
|
i += 1
|
|
return sources
|
|
|
|
def gen_wrappers(xnnpack_path):
|
|
xnnpack_sources = collections.defaultdict(list)
|
|
sources = update_sources(xnnpack_path)
|
|
|
|
microkernels_sources = update_sources(xnnpack_path, "XNNPACK/cmake/microkernels.cmake")
|
|
for key in microkernels_sources:
|
|
sources[key] = microkernels_sources[key]
|
|
|
|
for name in WRAPPER_SRC_NAMES:
|
|
xnnpack_sources[WRAPPER_SRC_NAMES[name]].extend(sources[name])
|
|
|
|
for condition, filenames in xnnpack_sources.items():
|
|
print(condition)
|
|
for filename in filenames:
|
|
filepath = os.path.join(xnnpack_path, "xnnpack_wrappers", filename)
|
|
|
|
if not os.path.isdir(os.path.dirname(filepath)):
|
|
os.makedirs(os.path.dirname(filepath))
|
|
with open(filepath, "w") as wrapper:
|
|
print("/* {} */".format(BANNER), file=wrapper)
|
|
print(file=wrapper)
|
|
|
|
# Architecture- or platform-dependent preprocessor flags can be
|
|
# defined here. Note: platform_preprocessor_flags can't be used
|
|
# because they are ignored by arc focus & buck project.
|
|
|
|
if condition is None:
|
|
print("#include <%s>" % filename, file=wrapper)
|
|
else:
|
|
# Include source file only if condition is satisfied
|
|
print("#if %s" % condition, file=wrapper)
|
|
print("#include <%s>" % filename, file=wrapper)
|
|
print("#endif /* %s */" % condition, file=wrapper)
|
|
|
|
# update xnnpack_wrapper_defs.bzl file under the same folder
|
|
with open(os.path.join(os.path.dirname(__file__), "xnnpack_wrapper_defs.bzl"), 'w') as wrapper_defs:
|
|
print('"""', file=wrapper_defs)
|
|
print(BANNER, file=wrapper_defs)
|
|
print('"""', file=wrapper_defs)
|
|
for name in WRAPPER_SRC_NAMES:
|
|
print('\n' + name + ' = [', file=wrapper_defs)
|
|
for file_name in sources[name]:
|
|
print(' "xnnpack_wrappers/{}",'.format(file_name), file=wrapper_defs)
|
|
print(']', file=wrapper_defs)
|
|
|
|
# update xnnpack_src_defs.bzl file under the same folder
|
|
with open(os.path.join(os.path.dirname(__file__), "xnnpack_src_defs.bzl"), 'w') as src_defs:
|
|
print('"""', file=src_defs)
|
|
print(BANNER, file=src_defs)
|
|
print('"""', file=src_defs)
|
|
for name in SRC_NAMES:
|
|
print('\n' + name + ' = [', file=src_defs)
|
|
for file_name in sources[name]:
|
|
print(' "XNNPACK/src/{}",'.format(file_name), file=src_defs)
|
|
print(']', file=src_defs)
|
|
|
|
|
|
def main(argv):
|
|
if argv is None or len(argv) == 0:
|
|
gen_wrappers(".")
|
|
else:
|
|
gen_wrappers(argv[0])
|
|
|
|
# The first argument is the place where the "xnnpack_wrappers" folder will be created.
|
|
# Run it without arguments will generate "xnnpack_wrappers" in the current path.
|
|
# The two .bzl files will always be generated in the current path.
|
|
if __name__ == "__main__":
|
|
main(sys.argv[1:])
|