Divide search for decomposition to be able to obtain a list of possible decompositions.

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}")
+    N = 128
+    dx, dy, dz = get_best_3d_decomposition(N)
+    print(f"Optimal decomposition: Dx = {dx}, Dy = {dy}, Dz = {dz}")
\ No newline at end of file
+    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