Fixed tests and moved checks to root

Author: astroC86

examples/plot_matrixmult.py

@@ -41,8 +41,8 @@
 rank = comm.Get_rank() # rank of current process
 size = comm.Get_size() # number of processes
 
-p_prime = int(math.ceil(math.sqrt(size)))
-repl_factor = int(math.ceil(size / p_prime))
+p_prime     = math.isqrt(size)
+repl_factor = p_prime
 
 if (p_prime * repl_factor) != size:
     print(f"Number of processes must be a square number, provided {size} instead...")
@@ -183,13 +183,27 @@
 
 # Local benchmarks
 y = y.asarray(masked=True)
-if N > 1:
-    y = y.reshape(p_prime, M, blk_cols)
-y = np.hstack([yblock for yblock in y])
+col_counts = [min(blk_cols, N - j * blk_cols) for j in range(p_prime)]
+y_blocks = []
+offset = 0
+for cnt in col_counts:
+    block_size = M * cnt
+    y_blocks.append(
+        y[offset: offset + block_size].reshape(M, cnt)
+    )
+    offset += block_size
+y = np.hstack(y_blocks)
+
 xadj = xadj.asarray(masked=True)
-if N > 1:
-    xadj = xadj.reshape(p_prime, K, blk_cols)
-xadj = np.hstack([xadjblock for xadjblock in xadj])
+xadj_blocks = []
+offset = 0
+for cnt in col_counts:
+    block_size = K * cnt
+    xadj_blocks.append(
+        xadj[offset: offset + block_size].reshape(K, cnt)
+    )
+    offset += block_size
+xadj = np.hstack(xadj_blocks)
 
 if rank == 0:
     y_loc = (A @ X).squeeze()

tests/test_matrixmult.py

@@ -35,9 +35,9 @@ def test_SUMMAMatrixMult(M, K, N, dtype_str):
     cmplx = 1j if np.issubdtype(dtype, np.complexfloating) else 0
     base_float_dtype = np.float32 if dtype == np.complex64 else np.float64
 
-    p_prime = int(math.ceil(math.sqrt(size)))
-    C = int(math.ceil(size / p_prime))
-    assert p_prime * C == size
+    p_prime = math.isqrt(size)
+    C = p_prime
+    assert p_prime * C == size, f"Number of processes must be a square number, provided {size} instead..."
 
     my_group = rank % p_prime
     my_layer = rank // p_prime
@@ -90,25 +90,42 @@ def test_SUMMAMatrixMult(M, K, N, dtype_str):
     # Adjoint operation: xadj = A.H @ y (distributed)
     xadj_dist = Aop.H @ y_dist
 
-    y    = y_dist.asarray(masked=True)
-    y    = y.reshape(p_prime, M, blk_cols_X)
+    y = y_dist.asarray(masked=True)
+    col_counts = [min(blk_cols_X, N - j * blk_cols_X) for j in range(p_prime)]
+    y_blocks = []
+    offset = 0
+    for cnt in col_counts:
+        block_size = M * cnt
+        y_blocks.append(
+            y[offset: offset + block_size].reshape(M, cnt)
+        )
+        offset += block_size
+    y = np.hstack(y_blocks)
 
     xadj = xadj_dist.asarray(masked=True)
-    xadj = xadj.reshape(p_prime, K, blk_cols_X)
+    xadj_blocks = []
+    offset = 0
+    for cnt in col_counts:
+        block_size = K * cnt
+        xadj_blocks.append(
+            xadj[offset: offset + block_size].reshape(K, cnt)
+        )
+        offset += block_size
+    xadj = np.hstack(xadj_blocks)
 
     if rank == 0:
-        y_loc = (A_glob @ X_glob).squeeze()
+        y_loc = A_glob @ X_glob
         assert_allclose(
-            y,
-            y_loc,
+            y.squeeze(),
+            y_loc.squeeze(),
             rtol=np.finfo(np.dtype(dtype)).resolution,
             err_msg=f"Rank {rank}: Forward verification failed."
         )
 
-        xadj_loc = (A_glob.conj().T @ y_loc.conj()).conj().squeeze()
+        xadj_loc = A_glob.conj().T @ y_loc
         assert_allclose(
-            xadj,
-            xadj_loc,
+            xadj.squeeze(),
+            xadj_loc.squeeze(),
             rtol=np.finfo(np.dtype(dtype)).resolution,
             err_msg=f"Rank {rank}: Ajoint verification failed."
         )