Implement change team and child heap management

We maintain that child teams can use the "remaining" memory
from the current team's heap, which requires some addtional
overhead during allocation and deallocation of coarrays.

Author: everythingfunctional

src/caffeine/allocation_s.f90

@@ -28,6 +28,12 @@
     coarray_size = product(ubounds-lbounds+1)*element_size
 
     me = caf_this_image(current_team%info%gex_team)
+    if (caf_have_child_teams()) then
+      ! Free the child team space to make sure we have space to allocate the coarray
+      if (me == 1) then
+        call caf_deallocate(current_team%info%heap_mspace, current_team%info%child_heap_info%allocated_memory)
+      end if
+    end if
     if (me == 1) then
       handle_size = c_sizeof(unused)
       total_size = handle_size + coarray_size
@@ -52,6 +58,9 @@
     call add_to_team_list(current_team, coarray_handle)
 
     allocated_memory = coarray_handle%info%coarray_data
+    if (caf_have_child_teams()) then
+      call caf_establish_child_heap
+    end if
   end procedure
 
   module procedure prif_allocate
@@ -118,6 +127,13 @@
         call caf_deallocate(current_team%info%heap_mspace, c_loc(coarray_handles(i)%info))
     end do
     if (present(stat)) stat = 0
+    if (caf_have_child_teams()) then
+      ! reclaim any free space possible for the child teams to use
+      if (caf_this_image(current_team%info%gex_team) == 1) then
+        call caf_deallocate(current_team%info%heap_mspace, current_team%info%child_heap_info%allocated_memory)
+      end if
+      call caf_establish_child_heap
+    end if
   end procedure
 
   module procedure prif_deallocate

src/caffeine/caffeine.c

@@ -104,11 +104,28 @@ void* caf_allocate(mspace heap, size_t bytes)
    return mspace_memalign(heap, 8, bytes);
 }
 
+void* caf_allocate_remaining(mspace heap, void** allocated_space, size_t* allocated_size)
+{
+  // The following doesn't necessarily give us all remaining space
+  // nor necessarily the largest open space, but in practice is likely
+  // to work out that way
+  struct mallinfo heap_info = mspace_mallinfo(heap);
+  *allocated_size = heap_info.keepcost;
+
+  *allocated_space = mspace_memalign(heap, 8, *allocated_size);
+}
+
 void caf_deallocate(mspace heap, void* mem)
 {
   mspace_free(heap, mem);
 }
 
