decomposition.F90

Go to the documentation of this file.

module decomposition_mod
  use datadefn_mod, only : string_length
  use monc_component_mod, only : component_descriptor_type
  use state_mod, only : model_state_type
  use grids_mod, only : local_grid_type, x_index, y_index, z_index
  use collections_mod, only : map_type
  use conversions_mod, only : conv_to_string
  use logging_mod, only : log_info, log_debug, log_warn, log_get_logging_level,&
       log_master_log, log_log
  use optionsdatabase_mod, only : options_get_string, options_get_integer
  use mpi
  implicit none

#ifndef TEST_MODE
  private
#endif

  public decomposition_get_descriptor

contains

  type(component_descriptor_type) function decomposition_get_descriptor()
    decomposition_get_descriptor%name="decomposition"
    decomposition_get_descriptor%version=0.1
    decomposition_get_descriptor%initialisation=>init_callback
  end function decomposition_get_descriptor

  subroutine init_callback(current_state)
    type(model_state_type), target, intent(inout) :: current_state

    character(len=STRING_LENGTH) :: method

    method=options_get_string(current_state%options_database, "decomposition_method")

    if (method .eq. "onedim") then
       current_state%parallel%decomposition_procedure => one_dim_decomposition
    else if (method .eq. "twodim") then
       current_state%parallel%decomposition_procedure => two_dim_decomposition
    else
       current_state%parallel%decomposition_procedure => serial_decomposition
       if (method .ne. "serial") call log_log(log_warn, "Decomposition method "//trim(method)//&
            " not recognised so defaulting to serial")
    end if
  end subroutine init_callback

  subroutine two_dim_decomposition(current_state)
    type(model_state_type), intent(inout) :: current_state

    integer :: ierr
    integer, dimension(2) :: coords, distributed_dims

    current_state%local_grid%active = current_state%global_grid%active
    current_state%local_grid%dimensions = current_state%global_grid%dimensions

    if (current_state%global_grid%dimensions .ne. 3) then
       call log_master_log(log_warn, "Two dimension decomposition selected with only "// &
            trim(conv_to_string(current_state%global_grid%dimensions)) //" so defaulting to one dimension instead")
       call one_dim_decomposition(current_state)
       return
    end if

    distributed_dims= (/0,0/)
    call mpi_dims_create(current_state%parallel%processes, 2, distributed_dims, ierr)

    current_state%parallel%dim_sizes(z_index)=1
    current_state%parallel%dim_sizes(y_index) = distributed_dims(1)
    current_state%parallel%dim_sizes(x_index) = distributed_dims(2)

    if (.not. does_two_dim_work_with_domain_size(current_state, current_state%parallel%dim_sizes(y_index), &
         current_state%parallel%dim_sizes(x_index))) then
       call log_master_log(log_warn, "Defaulting to one dimension decomposition due to solution size too small")
       call one_dim_decomposition(current_state)
       return
    end if
    
    call mpi_cart_create(current_state%parallel%monc_communicator, 2, (/distributed_dims(1), distributed_dims(2)/),&
         (/.false., .false./), .false., current_state%parallel%neighbour_comm, ierr)
    call mpi_cart_coords(current_state%parallel%neighbour_comm, current_state%parallel%my_rank, 2, coords, ierr)

    call apply_halo_information_and_allocate_neighbours(current_state)
    call apply_z_dimension_information(current_state)

    ! sets the limits for each rank in X and Y
    call apply_dimension_bounds(current_state, y_index, merge(current_state%global_grid%size(y_index), 1, &
         current_state%global_grid%active(y_index)), distributed_dims(1), coords(1))
    call apply_dimension_bounds(current_state, x_index, merge(current_state%global_grid%size(x_index), 1, &
         current_state%global_grid%active(x_index)), distributed_dims(2), coords(2))

    ! get all information about the neighbours and corner neighbours
    call apply_two_dim_neighbour_information(current_state, distributed_dims(1), distributed_dims(2), coords)

    call apply_data_start_end_bounds(current_state%local_grid)

    if (log_get_logging_level() .ge. log_debug) then
       call log_log(log_debug, "PID "//trim(conv_to_string(current_state%parallel%my_rank))//": y="//&
            trim(conv_to_string(current_state%parallel%my_coords(y_index)))//&
            " x="//trim(conv_to_string(current_state%parallel%my_coords(x_index)))//" ny-1="//&
            trim(conv_to_string(current_state%local_grid%neighbours(y_index,1)))//&
            " ny+1="//trim(conv_to_string(current_state%local_grid%neighbours(y_index,3)))//" nx-1="//&
            trim(conv_to_string(current_state%local_grid%neighbours(x_index,1)))//" nx+1="//&
            trim(conv_to_string(current_state%local_grid%neighbours(x_index,3))))
    end if
    call display_decomposition_information(current_state, "TwoDim")
  end subroutine two_dim_decomposition

  subroutine apply_two_dim_neighbour_information(current_state, y_procs, x_procs, coords)

    type(model_state_type), intent(inout) :: current_state
    integer, intent(in) :: y_procs, x_procs, coords(2)

    integer :: ierr, i, halo_depth

    current_state%parallel%wrapped_around=.false.

    call mpi_cart_shift(current_state%parallel%neighbour_comm, 0, 1, current_state%local_grid%neighbours(y_index,1), &
         current_state%local_grid%neighbours(y_index,3), ierr)
    if (current_state%local_grid%neighbours(y_index,1) .lt. 0) then
      call mpi_cart_rank(current_state%parallel%neighbour_comm,&
           (/y_procs-1, coords(2)/),  current_state%local_grid%neighbours(y_index,1), ierr)
      current_state%parallel%wrapped_around(y_index, 1)=.true.
    end if
    
    if (current_state%local_grid%neighbours(y_index,3) .lt. 0) then
      call mpi_cart_rank(current_state%parallel%neighbour_comm, &
         (/0, coords(2)/),  current_state%local_grid%neighbours(y_index,3), ierr)
      current_state%parallel%wrapped_around(y_index, 2)=.true.
    end if    

    ! find out who is on the left and on the right
    call mpi_cart_shift(current_state%parallel%neighbour_comm, 1, 1, &
         current_state%local_grid%neighbours(x_index,1), &
         current_state%local_grid%neighbours(x_index,3), ierr)
    if (current_state%local_grid%neighbours(x_index,1) .lt. 0) then
      call mpi_cart_rank(current_state%parallel%neighbour_comm, &
         (/coords(1), x_procs-1/),  current_state%local_grid%neighbours(x_index,1), ierr)
      current_state%parallel%wrapped_around(x_index, 1)=.true.
    end if
    if (current_state%local_grid%neighbours(x_index,3) .lt. 0) then
      call mpi_cart_rank(current_state%parallel%neighbour_comm, &
         (/coords(1), 0/),  current_state%local_grid%neighbours(x_index,3), ierr)
      current_state%parallel%wrapped_around(x_index, 2)=.true.
    end if

    ! the second column is on the same rank than the first one
    halo_depth = options_get_integer(current_state%options_database, "halo_depth")
    
    do i = 1, halo_depth -1
       current_state%local_grid%neighbours(y_index,i+1) = &
            current_state%local_grid%neighbours(y_index,i)
       current_state%local_grid%neighbours(y_index,i+halo_depth + 1) = &
            current_state%local_grid%neighbours(y_index,i+halo_depth)
       current_state%local_grid%neighbours(x_index,i+1) = &
            current_state%local_grid%neighbours(x_index,i)
       current_state%local_grid%neighbours(x_index,i+halo_depth + 1) = &
            current_state%local_grid%neighbours(x_index,i+halo_depth)
    end do
    ! obtain the rank of the corner neighbours
    call mpi_cart_rank(current_state%parallel%neighbour_comm, &
         (/merge(coords(1)-1, y_procs-1, coords(1) .ge. 1),   &
         merge(coords(2)-1, x_procs-1, coords(2) .ge. 1)/),   &
         current_state%local_grid%corner_neighbours(1,1), ierr)
    call mpi_cart_rank(current_state%parallel%neighbour_comm, &
         (/merge(coords(1)-1, y_procs-1, coords(1) .ge. 1),   &
         merge(coords(2)+1, 0, coords(2) .lt. x_procs-1)/),   &
         current_state%local_grid%corner_neighbours(2,1), ierr)
    call mpi_cart_rank(current_state%parallel%neighbour_comm, &
         (/merge(coords(1)+1, 0, coords(1) .lt. y_procs-1),   &
         merge(coords(2)-1, x_procs-1, coords(2) .ge. 1)/),   &
         current_state%local_grid%corner_neighbours(3,1), ierr)
    call mpi_cart_rank(current_state%parallel%neighbour_comm, &
         (/merge(coords(1)+1, 0, coords(1) .lt. y_procs-1),   &
         merge(coords(2)+1, 0, coords(2) .lt. x_procs-1)/),   &
         current_state%local_grid%corner_neighbours(4,1), ierr)
    
    !! TODO: hardcoded for halo depth two?
    current_state%local_grid%corner_neighbours(:,2)=current_state%local_grid%corner_neighbours(:,1)
  end subroutine apply_two_dim_neighbour_information

  logical function does_two_dim_work_with_domain_size(current_state, y_dims, x_dims)
    type(model_state_type), intent(inout) :: current_state
    integer, intent(in) :: y_dims, x_dims

    if (current_state%global_grid%active(y_index)) then
       if (floor(real(current_state%global_grid%size(Y_INDEX)) / y_dims) .lt. 2) then
          does_two_dim_work_with_domain_size=.false.
          return
       end if
    end if
    if (current_state%global_grid%active(x_index)) then
       if (floor(real(current_state%global_grid%size(X_INDEX)) / x_dims) .lt. 2) then
          does_two_dim_work_with_domain_size=.false.
          return
       end if
    end if
    does_two_dim_work_with_domain_size=.true.
  end function does_two_dim_work_with_domain_size

  subroutine apply_dimension_bounds(current_state, dim, dim_size, dim_processes, dim_my_rank)
    type(model_state_type), intent(inout) :: current_state
    integer, intent(in) :: dim, dim_size, dim_my_rank, dim_processes

    integer :: dimension_division, dimension_extra

    dimension_division = dim_size / dim_processes
    dimension_extra = dim_size - (dimension_division * dim_processes)

    current_state%local_grid%start(dim) = dimension_division*dim_my_rank + merge(dimension_extra, dim_my_rank, &
         dimension_extra .lt. dim_my_rank) + 1
    current_state%local_grid%end(dim) = (current_state%local_grid%start(dim)-1) + dimension_division + &
         merge(1, 0, dim_my_rank .lt. dimension_extra)
    current_state%local_grid%size(dim)=(current_state%local_grid%end(dim) - current_state%local_grid%start(dim)) + 1
    current_state%parallel%my_coords(dim) = dim_my_rank
  end subroutine apply_dimension_bounds

  subroutine apply_z_dimension_information(current_state)
    type(model_state_type), intent(inout) :: current_state

    current_state%local_grid%start(z_index) = 1
    current_state%local_grid%end(z_index) = current_state%global_grid%size(z_index)
    current_state%local_grid%size(z_index) = current_state%global_grid%size(z_index)
    current_state%parallel%my_coords(z_index) = 0
    current_state%parallel%dim_sizes(z_index) = 1
  end subroutine apply_z_dimension_information

  subroutine one_dim_decomposition(current_state)
    type(model_state_type), intent(inout) :: current_state

    integer :: x_size, y_size, split_size, dimension_division, dimension_extra

    call apply_halo_information_and_allocate_neighbours(current_state)

    current_state%parallel%wrapped_around=.false.

    x_size = merge(current_state%global_grid%size(x_index), 1, current_state%global_grid%active(x_index))
    y_size = merge(current_state%global_grid%size(y_index), 1, current_state%global_grid%active(y_index))

    split_size = merge(x_size, y_size, x_size .gt. y_size)
    dimension_division = split_size / current_state%parallel%processes
    dimension_extra = split_size - (dimension_division * current_state%parallel%processes)

    current_state%local_grid%active = current_state%global_grid%active
    current_state%local_grid%dimensions = current_state%global_grid%dimensions
    current_state%parallel%neighbour_comm = current_state%parallel%monc_communicator
    call apply_z_dimension_information(current_state)         

    if (x_size .gt. y_size) then
      ! Decompose in X
      current_state%local_grid%start(y_index)=1
      current_state%local_grid%end(y_index)=y_size
      current_state%local_grid%size(y_index)=y_size
      current_state%parallel%my_coords(y_index)=0
      current_state%parallel%dim_sizes(y_index)=1
      current_state%local_grid%start(x_index)=dimension_division*current_state%parallel%my_rank+merge(&
           dimension_extra, current_state%parallel%my_rank, dimension_extra .lt. current_state%parallel%my_rank) + 1
      current_state%local_grid%end(x_index)=(current_state%local_grid%start(x_index)-1) + dimension_division + merge(&
           1, 0, current_state%parallel%my_rank .lt. dimension_extra)
      current_state%local_grid%size(x_index)=(current_state%local_grid%end(x_index) - current_state%local_grid%start(x_index)) + 1
      current_state%parallel%my_coords(x_index) = current_state%parallel%my_rank
      current_state%parallel%dim_sizes(x_index) = current_state%parallel%processes
      current_state%local_grid%neighbours(y_index,:) = current_state%parallel%my_rank
      ! Currently assume same PID in single direction (TODO: relax this restraint)
      current_state%local_grid%neighbours(x_index,1:2) = merge(current_state%parallel%my_rank, current_state%parallel%processes, &
           current_state%parallel%my_rank .gt. 0) - 1
      current_state%parallel%wrapped_around(x_index, 1)=&
           current_state%local_grid%neighbours(x_index,1)==current_state%parallel%processes-1
      current_state%local_grid%neighbours(x_index,3:4) = merge(current_state%parallel%my_rank+1, 0, &
           current_state%parallel%my_rank .lt. current_state%parallel%processes-1)
      current_state%parallel%wrapped_around(x_index, 2)=current_state%local_grid%neighbours(x_index,3)==0
      current_state%local_grid%corner_neighbours(1,:)=current_state%local_grid%neighbours(x_index,1)
      current_state%local_grid%corner_neighbours(3,:)=current_state%local_grid%neighbours(x_index,1)
      current_state%local_grid%corner_neighbours(2,:)=current_state%local_grid%neighbours(x_index,3)
      current_state%local_grid%corner_neighbours(4,:)=current_state%local_grid%neighbours(x_index,3)
      current_state%parallel%wrapped_around(y_index, :)=.true.
    else 
      ! Decompose in Y
      current_state%local_grid%start(x_index)=1
      current_state%local_grid%end(x_index)=x_size
      current_state%local_grid%size(x_index)=x_size
      current_state%parallel%my_coords(x_index)=0
      current_state%parallel%dim_sizes(x_index)=1
      current_state%local_grid%start(y_index)=dimension_division*current_state%parallel%my_rank+merge(&
           dimension_extra, current_state%parallel%my_rank, dimension_extra .lt. current_state%parallel%my_rank) + 1
      current_state%local_grid%end(y_index)=(current_state%local_grid%start(y_index)-1) + dimension_division + merge(&
           1, 0, current_state%parallel%my_rank .lt. dimension_extra)
      current_state%local_grid%size(y_index)=(current_state%local_grid%end(y_index) - current_state%local_grid%start(y_index)) + 1
      current_state%parallel%my_coords(y_index)=current_state%parallel%my_rank
      current_state%parallel%dim_sizes(y_index) = current_state%parallel%processes
      current_state%local_grid%neighbours(x_index,:) = current_state%parallel%my_rank
      current_state%local_grid%neighbours(y_index,1:2) = merge(current_state%parallel%my_rank, current_state%parallel%processes, &
           current_state%parallel%my_rank .gt. 0) - 1
      current_state%parallel%wrapped_around(y_index, 1)=&
           current_state%local_grid%neighbours(y_index,1)==current_state%parallel%processes-1
      current_state%local_grid%neighbours(y_index,3:4) = merge(current_state%parallel%my_rank+1, 0, &
           current_state%parallel%my_rank .lt. current_state%parallel%processes-1)
      current_state%parallel%wrapped_around(y_index, 2)=current_state%local_grid%neighbours(y_index,3)==0
      current_state%local_grid%corner_neighbours(1:2,:)=current_state%local_grid%neighbours(y_index,1)
      current_state%local_grid%corner_neighbours(3:4,:)=current_state%local_grid%neighbours(y_index,3)
      current_state%parallel%wrapped_around(x_index, :)=.true.
    end if

    call apply_data_start_end_bounds(current_state%local_grid)
    call display_decomposition_information(current_state, "OneDim")
  end subroutine one_dim_decomposition

  subroutine serial_decomposition(current_state)
    type(model_state_type), intent(inout) :: current_state

    call apply_halo_information_and_allocate_neighbours(current_state)

    current_state%local_grid%active = current_state%global_grid%active
    current_state%local_grid%dimensions = current_state%global_grid%dimensions
    current_state%parallel%neighbour_comm = current_state%parallel%monc_communicator
    call apply_z_dimension_information(current_state)

    current_state%local_grid%start(y_index)=1
    current_state%local_grid%start(x_index)=1
    current_state%parallel%my_coords(y_index)=0
    current_state%parallel%my_coords(x_index)=0
    current_state%parallel%dim_sizes(y_index)=0
    current_state%parallel%dim_sizes(x_index)=0

    current_state%local_grid%end(x_index)= &
         merge(current_state%global_grid%size(x_index), 1, current_state%global_grid%active(x_index))
    current_state%local_grid%size(x_index)=current_state%local_grid%end(x_index)

    current_state%local_grid%end(y_index)=&
         merge(current_state%global_grid%size(y_index), 1, current_state%global_grid%active(y_index))
    current_state%local_grid%size(y_index)=current_state%local_grid%end(y_index)

    current_state%local_grid%neighbours(x_index,:) = current_state%parallel%my_rank
    current_state%local_grid%neighbours(y_index,:) = current_state%parallel%my_rank
    current_state%local_grid%corner_neighbours=current_state%parallel%my_rank

    current_state%parallel%wrapped_around=.true.

    call apply_data_start_end_bounds(current_state%local_grid)
    call display_decomposition_information(current_state, "Serial")
  end subroutine serial_decomposition

  subroutine display_decomposition_information(current_state, decomp_name)
    type(model_state_type), intent(inout) :: current_state
    character(len=*), intent(in) :: decomp_name

    if (log_get_logging_level() .le. log_debug) then
       call log_log(log_debug, decomp_name//" rank="//trim(conv_to_string(current_state%parallel%my_rank))//&
            ": W=("//trim(conv_to_string(current_state%local_grid%start(z_index)))//"->"&
            //trim(conv_to_string(current_state%local_grid%end(z_index)))//") V=("//trim(conv_to_string(&
            current_state%local_grid%start(y_index)))//"->"//trim(conv_to_string(current_state%local_grid%end(&
            y_index)))//") U=("//trim(conv_to_string(current_state%local_grid%start(x_index)))//&
            "->"//trim(conv_to_string(current_state%local_grid%end(x_index)))//")")
       call log_log(log_debug, "Neighbours of "//trim(conv_to_string(current_state%parallel%my_rank))//": y=(["//&
            trim(conv_to_string(current_state%local_grid%neighbours(y_index,1)))//"],["//&
            trim(conv_to_string(current_state%local_grid%neighbours(y_index,2)))//"]) x=(["//&
            trim(conv_to_string(current_state%local_grid%neighbours(x_index,1)))//","//&
            trim(conv_to_string(current_state%local_grid%neighbours(x_index,2)))//"],["//&
            trim(conv_to_string(current_state%local_grid%neighbours(x_index,3)))//","//&
            trim(conv_to_string(current_state%local_grid%neighbours(x_index,4)))//"])")
       call log_master_log(log_debug, "Halo size z="//&
            trim(conv_to_string(current_state%local_grid%halo_size(z_index)))//&
            " y="//trim(conv_to_string(current_state%local_grid%halo_size(y_index)))//" x="//&
            trim(conv_to_string(current_state%local_grid%halo_size(x_index))))
    end if
    call log_master_log(log_info, "Decomposed "//trim(conv_to_string(current_state%parallel%processes))//&
         " processes via '"//decomp_name// "' into z="//trim(conv_to_string(current_state%parallel%dim_sizes(&
         z_index)))//" y="//trim(conv_to_string(current_state%parallel%dim_sizes(y_index)))//" x="//&
         trim(conv_to_string(current_state%parallel%dim_sizes(x_index))))
  end subroutine display_decomposition_information

  subroutine apply_data_start_end_bounds(local_grid)
    type(local_grid_type), intent(inout) :: local_grid

    integer :: i,n_dim

    !Total number of dimensions
    n_dim = 3

    do i = 1,n_dim
       local_grid%local_domain_start_index(i) = local_grid%halo_size(i) + 1
       local_grid%local_domain_end_index(i) = local_grid%halo_size(i) + local_grid%size(i)
    end do
    ! call log_master_log(LOG_DEBUG,"start_index(1)= "//&
    !      trim(conv_to_string(local_grid%local_domain_start_index(1)))//" end_index(1)="//&
    !      trim(conv_to_string(local_grid%local_domain_end_index(1)))//" start_index(2)="//&
    !      trim(conv_to_string(local_grid%local_domain_start_index(2)))//" end_index(2)="//&
    !      trim(conv_to_string(local_grid%local_domain_end_index(2)))//" start_index(3)="//&
    !      trim(conv_to_string(local_grid%local_domain_start_index(3)))//" end_index(3)="//&
    !      trim(conv_to_string(local_grid%local_domain_end_index(3))))

  end subroutine apply_data_start_end_bounds

  subroutine apply_halo_information_and_allocate_neighbours(current_state)
    type(model_state_type), target, intent(inout) :: current_state
    integer :: n_dim, n_corners,total_halo_size_XY_dim
    integer :: halo_depth

    ! Read halo_depth value
    halo_depth = options_get_integer(current_state%options_database, "halo_depth")

    ! There is no halo_depth in Z direction
    current_state%local_grid%halo_size(z_index) = 0
    current_state%local_grid%halo_size(y_index) = halo_depth
    current_state%local_grid%halo_size(x_index) = halo_depth

    ! Left and right or Up and down
    total_halo_size_xy_dim  = current_state%local_grid%halo_size(x_index)*2
    ! Total number of dimensions
    n_dim = 3
    ! Total number of corners
    n_corners = 4
    allocate(current_state%local_grid%neighbours(n_dim,total_halo_size_xy_dim), &
         current_state%local_grid%corner_neighbours(n_corners,halo_depth))
  end subroutine apply_halo_information_and_allocate_neighbours
end module decomposition_mod