import argparse
import ast
import importlib
import pathlib
def get_njit_funcs(pkg_dir):
"""
Identify all njit functions
Parameters
----------
pkg_dir : str
The path to the directory containing some .py files
Returns
-------
njit_funcs : list
A list of all njit functions, where each element is a tuple of the form
(module_name, func_name)
"""
ignore_py_files = ["__init__", "__pycache__"]
pkg_dir = pathlib.Path(pkg_dir)
module_names = []
for fname in pkg_dir.iterdir():
if fname.stem not in ignore_py_files and not fname.stem.startswith("."):
module_names.append(fname.stem)
njit_funcs = []
for module_name in module_names:
filepath = pkg_dir / f"{module_name}.py"
file_contents = ""
with open(filepath, encoding="utf8") as f:
file_contents = f.read()
module = ast.parse(file_contents)
for node in module.body:
if isinstance(node, ast.FunctionDef):
func_name = node.name
for decorator in node.decorator_list:
decorator_name = None
if isinstance(decorator, ast.Name):
decorator_name = decorator.id
if isinstance(decorator, ast.Call) and isinstance(
decorator.func, ast.Name
):
decorator_name = decorator.func.id
if decorator_name == "njit":
njit_funcs.append((module_name, func_name))
return njit_funcs
def check_fastmath(pkg_dir, pkg_name):
"""
Check if all njit functions have the `fastmath` flag set
Parameters
----------
pkg_dir : str
The path to the directory containing some .py files
pkg_name : str
The name of the package
Returns
-------
None
"""
missing_fastmath = []
for module_name, func_name in get_njit_funcs(pkg_dir):
module = importlib.import_module(f".{module_name}", package=pkg_name)
func = getattr(module, func_name)
if "fastmath" not in func.targetoptions.keys():
missing_fastmath.append(f"{module_name}.{func_name}")
if len(missing_fastmath) > 0:
msg = (
"Found one or more `@njit` functions that are missing the `fastmath` flag. "
+ f"The functions are:\n {missing_fastmath}\n"
)
raise ValueError(msg)
return
class FunctionCallVisitor(ast.NodeVisitor):
"""
A class to traverse the AST of the modules of a package to collect
the call stacks of njit functions.
Parameters
----------
pkg_dir : str
The path to the package directory containing some .py files.
pkg_name : str
The name of the package.
Attributes
----------
module_names : list
A list of module names to track the modules as the visitor traverses them.
call_stack : list
A list of njit functions, representing a chain of function calls,
where each element is a string of the form "module_name.func_name".
out : list
A list of unique `call_stack`s.
njit_funcs : list
A list of all njit functions in `pkg_dir`'s modules. Each element is a tuple
of the form `(module_name, func_name)`.
njit_modules : set
A set that contains the names of all modules, each of which contains at least
one njit function.
njit_nodes : dict
A dictionary mapping njit function names to their corresponding AST nodes.
A key is a string, and it is of the form "module_name.func_name", and its
corresponding value is the AST node- with type ast.FunctionDef- of that
function.
ast_modules : dict
A dictionary mapping module names to their corresponding AST objects. A key
is the name of a module, and its corresponding value is the content of that
module as an AST object.
Methods
-------
push_module(module_name)
Push the name of a module onto the stack `module_names`.
pop_module()
Pop the last module name from the stack `module_names`.
push_call_stack(module_name, func_name)
Push a function call onto the stack of function calls, `call_stack`.
pop_call_stack()
Pop the last function call from the stack of function calls, `call_stack`
goto_deeper_func(node)
Calls the visit method from class `ast.NodeVisitor` on all children of
the `node`.
goto_next_func(node)
Calls the visit method from class `ast.NodeVisitor` on all children of
the `node`.
push_out()
Push the current function call stack, `call_stack`, onto the output list, `out`,
unless it is already included in one of the so-far-collected call stacks.
visit_Call(node)
This method is called when the visitor encounters a function call in the AST. It
checks if the called function is a njit function and, if so, traverses its AST
to collect its call stack.
"""
def __init__(self, pkg_dir, pkg_name):
"""
Initialize the FunctionCallVisitor class. This method sets up the necessary
attributes and prepares the visitor for traversing the AST of STUMPY's modules.
Parameters
----------
pkg_dir : str
The path to the package directory containing some .py files.
pkg_name : str
The name of the package.
Returns
-------
None
"""
super().__init__()
self.module_names = []
self.call_stack = []
self.out = []
self.njit_funcs = get_njit_funcs(pkg_dir)
self.njit_modules = set(mod_name for mod_name, func_name in self.njit_funcs)
self.njit_nodes = {}
self.ast_modules = {}
filepaths = sorted(f for f in pathlib.Path(pkg_dir).iterdir() if f.is_file())
ignore = ["__init__.py", "__pycache__"]
for filepath in filepaths:
file_name = filepath.name
if (
file_name not in ignore
and not file_name.startswith("gpu")
and str(filepath).endswith(".py")
):
module_name = file_name.replace(".py", "")
file_contents = ""
with open(filepath, encoding="utf8") as f:
file_contents = f.read()
self.ast_modules[module_name] = ast.parse(file_contents)
for node in self.ast_modules[module_name].body:
if isinstance(node, ast.FunctionDef):
func_name = node.name
if (module_name, func_name) in self.njit_funcs:
self.njit_nodes[f"{module_name}.{func_name}"] = node
def push_module(self, module_name):
"""
Push a module name onto the stack of module names.
Parameters
----------
module_name : str
The name of the module to be pushed onto the stack.
Returns
-------
None
"""
self.module_names.append(module_name)
return
def pop_module(self):
"""
Pop the last module name from the stack of module names.
Parameters
----------
None
Returns
-------
None
"""
if self.module_names:
self.module_names.pop()
return
def push_call_stack(self, module_name, func_name):
"""
Push a function call onto the stack of function calls.
Parameters
----------
module_name : str
A module's name
func_name : str
A function's name
Returns
-------
None
"""
self.call_stack.append(f"{module_name}.{func_name}")
return
def pop_call_stack(self):
"""
Pop the last function call from the stack of function calls.
Parameters
----------
None
Returns
-------
None
"""
if self.call_stack:
self.call_stack.pop()
return
def goto_deeper_func(self, node):
"""
Calls the visit method from class `ast.NodeVisitor` on
all children of the `node`.
Parameters
----------
node : ast.AST
The AST node to be visited.
Returns
-------
None
"""
self.generic_visit(node)
return
def goto_next_func(self, node):
"""
Calls the visit method from class `ast.NodeVisitor` on
all children of the node.
Parameters
----------
node : ast.AST
The AST node to be visited.
Returns
-------
None
"""
self.generic_visit(node)
return
def push_out(self):
"""
Push the current function call stack onto the output list unless it
is already included in one of the so-far-collected call stacks.
Parameters
----------
None
Returns
-------
None
"""
unique = True
for cs in self.out:
if " ".join(self.call_stack) in " ".join(cs):
unique = False
break
if unique:
self.out.append(self.call_stack.copy())
return
def visit_Call(self, node):
"""
Called when visiting an AST node of type `ast.Call`.
Parameters
----------
node : ast.Call
The AST node representing a function call.
Returns
-------
None
"""
callee_name = ast.unparse(node.func)
module_changed = False
if "." in callee_name:
new_module_name, new_func_name = callee_name.split(".")[:2]
if new_module_name in self.njit_modules:
self.push_module(new_module_name)
module_changed = True
else:
if self.module_names:
new_module_name = self.module_names[-1]
new_func_name = callee_name
callee_name = f"{new_module_name}.{new_func_name}"
if callee_name in self.njit_nodes.keys():
callee_node = self.njit_nodes[callee_name]
self.push_call_stack(new_module_name, new_func_name)
self.goto_deeper_func(callee_node)
self.push_out()
self.pop_call_stack()
if module_changed:
self.pop_module()
self.goto_next_func(node)
return
def get_njit_call_stacks(pkg_dir, pkg_name):
"""
Get the call stacks of all njit functions in `pkg_dir`
Parameters
----------
pkg_dir : str
The path to the package directory containing some .py files
pkg_name : str
The name of the package
Returns
-------
out : list
A list of unique function call stacks. Each item is of type list,
representing a chain of function calls.
"""
visitor = FunctionCallVisitor(pkg_dir, pkg_name)
for module_name in visitor.njit_modules:
visitor.push_module(module_name)
for node in visitor.ast_modules[module_name].body:
if isinstance(node, ast.FunctionDef):
func_name = node.name
if (module_name, func_name) in visitor.njit_funcs:
visitor.push_call_stack(module_name, func_name)
visitor.visit(node)
visitor.pop_call_stack()
visitor.pop_module()
return visitor.out
def check_call_stack_fastmath(pkg_dir, pkg_name):
"""
Check if all njit functions in a call stack have the same `fastmath` flag.
This function raises a ValueError if it finds any inconsistencies in the
`fastmath` flags in at lease one call stack of njit functions.
Parameters
----------
pkg_dir : str
The path to the directory containing some .py files
pkg_name : str
The name of the package
Returns
-------
None
"""
inconsistent_call_stacks = []
njit_call_stacks = get_njit_call_stacks(pkg_dir, pkg_name)
for cs in njit_call_stacks:
module_name, func_name = cs[0].split(".")
module = importlib.import_module(f".{module_name}", package="stumpy")
func = getattr(module, func_name)
flag_ref = func.targetoptions["fastmath"]
for item in cs[1:]:
module_name, func_name = cs[0].split(".")
module = importlib.import_module(f".{module_name}", package="stumpy")
func = getattr(module, func_name)
flag = func.targetoptions["fastmath"]
if flag != flag_ref:
inconsistent_call_stacks.append(cs)
break
if len(inconsistent_call_stacks) > 0:
msg = (
"Found at least one call stack that has inconsistent `fastmath` flags. "
+ f"Those call stacks are:\n {inconsistent_call_stacks}\n"
)
raise ValueError(msg)
return
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--check", dest="pkg_dir")
args = parser.parse_args()
if args.pkg_dir:
pkg_dir = pathlib.Path(args.pkg_dir)
pkg_name = pkg_dir.name
check_fastmath(str(pkg_dir), pkg_name)
check_call_stack_fastmath(str(pkg_dir), pkg_name)