diff --git a/python/waLBerla/tools/config/setup.py b/python/waLBerla/tools/config/setup.py index dc8f1ab282e1b78b661ef38e423ce755b98e238f..9b7034b317387ab64c664a4f7a8fb86b68e98d03 100644 --- a/python/waLBerla/tools/config/setup.py +++ b/python/waLBerla/tools/config/setup.py @@ -21,18 +21,18 @@ def block_decomposition(processes): return tuple(result) -def decompose_into_3d(processes: int, search_range: int = 1) -> Tuple[int, int, int]: +def decompose_into_3d(processes: int, search_range: int = 1, verbose: bool = False) -> Tuple[Tuple[int, int, int]]: """ Decomposes the number of processes 'processes' into three dimensions (Dx, Dy, Dz) such that: - Dx * Dy * Dz == processes - - The decomposition is as balanced as possible (minimizing the imbalance between dimensions). + - The possible decomposition are sorted to priortize the once that are as balanced as possible. Parameters: processes (int): The total number of processes to decompose. search_range (int): The range of values around the cube root to search for possible decompositions. Returns: - Tuple[int, int, int]: A tuple representing the dimensions (Dx, Dy, Dz). + Tuple[Tuple[int, int, int]]: All tuples representing the dimensions (Dx, Dy, Dz). NOTE: First draft and comments were generated using ChatGBT. Code was improved by hand and with AI interaction afterwards. """ @@ -42,31 +42,49 @@ def decompose_into_3d(processes: int, search_range: int = 1) -> Tuple[int, int, # Step 2: Gradually increase the search range if no valid decomposition is found while True: possible_decompositions = ( - (dx, dy, dz) + (dx, dy, dz) # Ensure the order: Dx <= Dy <= Dz for dx, dy, dz in product( range(cube_root - search_range, cube_root + search_range + 1), repeat=3 ) if dx > 0 and dy > 0 and dz > 0 and dx * dy * dz == processes ) - try: - # Try to get the next valid decomposition - best_decomposition = min( - possible_decompositions, key=lambda dims: max(dims) - min(dims) + # Sort the available decompositions to priotize the most uniform ones + decomposition_priority = sorted( + possible_decompositions, key=lambda dims: max(dims) - min(dims) ) + if decomposition_priority: + if verbose: + print(f"A search range of {search_range} was required to find a valid decomposition.") + return decomposition_priority + + # If no valid decomposition found, increase the search range and try again + search_range *= 2 + - # Ensure the order: Dx <= Dy <= Dz - best_decomposition = tuple(sorted(best_decomposition)) - return best_decomposition +def get_best_3d_decomposition(processes: int) -> Tuple[int, int, int]: + """ + Decomposes the number of processes 'processes' into three dimensions (Dx, Dy, Dz) such that: + - Dx * Dy * Dz == processes + - The decomposition is as balanced as possible (minimizing the distance between dimensions). - except ValueError: - # If no valid decomposition found, increase the search range and try again - search_range *= 2 + Parameters: + processes (int): The total number of processes to decompose. + Returns: + Tuple[int, int, int]: A tuple representing the dimensions (Dx, Dy, Dz). + """ + return decompose_into_3d(processes, search_range=1)[0] if __name__ == "__main__": # Example usage - N = 112 * 4 - dx, dy, dz = decompose_into_3d(N) - print(f"Optimal decomposition: Dx = {dx}, Dy = {dy}, Dz = {dz}") \ No newline at end of file + N = 128 + dx, dy, dz = get_best_3d_decomposition(N) + print(f"Optimal decomposition: Dx = {dx}, Dy = {dy}, Dz = {dz}") + options = decompose_into_3d(N, verbose=True) + print(f"Possible decomposition for initial search range 1: {options}") + options = decompose_into_3d(N, 5, verbose=True) + print(f"Possible decomposition for initial search range 2: {options}") + options = decompose_into_3d(N, 10, verbose=True) + print(f"Possible decomposition for initial search range 3: {options}") \ No newline at end of file