Merge pull request #243 from bonachea/prif_0.6_dev

Add `prif_initial_team_index` procedures to `caffeine:main`

Author: bonachea

docs/implementation-status.md

@@ -89,6 +89,9 @@ are accepted, but in some cases, the associated runtime behavior is not fully im
 | `prif_image_index`                               | **YES** |  |
 | `prif_image_index_with_team`                     | **YES** |  |
 | `prif_image_index_with_team_number`              | *partial* | no support for sibling teams |
+| `prif_initial_team_index`                        | **YES** | expected in PRIF 0.6 |
+| `prif_initial_team_index_with_team`              | **YES** | expected in PRIF 0.6 |
+| `prif_initial_team_index_with_team_number`       | *partial* | expected in PRIF 0.6, no support for sibling teams |
 
 ---
 

src/caffeine/coarray_queries_s.F90

@@ -102,6 +102,61 @@ subroutine image_index_helper(coarray_handle, sub, num_images, image_index)
     end if 
   end procedure
 
+  !---------------------------------------------------------------------
+
+  subroutine initial_index_helper(coarray_handle, sub, team, initial_team_index)
+    implicit none
+    type(prif_coarray_handle), intent(in) :: coarray_handle
+    integer(c_int64_t), intent(in) :: sub(:)
+    type(prif_team_type), intent(in) :: team
+    integer(c_int), intent(out) :: initial_team_index
+
+    integer :: dim
+    integer(c_int) :: prior_size, image_index
+
+    call_assert(coarray_handle_check(coarray_handle))
+
+    associate (info => coarray_handle%info) 
+      call_assert(size(sub) == info%corank)
+      call_assert(sub(1) .ge. info%lcobounds(1) .and. sub(1) .le. info%ucobounds(1))
+      image_index = 1 + INT(sub(1) - info%lcobounds(1), c_int)
+      prior_size = 1
+      ! Future work: values of prior_size are invariant across calls w/ the same coarray_handle
+      !  We could store them in the coarray metadata at allocation rather than redundantly
+      ! computing them here, which would accelerate calls with corank > 1 by removing
+      ! corank multiply/add operations and the loop-carried dependence
+      do dim = 2, size(sub)
+        prior_size = prior_size * INT(info%ucobounds(dim-1) - info%lcobounds(dim-1) + 1, c_int)
+        call_assert(sub(dim) .ge. info%lcobounds(dim) .and. sub(dim) .le. info%ucobounds(dim))
+        image_index = image_index + INT(sub(dim) - info%lcobounds(dim), c_int) * prior_size
+       end do
+    end associate
+
+    call_assert(image_index .le. team%info%num_images)
+    initial_team_index = caf_image_to_initial(team%info%gex_team, image_index)
+    call_assert(initial_team_index .ge. 1 .and. initial_team_index .le. initial_team%num_images)
+  end subroutine
+
+  module procedure prif_initial_team_index
+    call initial_index_helper(coarray_handle, sub, current_team, initial_team_index)
+  end procedure
+
+  module procedure prif_initial_team_index_with_team
+    call initial_index_helper(coarray_handle, sub, team, initial_team_index)
+  end procedure
+
+  module procedure prif_initial_team_index_with_team_number
+    if (team_number == -1) then
+      call initial_index_helper(coarray_handle, sub, prif_team_type(initial_team), initial_team_index)
+    else if (team_number == current_team%info%team_number) then
+      call initial_index_helper(coarray_handle, sub, current_team, initial_team_index)
+    else
+      call unimplemented("prif_initial_team_index_with_team_number: no support for sibling teams")
+    end if 
+  end procedure
+
+  !---------------------------------------------------------------------
+
   module procedure prif_local_data_pointer
     call_assert(coarray_handle_check(coarray_handle))
 

src/prif.F90

@@ -24,6 +24,7 @@ module prif
   public :: prif_alias_create, prif_alias_destroy
   public :: prif_lcobound_with_dim, prif_lcobound_no_dim, prif_ucobound_with_dim, prif_ucobound_no_dim, prif_coshape
   public :: prif_image_index, prif_image_index_with_team, prif_image_index_with_team_number
