Merge pull request #19 from sourceryinstitute/expand-real-kinds

feat(data_partition): support real{32,64} gathers

Author: Damian Rouson

src/data-partition-implementation.F90

@@ -13,7 +13,7 @@
 
     associate( ni => num_images() )
 
-      call assert( ni<=cardinality, "sufficient data for distribution across images")
+      call assert( ni<=cardinality, "sufficient data for distribution across images", cardinality)
 
       allocate(first_datum(ni), last_datum(ni))
 
@@ -54,7 +54,39 @@ pure function overflow(im, excess) result(extra_datum)
     last_index = last_datum( image_number )
   end procedure
 
-  module procedure gather_real_1D_array
+  module procedure gather_real32_1D_array
+
+    if (present(dim)) call assert (dim==1, "dimensioned partitioned == 1")
+
+    associate( me => this_image() )
+      if (verbose) then
+        write(6,*) 'gather_real_1D_array(): executing on image', me
+        flush(6)
+      end if
+      associate( first=>first(me), last=>last(me) )
+        if (.not. present(result_image)) then
+          a(1:first-1)  = 0.
+          a(last+1:)  = 0.
+          call co_sum(a)
+        else
+          block
+            real(real32), allocatable, dimension(:) :: a_lower, a_upper
+            a_lower = a(1:first-1)
+            a_upper = a(last+1:)
+            a(1:first-1)  = 0.
+            a(last+1:)  = 0.
+            call co_sum(a, result_image=result_image)
+            if (result_image /= me) then
+              a(1:first-1) = a_lower
+              a(last+1:) = a_upper
+            end if
+          end block
+        end if
+      end associate
+    end associate
+  end procedure
+
+  module procedure gather_real64_1D_array
 
     if (present(dim)) call assert (dim==1, "dimensioned partitioned == 1")
 
@@ -86,7 +118,72 @@ pure function overflow(im, excess) result(extra_datum)
     end associate
   end procedure
 
-  module procedure gather_real_2D_array
+  module procedure gather_real32_2D_array
+
+    integer dim_
+    if (present(dim)) then
+      dim_ = dim
+    else
+      dim_ = 2
+    end if
+
+    associate( me => this_image() )
+      if (verbose) then
+        write(6,*) 'gather_real32_2D_array(): executing on image', me
+        flush(6)
+      end if
+      associate( first => first(me), last => last(me) )
+        if (.not. present(result_image)) then
+          select case(dim_)
+            case(1)
+              a(1:first-1, :) = 0.
+              a(last+1:, :) = 0.
+            case(2)
+              a(:, 1:first-1) = 0.
+              a(:, last+1:) = 0.
+            case default
+              error stop "gather_real32_2D_array: invalid dim argument"
+          end select
+          call co_sum(a)
+        else
+          block
+            real(real32), allocatable, dimension(:,:) :: a_lower, a_upper
+            select case(dim_)
+              case(1)
+                a_lower = a(1:first-1, :)
+                a_upper = a(last+1:, :)
+                a(1:first-1, :) = 0.
+                a(last+1:, :) = 0.
+              case(2)
+                a_lower = a(:, 1:first-1)
+                a_upper = a(:, last+1:)
+                a(:, 1:first-1) = 0.
+                a(:, last+1:) = 0.
+              case default
+                error stop "gather_real32_2D_array: invalid dim argument"
+            end select
+
+            call co_sum(a, result_image=result_image)
+
+            if (result_image /= me) then
+              select case(dim_)
+                case(1)
+                  a(1:first-1, :) = a_lower
+                  a(last+1:, :) = a_upper
+                case(2)
+                  a(:, 1:first-1) = a_lower
+                  a(:, last+1:) = a_upper
+                case default
+                  error stop "gather_real32_2D_array: invalid dim argument"
+              end select
+            end if
+          end block
+        end if
+      end associate
+    end associate
+  end procedure
+
+  module procedure gather_real64_2D_array
 
     integer dim_
     if (present(dim)) then
