From aeb85247d5921f8b4ae1ebdacf36f4d9d79bed47 Mon Sep 17 00:00:00 2001 From: NimaSarajpoor Date: Tue, 6 Aug 2024 02:50:05 -0400 Subject: [PATCH 1/9] Empty commit From 74392173fddb0383c28b6ca4823c3d7b0de02c76 Mon Sep 17 00:00:00 2001 From: NimaSarajpoor Date: Sun, 22 Sep 2024 23:17:02 -0400 Subject: [PATCH 2/9] create script to check fastmath flag in njit decorator --- utils.py | 95 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 95 insertions(+) create mode 100644 utils.py diff --git a/utils.py b/utils.py new file mode 100644 index 000000000..92a14033d --- /dev/null +++ b/utils.py @@ -0,0 +1,95 @@ +import ast + + +def check_fastmath(decorator): + """ + For the given `decorator` node with type `ast.Call`, + return the value of the `fastmath` argument if it exists. + Otherwise, return `None`. + """ + fastmath_value = None + for n in ast.iter_child_nodes(decorator): + if isinstance(n, ast.keyword) and n.arg == 'fastmath': + if isinstance(n.value, ast.Constant): + fastmath_value = n.value.value + elif isinstance(n.value, ast.Set): + fastmath_value = set(item.value for item in n.value.elts) + else: + pass + break + + return fastmath_value + + + +def check_njit(fd): + """ + For the given `fd` node with type `ast.FunctionDef`, + return the node of the `njit` decorator if it exists. + Otherwise, return `None`. + """ + decorator_node = None + for decorator in fd.decorator_list: + if not isinstance(decorator, ast.Call): + continue + + obj = decorator.func + if isinstance(obj, ast.Attribute): + name = obj.attr + elif isinstance(obj, ast.Subscript): + name = obj.value.id + elif isinstance(obj, ast.Name): + name = obj.id + else: + msg = f"The type {type(obj)} is not supported." + raise ValueError(msg) + + if name == "njit": + decorator_node = decorator + break + + return decorator_node + + +def check_functions(filepath): + """ + For the given `filepath`, return the function names, + whether the function is decorated with `@njit`, + and the value of the `fastmath` argument if it exists + + Parameters + ---------- + filepath : str + The path to the file + + Returns + ------- + func_names : list + List of function names + + is_njit : list + List of boolean values indicating whether the function is decorated with `@njit` + + fastmath_value : list + List of values of the `fastmath` argument if it exists + """ + file_contents = "" + with open(filepath, encoding="utf8") as f: + file_contents = f.read() + module = ast.parse(file_contents) + + function_definitions = [ + node for node in module.body if isinstance(node, ast.FunctionDef) + ] + + func_names = [fd.name for fd in function_definitions] + + njit_nodes = [check_njit(fd) for fd in function_definitions] + is_njit = [node is not None for node in njit_nodes] + + fastmath_value = [None] * len(njit_nodes) + for i, node in enumerate(njit_nodes): + if node is not None: + fastmath_value[i] = check_fastmath(node) + + return func_names, is_njit, fastmath_values \ No newline at end of file From 82bfeb19e06651910f599d1680f8aa8eead7f001 Mon Sep 17 00:00:00 2001 From: NimaSarajpoor Date: Sun, 22 Sep 2024 23:42:30 -0400 Subject: [PATCH 3/9] fix for black and flake8 --- utils.py | 33 ++++++++++++++++----------------- 1 file changed, 16 insertions(+), 17 deletions(-) diff --git a/utils.py b/utils.py index 92a14033d..a3daba646 100644 --- a/utils.py +++ b/utils.py @@ -3,13 +3,13 @@ def check_fastmath(decorator): """ - For the given `decorator` node with type `ast.Call`, + For the given `decorator` node with type `ast.Call`, return the value of the `fastmath` argument if it exists. Otherwise, return `None`. """ fastmath_value = None for n in ast.iter_child_nodes(decorator): - if isinstance(n, ast.keyword) and n.arg == 'fastmath': + if isinstance(n, ast.keyword) and n.arg == "fastmath": if isinstance(n.value, ast.Constant): fastmath_value = n.value.value elif isinstance(n.value, ast.Set): @@ -21,10 +21,9 @@ def check_fastmath(decorator): return fastmath_value - def check_njit(fd): """ - For the given `fd` node with type `ast.FunctionDef`, + For the given `fd` node with type `ast.FunctionDef`, return the node of the `njit` decorator if it exists. Otherwise, return `None`. """ @@ -38,7 +37,7 @@ def check_njit(fd): name = obj.attr elif isinstance(obj, ast.Subscript): name = obj.value.id - elif isinstance(obj, ast.Name): + elif isinstance(obj, ast.Name): name = obj.id else: msg = f"The type {type(obj)} is not supported." @@ -47,29 +46,29 @@ def check_njit(fd): if name == "njit": decorator_node = decorator break - + return decorator_node def check_functions(filepath): """ - For the given `filepath`, return the function names, - whether the function is decorated with `@njit`, + For the given `filepath`, return the function names, + whether the function is decorated with `@njit`, and the value of the `fastmath` argument if it exists Parameters ---------- filepath : str The path to the file - + Returns ------- func_names : list List of function names - + is_njit : list List of boolean values indicating whether the function is decorated with `@njit` - + fastmath_value : list List of values of the `fastmath` argument if it exists """ @@ -77,19 +76,19 @@ def check_functions(filepath): with open(filepath, encoding="utf8") as f: file_contents = f.read() module = ast.parse(file_contents) - + function_definitions = [ - node for node in module.body if isinstance(node, ast.FunctionDef) + node for node in module.body if isinstance(node, ast.FunctionDef) ] func_names = [fd.name for fd in function_definitions] - njit_nodes = [check_njit(fd) for fd in function_definitions] + njit_nodes = [check_njit(fd) for fd in function_definitions] is_njit = [node is not None for node in njit_nodes] - fastmath_value = [None] * len(njit_nodes) + fastmath_values = [None] * len(njit_nodes) for i, node in enumerate(njit_nodes): if node is not None: - fastmath_value[i] = check_fastmath(node) + fastmath_values[i] = check_fastmath(node) - return func_names, is_njit, fastmath_values \ No newline at end of file + return func_names, is_njit, fastmath_values From 5cd3a35059e84e4b8f19ae85bcc24b822e0fcaba Mon Sep 17 00:00:00 2001 From: NimaSarajpoor Date: Mon, 23 Sep 2024 22:10:06 -0400 Subject: [PATCH 4/9] add function to obtain callees of functions in a script --- utils.py | 52 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/utils.py b/utils.py index a3daba646..6c9c41423 100644 --- a/utils.py +++ b/utils.py @@ -92,3 +92,55 @@ def check_functions(filepath): fastmath_values[i] = check_fastmath(node) return func_names, is_njit, fastmath_values + + +def _get_callees(node, all_functions): + for n in ast.iter_child_nodes(node): + if isinstance(n, ast.Call): + obj = n.func + if isinstance(obj, ast.Attribute): + name = obj.attr + elif isinstance(obj, ast.Subscript): + name = obj.value.id + elif isinstance(obj, ast.Name): + name = obj.id + else: + msg = f"The type {type(obj)} is not supported" + raise ValueError(msg) + + all_functions.append(name) + + _get_callees(n, all_functions) + + +def get_all_callees(fd): + """ + For a given node of type ast.FunctionDef, visit all of its child nodes, + and return a list of all of its callees + """ + all_functions = [] + _get_callees(fd, all_functions) + + return all_functions + + +def check_callees(filepath): + """ + For the given `filepath`, return a dictionary with the key + being the function name and the value being a set of function names + that are called by the function + """ + file_contents = "" + with open(filepath, encoding="utf8") as f: + file_contents = f.read() + module = ast.parse(file_contents) + + function_definitions = [ + node for node in module.body if isinstance(node, ast.FunctionDef) + ] + + callees = {} + for fd in function_definitions: + callees[fd.name] = set(get_all_callees(fd)) + + return callees From 033bc3bac07fe4276813012737ea0ba17f107c6b Mon Sep 17 00:00:00 2001 From: NimaSarajpoor Date: Sat, 28 Sep 2024 20:43:51 -0400 Subject: [PATCH 5/9] add code to get dictionary of functions metadata --- njit_fastmath.py | 69 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 69 insertions(+) create mode 100644 njit_fastmath.py diff --git a/njit_fastmath.py b/njit_fastmath.py new file mode 100644 index 000000000..5c83ba09d --- /dev/null +++ b/njit_fastmath.py @@ -0,0 +1,69 @@ + +import pathlib + +from utils import check_callees, check_functions + +stumpy_path = pathlib.Path(__file__).parent / "stumpy" +filepaths = sorted(f for f in pathlib.Path(stumpy_path).iterdir() if f.is_file()) + +all_functions = {} +callees = {} + +ignore = ["__init__.py", "__pycache__"] +for filepath in filepaths: + name = filepath.name + if name not in ignore and str(filepath).endswith(".py"): + callees[name] = check_callees(filepath) + + func_names, is_njit, fastmath_values = check_functions(filepath) + all_functions[name] = { + "func_names": func_names, + "is_njit": is_njit, + "fastmath_values": fastmath_values, + } + + +stumpy_functions = set() +for fname, func_data in all_functions.items(): + prefix = fname.replace(".py", "") + stumpy_functions.update([prefix + '.' + x for x in func_data["func_names"]]) +stumpy_functions = list(stumpy_functions) +stumpy_functions_no_prefix = [x.split('.')[1] for x in stumpy_functions] + + +# create a dictionary where keys are function names in stumpy_functions, and the value +# is a tuple, where the first element says whether the function is decorated with njit +# and the second element is the list of callees + +callers_callees = {} +for func_name in stumpy_functions: + callers_callees[func_name] = None + + prefix, func = func_name.split('.') + fname = prefix + '.py' + + idx = all_functions[fname]["func_names"].index(func) + is_njit = all_functions[fname]["is_njit"][idx] + fastmath_val = all_functions[fname]["fastmath_values"][idx] + + callees_functions = callees[fname][func] + + pruned_callees_functions = [] + for callee in callees_functions: + if callee in all_functions[fname]["func_names"]: + pruned_callees_functions.append(prefix + '.' + callee) + elif callee in stumpy_functions: + idx = stumpy_functions_no_prefix.index(callee) + pruned_callees_functions.append(stumpy_functions[idx]) + else: + continue + + callers_callees[func_name] = (is_njit, fastmath_val, pruned_callees_functions) + + + + +for k, v in callers_callees.items(): + print('---------------------------------') + print('> func_name: ', k) + print('metadata: ', v) \ No newline at end of file From ab280e56e5b2aab93248418427897d806743e129 Mon Sep 17 00:00:00 2001 From: NimaSarajpoor Date: Sat, 28 Sep 2024 20:44:29 -0400 Subject: [PATCH 6/9] fix format --- njit_fastmath.py | 23 +++++++---------------- 1 file changed, 7 insertions(+), 16 deletions(-) diff --git a/njit_fastmath.py b/njit_fastmath.py index 5c83ba09d..7a03110f0 100644 --- a/njit_fastmath.py +++ b/njit_fastmath.py @@ -1,4 +1,3 @@ - import pathlib from utils import check_callees, check_functions @@ -26,9 +25,9 @@ stumpy_functions = set() for fname, func_data in all_functions.items(): prefix = fname.replace(".py", "") - stumpy_functions.update([prefix + '.' + x for x in func_data["func_names"]]) + stumpy_functions.update([prefix + "." + x for x in func_data["func_names"]]) stumpy_functions = list(stumpy_functions) -stumpy_functions_no_prefix = [x.split('.')[1] for x in stumpy_functions] +stumpy_functions_no_prefix = [x.split(".")[1] for x in stumpy_functions] # create a dictionary where keys are function names in stumpy_functions, and the value @@ -38,20 +37,20 @@ callers_callees = {} for func_name in stumpy_functions: callers_callees[func_name] = None - - prefix, func = func_name.split('.') - fname = prefix + '.py' + + prefix, func = func_name.split(".") + fname = prefix + ".py" idx = all_functions[fname]["func_names"].index(func) is_njit = all_functions[fname]["is_njit"][idx] fastmath_val = all_functions[fname]["fastmath_values"][idx] callees_functions = callees[fname][func] - + pruned_callees_functions = [] for callee in callees_functions: if callee in all_functions[fname]["func_names"]: - pruned_callees_functions.append(prefix + '.' + callee) + pruned_callees_functions.append(prefix + "." + callee) elif callee in stumpy_functions: idx = stumpy_functions_no_prefix.index(callee) pruned_callees_functions.append(stumpy_functions[idx]) @@ -59,11 +58,3 @@ continue callers_callees[func_name] = (is_njit, fastmath_val, pruned_callees_functions) - - - - -for k, v in callers_callees.items(): - print('---------------------------------') - print('> func_name: ', k) - print('metadata: ', v) \ No newline at end of file From c51b7a32e25e6ec41105566179b8d139e07d7842 Mon Sep 17 00:00:00 2001 From: NimaSarajpoor Date: Sat, 28 Sep 2024 21:11:37 -0400 Subject: [PATCH 7/9] add code to obtain callees_callers dictionary --- njit_fastmath.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/njit_fastmath.py b/njit_fastmath.py index 7a03110f0..495de4cde 100644 --- a/njit_fastmath.py +++ b/njit_fastmath.py @@ -58,3 +58,21 @@ continue callers_callees[func_name] = (is_njit, fastmath_val, pruned_callees_functions) + + +# Create callees_callers dictionary using callers_callees dictionary +callees_callers = {} +for func_name, func_metadata in callers_callees.items(): + callees_callers[func_name] = [func_metadata[0], func_metadata[1], []] + + +for func_name, func_metadata in callers_callees.items(): + for callee in func_metadata[2]: + callees_callers[callee][-1].append(func_name) + + +for func_name, func_metadata in callees_callers.items(): + callees_callers[func_name][2] = set(callees_callers[func_name][2]) + callees_callers[func_name] = tuple(callees_callers[func_name]) + + From 8464b704e8ea2c2481c987e176296d8935f44ea9 Mon Sep 17 00:00:00 2001 From: NimaSarajpoor Date: Sat, 28 Sep 2024 21:12:09 -0400 Subject: [PATCH 8/9] fix format --- njit_fastmath.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/njit_fastmath.py b/njit_fastmath.py index 495de4cde..4f13e4616 100644 --- a/njit_fastmath.py +++ b/njit_fastmath.py @@ -62,7 +62,7 @@ # Create callees_callers dictionary using callers_callees dictionary callees_callers = {} -for func_name, func_metadata in callers_callees.items(): +for func_name, func_metadata in callers_callees.items(): callees_callers[func_name] = [func_metadata[0], func_metadata[1], []] @@ -74,5 +74,3 @@ for func_name, func_metadata in callees_callers.items(): callees_callers[func_name][2] = set(callees_callers[func_name][2]) callees_callers[func_name] = tuple(callees_callers[func_name]) - - From f548907d051d766bafb560411928b1b0bb9840a6 Mon Sep 17 00:00:00 2001 From: NimaSarajpoor Date: Sat, 28 Sep 2024 23:02:20 -0400 Subject: [PATCH 9/9] revise script to get two dictionarys with one being callee to callers map --- njit_fastmath.py | 86 +++++++++++++++++++++++------------------------- 1 file changed, 41 insertions(+), 45 deletions(-) diff --git a/njit_fastmath.py b/njit_fastmath.py index 4f13e4616..9e46c5e89 100644 --- a/njit_fastmath.py +++ b/njit_fastmath.py @@ -6,71 +6,67 @@ filepaths = sorted(f for f in pathlib.Path(stumpy_path).iterdir() if f.is_file()) all_functions = {} -callees = {} ignore = ["__init__.py", "__pycache__"] for filepath in filepaths: - name = filepath.name - if name not in ignore and str(filepath).endswith(".py"): - callees[name] = check_callees(filepath) + file_name = filepath.name + if file_name not in ignore and str(filepath).endswith(".py"): + prefix = file_name.replace(".py", "") func_names, is_njit, fastmath_values = check_functions(filepath) - all_functions[name] = { + func_names = [f"{prefix}.{fn}" for fn in func_names] + + all_functions[file_name] = { "func_names": func_names, "is_njit": is_njit, "fastmath_values": fastmath_values, } +all_stumpy_functions = set() +for file_name, file_functions_metadata in all_functions.items(): + all_stumpy_functions.update(file_functions_metadata["func_names"]) -stumpy_functions = set() -for fname, func_data in all_functions.items(): - prefix = fname.replace(".py", "") - stumpy_functions.update([prefix + "." + x for x in func_data["func_names"]]) -stumpy_functions = list(stumpy_functions) -stumpy_functions_no_prefix = [x.split(".")[1] for x in stumpy_functions] +all_stumpy_functions = list(all_stumpy_functions) +all_stumpy_functions_no_prefix = [f.split(".")[-1] for f in all_stumpy_functions] -# create a dictionary where keys are function names in stumpy_functions, and the value -# is a tuple, where the first element says whether the function is decorated with njit -# and the second element is the list of callees +# output 1: func_metadata +func_metadata = {} +for file_name, file_functions_metadata in all_functions.items(): + for i, f in enumerate(file_functions_metadata["func_names"]): + is_njit = file_functions_metadata["is_njit"][i] + fastmath_value = file_functions_metadata["fastmath_values"][i] + func_metadata[f] = [is_njit, fastmath_value] -callers_callees = {} -for func_name in stumpy_functions: - callers_callees[func_name] = None - prefix, func = func_name.split(".") - fname = prefix + ".py" +# output 2: func_callers +func_callers = {} +for f in func_metadata.keys(): + func_callers[f] = [] - idx = all_functions[fname]["func_names"].index(func) - is_njit = all_functions[fname]["is_njit"][idx] - fastmath_val = all_functions[fname]["fastmath_values"][idx] +for filepath in filepaths: + file_name = filepath.name + if file_name in ignore or not str(filepath).endswith(".py"): + continue - callees_functions = callees[fname][func] + prefix = file_name.replace(".py", "") + callees = check_callees(filepath) - pruned_callees_functions = [] - for callee in callees_functions: - if callee in all_functions[fname]["func_names"]: - pruned_callees_functions.append(prefix + "." + callee) - elif callee in stumpy_functions: - idx = stumpy_functions_no_prefix.index(callee) - pruned_callees_functions.append(stumpy_functions[idx]) - else: + current_callers = set(callees.keys()) + for caller, callee_set in callees.items(): + s = list(callee_set.intersection(all_stumpy_functions_no_prefix)) + if len(s) == 0: continue - callers_callees[func_name] = (is_njit, fastmath_val, pruned_callees_functions) - - -# Create callees_callers dictionary using callers_callees dictionary -callees_callers = {} -for func_name, func_metadata in callers_callees.items(): - callees_callers[func_name] = [func_metadata[0], func_metadata[1], []] - + for c in s: + if c in current_callers: + c_name = prefix + "." + c + else: + idx = all_stumpy_functions_no_prefix.index(c) + c_name = all_stumpy_functions[idx] -for func_name, func_metadata in callers_callees.items(): - for callee in func_metadata[2]: - callees_callers[callee][-1].append(func_name) + func_callers[c_name].append(f"{prefix}.{caller}") -for func_name, func_metadata in callees_callers.items(): - callees_callers[func_name][2] = set(callees_callers[func_name][2]) - callees_callers[func_name] = tuple(callees_callers[func_name]) +for f, callers in func_callers.items(): + func_callers[f] = list(set(callers))