+  public :: prif_initial_team_index, prif_initial_team_index_with_team, prif_initial_team_index_with_team_number
   public :: prif_this_image_no_coarray, prif_this_image_with_coarray, prif_this_image_with_dim
   public :: prif_num_images, prif_num_images_with_team, prif_num_images_with_team_number
   public :: prif_failed_images, prif_stopped_images, prif_image_status
@@ -503,6 +504,32 @@ module subroutine prif_image_index_with_team_number(coarray_handle, sub, team_nu
       integer(c_int), intent(out) :: image_index
     end subroutine
 
+    module subroutine prif_initial_team_index(coarray_handle, sub, initial_team_index, stat)
+      implicit none
+      type(prif_coarray_handle), intent(in) :: coarray_handle
+      integer(c_int64_t), intent(in) :: sub(:)
+      integer(c_int), intent(out) :: initial_team_index
+      integer(c_int), intent(out), optional :: stat
+    end subroutine
+    
+    module subroutine prif_initial_team_index_with_team(coarray_handle, sub, team, initial_team_index, stat)
+      implicit none
+      type(prif_coarray_handle), intent(in) :: coarray_handle
+      integer(c_int64_t), intent(in) :: sub(:)
+      type(prif_team_type), intent(in) :: team
+      integer(c_int), intent(out) :: initial_team_index
+      integer(c_int), intent(out), optional :: stat
+    end subroutine
+    
+    module subroutine prif_initial_team_index_with_team_number(coarray_handle, sub, team_number, initial_team_index, stat)
+      implicit none
+      type(prif_coarray_handle), intent(in) :: coarray_handle
+      integer(c_int64_t), intent(in) :: sub(:)
+      integer(c_int64_t), intent(in) :: team_number
+      integer(c_int), intent(out) :: initial_team_index
+      integer(c_int), intent(out), optional :: stat
+    end subroutine
+
     module subroutine prif_num_images(num_images)
       implicit none
       integer(c_int), intent(out) :: num_images

test/prif_image_index_test.F90

@@ -6,9 +6,10 @@ module caf_image_index_test
                     prif_this_image_no_coarray, &
                     prif_form_team, prif_change_team, prif_end_team, &
                     prif_image_index_with_team, prif_image_index_with_team_number, &
+                    prif_initial_team_index, prif_initial_team_index_with_team, prif_initial_team_index_with_team_number, &
                     prif_this_image_with_coarray, prif_this_image_with_dim, &
                     prif_lcobound_no_dim, prif_ucobound_no_dim, &
-                    prif_num_images_with_team
+                    prif_num_images_with_team, PRIF_INITIAL_TEAM
     use veggies, only: result_t, test_item_t, assert_equals, assert_that, describe, it, succeed
 
     implicit none
@@ -19,7 +20,7 @@ function test_prif_image_index() result(tests)
         type(test_item_t) :: tests
 
         tests = describe( &
-          "prif_image_index", &
+          "prif_image_index and prif_initial_team_index", &
           [ it("returns 1 for the simplest case", check_simple_case) &
           , it("returns 1 when given the lower bounds", check_lower_bounds) &
           , it("returns 0 with invalid subscripts", check_invalid_subscripts) &
@@ -36,13 +37,17 @@ function check_this_image_coarray(coarray_handle, corank, team) result(result_)
         type(result_t) :: result_
 
         integer(c_int64_t) :: co, cosubscripts(corank), colbound(corank), coubound(corank)
-        integer(c_int) :: i, me
+        integer(c_int) :: i, me, me_initial
+        type(prif_team_type) :: initial_team
+
+        call prif_get_team(PRIF_INITIAL_TEAM, team=initial_team)
 
         result_ = succeed("")
 
         call prif_lcobound_no_dim(coarray_handle, colbound)
         call prif_ucobound_no_dim(coarray_handle, coubound)
         call prif_this_image_no_coarray(team, me)
+        call prif_this_image_no_coarray(initial_team, me_initial)
 
         call prif_this_image_with_coarray(coarray_handle, team=team, cosubscripts=cosubscripts)
         do i=1,corank
@@ -60,6 +65,15 @@ function check_this_image_coarray(coarray_handle, corank, team) result(result_)
           call prif_image_index(coarray_handle, cosubscripts, i)
         end if
         result_ = result_ .and. assert_equals(i, me)
+
+        ! and prif_initial_team_index
+        if (present(team)) then
+          call prif_initial_team_index_with_team(coarray_handle, cosubscripts, team, i)
+        else
+          call prif_initial_team_index(coarray_handle, cosubscripts, i)
+        end if
+        result_ = result_ .and. assert_equals(i, me_initial)
+
     end function
 
     function check_simple_case() result(result_)
@@ -80,6 +94,9 @@ function check_simple_case() result(result_)
         call prif_image_index(coarray_handle, [1_c_int64_t], image_index=answer)
         result_ = assert_equals(1_c_int, answer)
 
+        call prif_initial_team_index(coarray_handle, [1_c_int64_t], initial_team_index=answer)
+        result_ =  result_ .and. assert_equals(1_c_int, answer)
+
         result_ = result_ .and. &
           check_this_image_coarray(coarray_handle, 1)
 
@@ -104,6 +121,9 @@ function check_lower_bounds() result(result_)
         call prif_image_index(coarray_handle, [2_c_int64_t, 3_c_int64_t], image_index=answer)
         result_ = assert_equals(1_c_int, answer)
 
+        call prif_initial_team_index(coarray_handle, [2_c_int64_t, 3_c_int64_t], initial_team_index=answer)
+        result_ =  result_ .and. assert_equals(1_c_int, answer)
+
         result_ = result_ .and. &
           check_this_image_coarray(coarray_handle, 2)
 
@@ -139,7 +159,7 @@ function check_complicated_2d() result(result_)
 
         type(prif_coarray_handle) :: coarray_handle
         type(c_ptr) :: allocated_memory
-        integer(c_int) :: answer, ni
+        integer(c_int) :: answer, ni, expected
         call prif_num_images(num_images=ni)
 
         call prif_allocate_coarray( &
@@ -150,7 +170,13 @@ function check_complicated_2d() result(result_)
                 coarray_handle = coarray_handle, &
                 allocated_memory = allocated_memory)
         call prif_image_index(coarray_handle, [1_c_int64_t, 3_c_int64_t], image_index=answer)
-        result_ = assert_equals(merge(3_c_int,0_c_int,ni >= 3), answer)
+        expected = merge(3_c_int,0_c_int,ni >= 3)
+        result_ = assert_equals(expected, answer)
+
+        if (expected > 0) then
+          call prif_initial_team_index(coarray_handle, [1_c_int64_t, 3_c_int64_t], initial_team_index=answer)
+          result_ =  result_ .and. assert_equals(expected, answer)
+        end if
 
         result_ = result_ .and. &
           check_this_image_coarray(coarray_handle, 2)
@@ -163,7 +189,7 @@ function check_complicated_3d() result(result_)
 
         type(prif_coarray_handle) :: coarray_handle
         type(c_ptr) :: allocated_memory
-        integer(c_int) :: answer, ni
+        integer(c_int) :: answer, ni, expected
         type(prif_team_type) :: initial_team
         call prif_get_team(team=initial_team)
         call prif_num_images_with_team(team=initial_team, num_images=ni)
@@ -178,7 +204,15 @@ function check_complicated_3d() result(result_)
         call prif_image_index_with_team(coarray_handle, &
                            [2_c_int64_t, 1_c_int64_t, 1_c_int64_t], &
                            team=initial_team, image_index=answer)
-        result_ = assert_equals(merge(8_c_int,0_c_int,ni >= 8), answer)
+        expected = merge(8_c_int,0_c_int,ni >= 8)
+        result_ = assert_equals(expected, answer)
+
+        if (expected > 0) then
+          call prif_initial_team_index_with_team(coarray_handle, &
+                           [2_c_int64_t, 1_c_int64_t, 1_c_int64_t], &
+                           team=initial_team, initial_team_index=answer)
+          result_ =  result_ .and. assert_equals(expected, answer)
+        endif
 
         result_ = result_ .and. &
           check_this_image_coarray(coarray_handle, 3)
@@ -209,12 +243,14 @@ function check_complicated_2d_team() result(result_)
                 coarray_handle = coarray_handle, &
                 allocated_memory = allocated_memory)
 
-        which_team = merge(1_c_int64_t, 2_c_int64_t, mod(me, 2) == 0)
+        which_team = merge(2_c_int64_t, 1_c_int64_t, mod(me, 2) == 0)
         call prif_form_team(team_number = which_team, team = child_team)
         call prif_change_team(child_team)
 
           call prif_num_images_with_team(team=child_team, num_images=cni)
 
+          ! image_index lcobound
+
           call prif_image_index_with_team(coarray_handle, &
                               [0_c_int64_t, 2_c_int64_t], &
                               team=initial_team, image_index=answer)
@@ -245,6 +281,40 @@ function check_complicated_2d_team() result(result_)
           result_ = result_ .and. &
             assert_equals(1_c_int, answer)
 
+          ! initial_team_index lcobound
+
+          call prif_initial_team_index_with_team(coarray_handle, &
+                              [0_c_int64_t, 2_c_int64_t], &
+                              initial_team, answer)
+          result_ = result_ .and. &
+            assert_equals(1_c_int, answer)
+
+          call prif_initial_team_index_with_team_number(coarray_handle, &
+                              [0_c_int64_t, 2_c_int64_t], &
+                              -1_c_int64_t, answer)
+          result_ = result_ .and. &
+            assert_equals(1_c_int, answer)
+
+          call prif_initial_team_index_with_team(coarray_handle, &
+                              [0_c_int64_t, 2_c_int64_t], &
+                              child_team, answer)
+          result_ = result_ .and. &
+            assert_equals(merge(1_c_int,2_c_int,which_team==1), answer)
+
+          call prif_initial_team_index_with_team_number(coarray_handle, &
+                              [0_c_int64_t, 2_c_int64_t], &
+                              which_team, answer)
+          result_ = result_ .and. &
+            assert_equals(merge(1_c_int,2_c_int,which_team==1), answer)
+
+          call prif_initial_team_index(coarray_handle, &
+                              [0_c_int64_t, 2_c_int64_t], &
+                              answer)
+          result_ = result_ .and. &
+            assert_equals(merge(1_c_int,2_c_int,which_team==1), answer)
+
+          ! image_index 3
+
           call prif_image_index_with_team(coarray_handle, &
                               [0_c_int64_t, 3_c_int64_t], &
                               team=initial_team, image_index=answer)
@@ -275,6 +345,40 @@ function check_complicated_2d_team() result(result_)
           result_ = result_ .and. &
             assert_equals(merge(3_c_int,0_c_int,cni >= 3), answer)
 
+          ! initial_team_index 3
+          if (ni >= 3) then
+            call prif_initial_team_index_with_team(coarray_handle, &
+                                [0_c_int64_t, 3_c_int64_t], &
+                                team=initial_team, initial_team_index=answer)
+            result_ = result_ .and. &
+              assert_equals(3_c_int, answer)
+  
+            call prif_initial_team_index_with_team_number(coarray_handle, &
+                                [0_c_int64_t, 3_c_int64_t], &
+                                team_number=-1_c_int64_t, initial_team_index=answer)
+            result_ = result_ .and. &
+              assert_equals(3_c_int, answer)
+          end if
+          if (cni >= 3) then
+            call prif_initial_team_index_with_team(coarray_handle, &
+                                [0_c_int64_t, 3_c_int64_t], &
+                                team=child_team, initial_team_index=answer)
+            result_ = result_ .and. &
+              assert_equals(merge(5_c_int,6_c_int,which_team==1), answer)
+  
+            call prif_initial_team_index_with_team_number(coarray_handle, &
+                                [0_c_int64_t, 3_c_int64_t], &
+                                team_number=which_team, initial_team_index=answer)
+            result_ = result_ .and. &
+              assert_equals(merge(5_c_int,6_c_int,which_team==1), answer)
+  
+            call prif_initial_team_index(coarray_handle, &
+                                [0_c_int64_t, 3_c_int64_t], &
+                                initial_team_index=answer)
+            result_ = result_ .and. &
+              assert_equals(merge(5_c_int,6_c_int,which_team==1), answer)
+          end if
+
           result_ = result_ .and. &
             check_this_image_coarray(coarray_handle, 2, initial_team)
           result_ = result_ .and. &