wehub-resource-sync
129 lines
3.6 KiB
Python
129 lines
3.6 KiB
Python
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""
|
|
This script simply removes grad kernels. You should use this script
|
|
when cmake ON_INFER=ON, which can greatly reduce the volume of the inference library.
|
|
"""
|
|
|
|
import glob
|
|
import os
|
|
|
|
|
|
def is_balanced(content):
|
|
"""
|
|
Check whether sequence contains valid parenthesis.
|
|
Args:
|
|
content (str): content of string.
|
|
|
|
Returns:
|
|
boolean: True if sequence contains valid parenthesis.
|
|
"""
|
|
|
|
if content.find('{') == -1:
|
|
return False
|
|
stack = []
|
|
push_chars, pop_chars = '({', ')}'
|
|
for c in content:
|
|
if c in push_chars:
|
|
stack.append(c)
|
|
elif c in pop_chars:
|
|
if not len(stack):
|
|
return False
|
|
else:
|
|
stack_top = stack.pop()
|
|
balancing_bracket = push_chars[pop_chars.index(c)]
|
|
if stack_top != balancing_bracket:
|
|
return False
|
|
return not stack
|
|
|
|
|
|
def grad_kernel_definition(content, kernel_pattern, grad_pattern):
|
|
"""
|
|
Args:
|
|
content(str): file content
|
|
kernel_pattern(str): kernel pattern
|
|
grad_pattern(str): grad pattern
|
|
|
|
Returns:
|
|
(list, int): grad kernel definitions in file and count.
|
|
"""
|
|
|
|
results = []
|
|
count = 0
|
|
start = 0
|
|
lens = len(content)
|
|
while True:
|
|
index = content.find(kernel_pattern, start)
|
|
if index == -1:
|
|
return results, count
|
|
i = index + 1
|
|
while i <= lens:
|
|
check_str = content[index:i]
|
|
if is_balanced(check_str):
|
|
if check_str.find(grad_pattern) != -1:
|
|
results.append(check_str)
|
|
count += 1
|
|
start = i
|
|
break
|
|
i += 1
|
|
else:
|
|
return results, count
|
|
|
|
|
|
def remove_grad_kernels(dry_run=False):
|
|
"""
|
|
Args:
|
|
dry_run(bool): whether just print
|
|
|
|
Returns:
|
|
int: number of kernel(grad) removed
|
|
"""
|
|
|
|
pd_kernel_pattern = 'PD_REGISTER_STRUCT_KERNEL'
|
|
register_op_pd_kernel_count = 0
|
|
matches = []
|
|
|
|
tool_dir = os.path.dirname(os.path.abspath(__file__))
|
|
all_op = glob.glob(
|
|
os.path.join(tool_dir, '../paddle/fluid/operators/**/*.cc'),
|
|
recursive=True,
|
|
)
|
|
all_op += glob.glob(
|
|
os.path.join(tool_dir, '../paddle/fluid/operators/**/*.cu'),
|
|
recursive=True,
|
|
)
|
|
|
|
for op_file in all_op:
|
|
with open(op_file, 'r', encoding='utf-8') as f:
|
|
content = ''.join(f.readlines())
|
|
|
|
pd_kernel, pd_kernel_count = grad_kernel_definition(
|
|
content, pd_kernel_pattern, '_grad,'
|
|
)
|
|
|
|
register_op_pd_kernel_count += pd_kernel_count
|
|
|
|
matches.extend(pd_kernel)
|
|
|
|
for to_remove in matches:
|
|
content = content.replace(to_remove, '')
|
|
if dry_run:
|
|
print(op_file, to_remove)
|
|
|
|
if not dry_run:
|
|
with open(op_file, 'w', encoding='utf-8') as f:
|
|
f.write(content)
|
|
|
|
return register_op_pd_kernel_count
|