I implemented @ztik suggestion and this is the result I came up with. It seems to work:
program main
use mpi
implicit none
integer :: ierr, me, nProcs, color, i,j,k, my_comm, dims(2), global_contxt
integer :: cnt, n_colors, map(2,2)
integer, allocatable :: contxts(:)
integer, parameter :: group_size = 4
call MPI_Init(ierr)
call MPI_Comm_size(MPI_COMM_WORLD, nProcs, ierr)
call MPI_Comm_rank(MPI_COMM_WORLD, me, ierr)
color = get_color(group_size)
n_colors = nProcs / group_size
allocate(contxts(n_colors))
dims = calc_2d_dim(group_size)
call BLACS_GET(0, 0, global_contxt)
if(me == 0) write (*,*) global_contxt
contxts = global_contxt
do k = 1,n_colors
! shift for each context
cnt = group_size * (k-1)
if(me==0) write (*,*) "##############", cnt
! create map
do i=1,2
do j=1,2
map(i,j) = cnt
cnt = cnt + 1
enddo
enddo
call BLACS_GRIDMAP(contxts(k), map, 2, 2, 2)
do i = 0,nProcs
if(i == me) then
write (*,*) me, contxts(k)
endif
call MPI_Barrier(MPI_COMM_WORLD, ierr)
enddo
enddo
call MPI_Finalize(ierr)
contains
function get_color(group_size) result(color)
implicit none
integer, intent(in) :: group_size
integer :: me, nProcs, color, ierr, i, cnt
call MPI_Comm_size(MPI_COMM_WORLD, nProcs, ierr)
call MPI_Comm_rank(MPI_COMM_WORLD, me, ierr)
if(mod(nProcs, group_size) /= 0) then
write (*,*) "Nprocs not divisable by group_size", mod(nProcs, group_size)
call MPI_Abort(MPI_COMM_WORLD, 0, ierr)
endif
color = 0
do i = 1,me
if(mod(i, group_size) == 0) then
color = color + 1
endif
enddo
end function get_color
function calc_2d_dim(sz) result(dim)
implicit none
integer, intent(in) :: sz
integer :: dim(2), cand
cand = nint(sqrt(real(sz)))
do while(mod(sz, cand) /= 0)
cand = cand - 1
enddo
dim(1) = sz/cand
dim(2) = cand
end function calc_2d_dim
end program main