+void caf_establish_mspace(mspace* heap, void* heap_start, size_t heap_size)
+{
+  *heap = create_mspace_with_base(heap_start, heap_size, 0);
+  mspace_set_footprint_limit(*heap, heap_size);
+}
+
 // take address in a segment and convert to an address on given image
 intptr_t caf_convert_base_addr(void* addr, int image)
 {

src/caffeine/prif_private_s.f90

@@ -68,13 +68,29 @@ function caf_allocate(mspace, bytes) result(ptr) bind(c)
        type(c_ptr) :: ptr
     end function
 
+    subroutine caf_allocate_remaining(mspace, allocated_space, allocated_size) bind(c)
+      import c_size_t, c_ptr
+      implicit none
+      type(c_ptr), intent(in), value :: mspace
+      type(c_ptr), intent(out) :: allocated_space
+      integer(c_size_t), intent(out) :: allocated_size
+    end subroutine
+
     subroutine caf_deallocate(mspace, mem) bind(c)
       import c_ptr
       implicit none
       type(c_ptr), intent(in), value :: mspace
       type(c_ptr), intent(in), value :: mem
     end subroutine
 
+    subroutine caf_establish_mspace(mspace, mem, mem_size)
+      import c_size_t, c_ptr
+      implicit none
+      type(c_ptr), intent(out) :: mspace
+      type(c_ptr), intent(in), value :: mem
+      integer(c_size_t), intent(in), value :: mem_size
+    end subroutine
+
     ! ___________________ PRIF Queries ______________________
 
     module function caf_convert_base_addr(addr, image) result(ptr) bind(c)
@@ -253,4 +269,21 @@ pure function optional_value(var) result(c_val)
     end if
   end function
 
+  subroutine caf_establish_child_heap
+    if (caf_this_image(current_team%info%gex_team) == 1) then
+      call caf_allocate_remaining( &
+          current_team%info%heap_mspace, &
+          current_team%info%child_heap_info%allocated_memory, &
+          current_team%info%child_heap_info%size)
+      current_team%info%child_heap_info%offset = &
+          as_int(current_team%info%child_heap_info%allocated_memory) - current_team%info%heap_start
+    end if
+    call prif_co_broadcast(current_team%info%child_heap_info%offset, 1)
+    call prif_co_broadcast(current_team%info%child_heap_info%size, 1)
+  end subroutine
+
+  logical function caf_have_child_teams()
+    caf_have_child_teams = associated(current_team%info%child_heap_info)
+  end function
+
 end submodule prif_private_s

src/caffeine/teams_s.f90

@@ -1,47 +1,35 @@
 ! Copyright (c), The Regents of the University of California
 ! Terms of use are as specified in LICENSE.txt
 submodule(prif:prif_private_s) teams_s
-  use iso_c_binding, only: c_null_funptr, c_f_pointer
+  use iso_c_binding, only: c_null_funptr, c_f_pointer, c_loc
 
   implicit none
 contains
 
   module procedure prif_change_team
-    call unimplemented("prif_change_team")
+    if (caf_this_image(team%info%gex_team) == 1) then ! need to setup the heap for the team
+      team%info%heap_start = current_team%info%child_heap_info%offset + current_team%info%heap_start
+      team%info%heap_size = current_team%info%child_heap_info%size
+      call caf_establish_mspace( &
+          team%info%heap_mspace, &
+          as_c_ptr(team%info%heap_start), &
+          current_team%info%child_heap_info%size)
+    end if
+    current_team = team
+    if (caf_have_child_teams()) then ! need to establish heap for child teams
+      call caf_establish_child_heap
+    end if
   end procedure
 
   module procedure prif_end_team
     call unimplemented("prif_end_team")
   end procedure
 
   module procedure prif_form_team
-    if (.not.associated(current_team%info%child_heap_info)) then
-      block
-        type(c_ptr) :: allocated_memory
-        integer :: num_imgs
-        type(child_team_info) :: dummy_element
-
-        call prif_num_images(num_images = num_imgs)
-        call prif_allocate_coarray( &
-            lcobounds = [1_c_intmax_t], &
-            ucobounds = [int(num_imgs, c_intmax_t)], &
-            lbounds = [integer(c_intmax_t)::], &
-            ubounds = [integer(c_intmax_t)::], &
-            element_size = int(storage_size(dummy_element)/8, c_size_t), &
-            final_func = c_null_funptr, &
-            coarray_handle = current_team%info%child_team_handle, &
-            allocated_memory = allocated_memory)
-        call c_f_pointer(allocated_memory, current_team%info%child_heap_info)
-
-        if (caf_this_image(current_team%info%gex_team) == 1) then
-          ! TODO: Take up remaining space and adjust during (de)allocations
-          current_team%info%child_heap_info%size = int(current_team%info%heap_size * 0.1d0, c_size_t)
-          current_team%info%child_heap_info%offset = &
-              as_int(caf_allocate( &
-                  current_team%info%heap_mspace, current_team%info%child_heap_info%size)) &
-              - current_team%info%heap_start
-        end if
-      end block
+    ! indicates this is the first time we're creating a child team
+    if (.not.caf_have_child_teams()) then
+      allocate(current_team%info%child_heap_info)
+      call caf_establish_child_heap
     end if
 
     block

src/prif.f90

@@ -1034,11 +1034,11 @@ module subroutine prif_atomic_ref_logical_indirect(image_num, atom_remote_ptr, v
     integer(c_size_t) :: heap_size
     type(team_data), pointer :: parent_team => null()
     type(handle_data), pointer :: coarrays => null()
-    type(prif_coarray_handle) :: child_team_handle
     type(child_team_info), pointer :: child_heap_info => null()
   end type
 
   type :: child_team_info
+    type(c_ptr) :: allocated_memory
     integer(c_ptrdiff_t) :: offset
     integer(c_size_t) :: size
   end type