@@ -97,7 +194,7 @@ pure function overflow(im, excess) result(extra_datum)
 
     associate( me => this_image() )
       if (verbose) then
-        write(6,*) 'gather_real_2D_array(): executing on image', me
+        write(6,*) 'gather_real64_2D_array(): executing on image', me
         flush(6)
       end if
       associate( first => first(me), last => last(me) )
@@ -110,7 +207,7 @@ pure function overflow(im, excess) result(extra_datum)
               a(:, 1:first-1) = 0.
               a(:, last+1:) = 0.
             case default
-              error stop "gather_real_2D_array: invalid dim argument"
+              error stop "gather_real64_2D_array: invalid dim argument"
           end select
           call co_sum(a)
         else
@@ -128,7 +225,7 @@ pure function overflow(im, excess) result(extra_datum)
                 a(:, 1:first-1) = 0.
                 a(:, last+1:) = 0.
               case default
-                error stop "gather_real_2D_array: invalid dim argument"
+                error stop "gather_real64_2D_array: invalid dim argument"
             end select
 
             call co_sum(a, result_image=result_image)
@@ -142,7 +239,7 @@ pure function overflow(im, excess) result(extra_datum)
                   a(:, 1:first-1) = a_lower
                   a(:, last+1:) = a_upper
                 case default
-                  error stop "gather_real_2D_array: invalid dim argument"
+                  error stop "gather_real64_2D_array: invalid dim argument"
               end select
             end if
           end block

src/data-partition-interface.f90

@@ -1,7 +1,7 @@
 module data_partition_interface
   !! distribute data identification numbers across images such that the number of
   !! items differs by at most 1 between any two images.
-  use iso_fortran_env, only : real64
+  use iso_fortran_env, only : real32, real64
   implicit none
 
   private
@@ -14,8 +14,8 @@ module data_partition_interface
     procedure, nopass :: define_partitions
     procedure, nopass :: first
     procedure, nopass :: last
-    procedure, nopass, private :: gather_real_2D_array, gather_real_1D_array
-    generic :: gather => gather_real_2D_array, gather_real_1D_array
+    procedure, nopass, private :: gather_real32_2D_array, gather_real64_2D_array,  gather_real32_1D_array, gather_real64_1D_array
+    generic :: gather => gather_real32_2D_array, gather_real64_2D_array,  gather_real32_1D_array, gather_real64_1D_array
   end type
 
   integer, allocatable :: first_datum(:), last_datum(:)
@@ -45,14 +45,28 @@ pure module function last(image_number) result(last_index)
     !! 1. Near the beginning/end of execution to amortize costs across an entire run or
     !! 2. Temporarily while developing/debugging code.
 
-    module subroutine gather_real_1D_array( a, result_image, dim )
+    module subroutine gather_real32_1D_array( a, result_image, dim )
+      !! Gather the elements of an 1D array distributed along dimension dim onto result_image
+      real(real32), intent(inout) :: a(:)
+      integer, intent(in), optional :: result_image
+      integer, intent(in), optional :: dim
+    end subroutine
+
+    module subroutine gather_real64_1D_array( a, result_image, dim )
       !! Gather the elements of an 1D array distributed along dimension dim onto result_image
       real(real64), intent(inout) :: a(:)
       integer, intent(in), optional :: result_image
       integer, intent(in), optional :: dim
     end subroutine
 
-    module subroutine gather_real_2D_array( a, result_image, dim )
+    module subroutine gather_real32_2D_array( a, result_image, dim )
+      !! Gather the elements of an 2D array distributed along dimension dim onto result_image
+      real(real32), intent(inout) :: a(:,:)
+      integer, intent(in), optional :: result_image
+      integer, intent(in), optional :: dim
+    end subroutine
+
+    module subroutine gather_real64_2D_array( a, result_image, dim )
       !! Gather the elements of an 2D array distributed along dimension dim onto result_image
       real(real64), intent(inout) :: a(:,:)
       integer, intent(in), optional :: result_image