wehub-resource-sync
128 lines
3.7 KiB
Python
128 lines
3.7 KiB
Python
#!/usr/bin/env python3
|
|
import os
|
|
import sys
|
|
|
|
execute_treat_list = [
|
|
"source/backend/cpu/CPUOPRegister.cpp",
|
|
"source/backend/metal/MetalOPRegister.mm"
|
|
]
|
|
|
|
|
|
shape_treat_file = "source/shape/ShapeRegister.cpp"
|
|
import re
|
|
|
|
def extract_geometry_and_op_types(code: str):
|
|
register_pattern = r'REGISTER_GEOMETRY\s*\(\s*(\w+)\s*,\s*(\w+)\s*\)\s*;'
|
|
matches = re.findall(register_pattern, code)
|
|
|
|
results = []
|
|
|
|
for geometry_type, func_name in matches:
|
|
func_pattern = rf'static\s+void\s+{re.escape(func_name)}\s*\(\s*\)\s*\{{(.*?)\}}'
|
|
func_match = re.search(func_pattern, code, re.DOTALL)
|
|
|
|
op_types = []
|
|
if func_match:
|
|
func_body = func_match.group(1)
|
|
op_type_pattern = r'OpType_\w+'
|
|
op_types = re.findall(op_type_pattern, func_body)
|
|
|
|
results.append((geometry_type, op_types))
|
|
|
|
return results
|
|
|
|
def load_op_set(op_file):
|
|
ops = set()
|
|
if os.path.exists(op_file):
|
|
with open(op_file, 'r') as f:
|
|
for line in f:
|
|
op = line.strip()
|
|
if op and not op.startswith('#'):
|
|
ops.add("OpType_" + op)
|
|
return ops
|
|
if __name__ == '__main__':
|
|
if len(sys.argv) <3:
|
|
print("Usage: python3 prue_mnn_ops.py op.txt ${MNN_ROOT}")
|
|
exit()
|
|
ops = load_op_set(sys.argv[1])
|
|
# push default op
|
|
exes = ops
|
|
exes.add("OpType_Cast")
|
|
exes.add("OpType_FloatToInt8")
|
|
exes.add("OpType_Int8ToFloat")
|
|
exes.add("OpType_Int8ToFloat")
|
|
exes.add("OpType_Raster")
|
|
exes.add("OpType_While")
|
|
|
|
mnnroot = sys.argv[2]
|
|
for filename in execute_treat_list:
|
|
lines = []
|
|
with open(mnnroot + "/" + filename, 'r') as f:
|
|
lines = f.read().split('\n')
|
|
full = ""
|
|
for line in lines:
|
|
valid = line.find('OpType') < 0
|
|
if not valid:
|
|
for op in exes:
|
|
if line.find('_' + op + '_') >= 0:
|
|
valid = True
|
|
if valid:
|
|
full += line + '\n'
|
|
with open(mnnroot + "/" + filename, 'w') as f:
|
|
f.write(full)
|
|
lines = []
|
|
with open(mnnroot + "/" + shape_treat_file, 'r') as f:
|
|
lines = f.read().split('\n')
|
|
full = ""
|
|
for line in lines:
|
|
valid = line.find('OpType') < 0
|
|
if not valid:
|
|
for op in ops:
|
|
if line.find('_' + op + '_') >= 0:
|
|
valid = True
|
|
if valid:
|
|
full += line + '\n'
|
|
with open(mnnroot + "/" + shape_treat_file, 'w') as f:
|
|
f.write(full)
|
|
|
|
# Search Geometry File to make op type and file link
|
|
geoDir = mnnroot + "/" + "source/geometry"
|
|
fileNames = os.listdir(geoDir)
|
|
funcNames = []
|
|
geometry_treat_file = mnnroot + "/" + "source/geometry/GeometryOPRegister.cpp"
|
|
need_delete_geo = []
|
|
for fi in fileNames:
|
|
if ".hpp" in fi:
|
|
continue
|
|
f = os.path.join(geoDir, fi)
|
|
with open(f) as fr:
|
|
res = extract_geometry_and_op_types(fr.read())
|
|
for r, oplist in res:
|
|
valid = False
|
|
for op in oplist:
|
|
print(op)
|
|
if op in ops:
|
|
valid = True
|
|
break
|
|
if not valid:
|
|
need_delete_geo.append(r)
|
|
print("Remove Geometry Op:", need_delete_geo)
|
|
lines = []
|
|
with open(geometry_treat_file, 'r') as f:
|
|
lines = f.read().split('\n')
|
|
full = ""
|
|
for line in lines:
|
|
valid = True
|
|
for op in need_delete_geo:
|
|
if line.find('_' + op + '_') >=0:
|
|
valid = False
|
|
if valid:
|
|
full += line + '\n'
|
|
with open(geometry_treat_file, 'w') as f:
|
|
f.write(full)
|
|
|
|
|
|
|
|
|
|
|