For simplicity I’ll remove custom functions etc. My aim is that dividing matrices with dynamically according to the rank size.
For example we have a matrix:

std::vector<std::vector<int>> A = {
      {0, 0, 0, 0}, {1, 1, 1, 1}, {2, 2, 2, 2}, {3, 3, 3, 3}};

if rank size 2, matrix will be divide into:

• Thread 0: Local A = ((0,0,0,0), (1,1,1,1))
• Thread 1: Local A = ((0,0,0,0), (1,1,1,1))
• Thread 2: Local A = ((2,2,2,2), (3,3,3,3))
• Thread 3: Local A = ((2,2,2,2), (3,3,3,3))

my code:

#include <mpi.h>

#include <iostream>
#include <vector>

void print_matrix(const std::vector<std::vector<int>>& mat,
                  const std::string& name, int rank) {
  std::cout << "Process " << rank << " received matrix " << name
            << " with dimensions: " << mat.size() << "x"
            << (mat.empty() ? 0 : mat[0].size()) << std::endl;
  std::cout << "Process " << rank << " printing " << name << ":" << std::endl;
  for (const auto& row : mat) {
    for (int val : row) {
      std::cout << val << " ";
    }
    std::cout << std::endl;
  }
}

int main(int argc, char** argv) {
  MPI_Init(&argc, &argv);
  int rank, size;
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &size);

  int n = 4;                  // Assuming a 4x4 matrix
  int rows_per_proc = n / 2;  // Each process handles half the rows

  MPI_Comm group_comm;
  int group = rank % 2;
  MPI_Comm_split(MPI_COMM_WORLD, group, rank, &group_comm);

  int new_rank;
  MPI_Comm_rank(group_comm, &new_rank);

  std::vector<std::vector<int>> A = {
      {0, 0, 0, 0}, {1, 1, 1, 1}, {2, 2, 2, 2}, {3, 3, 3, 3}};
  std::vector<std::vector<int>> localA(rows_per_proc, std::vector<int>(n));

  // Create a datatype for a block of rows
  MPI_Datatype matrix_block;
  MPI_Type_vector(rows_per_proc, n, n, MPI_INT, &matrix_block);
  MPI_Type_commit(&matrix_block);

  int root = 0;

  // Use MPI_Scatter with the created datatype
  if (new_rank == root) {
    MPI_Scatter(&A[0][0], 1, matrix_block, &localA[0][0], 1, matrix_block, root,
                group_comm);
  } else {
    MPI_Scatter(nullptr, 0, matrix_block, &localA[0][0], 1, matrix_block, root,
                group_comm);
  }

  print_matrix(localA, "Matrix A", rank);

  // Free the custom datatype
  MPI_Type_free(&matrix_block);

  MPI_Finalize();
  return 0;
}

this code’s output is:

Process 0 received matrix Matrix A with dimensions: 2x4
Process 0 printing Matrix A:
0 0 0 0 
0 0 0 0 
Process 1 received matrix Matrix A with dimensions: 2x4
Process 1 printing Matrix A:
0 0 0 0 
0 0 0 0 
Process 2 received matrix Matrix A with dimensions: 2x4
Process 2 printing Matrix A:
1 1 1 1 
0 0 0 0 
Process 3 received matrix Matrix A with dimensions: 2x4
Process 3 printing Matrix A:
1 1 1 1 
0 0 0 0