Commit 60b3d1b6 authored by Jakub Klinkovský's avatar Jakub Klinkovský Committed by Jakub Klinkovský
Browse files

Improved algorithm in Partitioner.h to fix problems with integer overflow and... 

Improved algorithm in Partitioner.h to fix problems with integer overflow and produce more optimal partitioning
parent 2ac3133b

src/TNL/Containers/Partitioner.h

+25 −25
Original line number Original line Diff line number Diff line
@@ -29,28 +29,22 @@ public: 
   static SubrangeType

   static SubrangeType

   splitRange( Index globalSize, const MPI::Comm& communicator )

   splitRange( Index globalSize, const MPI::Comm& communicator )

   {

   {

      if( communicator != MPI_COMM_NULL ) {

      if( communicator == MPI_COMM_NULL )

         return { 0, 0 };



      const int rank = communicator.rank();

      const int rank = communicator.rank();

      const int partitions = communicator.size();

      const int partitions = communicator.size();

         const Index begin = TNL::min( globalSize, rank * globalSize / partitions );

         const Index end = TNL::min( globalSize, ( rank + 1 ) * globalSize / partitions );

         return SubrangeType( begin, end );

      }

      else

         return SubrangeType( 0, 0 );

   }





   // Gets the owner of given global index.

      const Index partSize = globalSize / partitions;

   __cuda_callable__

      const int remainder = globalSize % partitions;

   static int

      if( rank < remainder ) {

   getOwner( Index i, Index globalSize, int partitions )

         const Index begin = rank * ( partSize + 1 );

   {

         const Index end = begin + partSize + 1;

      int owner = i * partitions / globalSize;

         return { begin, end };

      if( owner < partitions - 1 && i >= getOffset( globalSize, owner + 1, partitions ) )

      }

         owner++;

      const Index begin = remainder * ( partSize + 1 ) + ( rank - remainder ) * partSize;

      TNL_ASSERT_GE( i, getOffset( globalSize, owner, partitions ), "BUG in getOwner" );

      const Index end = begin + partSize;

      TNL_ASSERT_LT( i, getOffset( globalSize, owner + 1, partitions ), "BUG in getOwner" );

      return { begin, end };

      return owner;

   }

   }





   // Gets the offset of data for given rank.

   // Gets the offset of data for given rank.
@@ -58,7 +52,11 @@ public: 
   static Index

   static Index

   getOffset( Index globalSize, int rank, int partitions )

   getOffset( Index globalSize, int rank, int partitions )

   {

   {

      return rank * globalSize / partitions;

      const Index partSize = globalSize / partitions;

      const int remainder = globalSize % partitions;

      if( rank < remainder )

         return rank * ( partSize + 1 );

      return remainder * ( partSize + 1 ) + ( rank - remainder ) * partSize;

   }

   }





   // Gets the size of data assigned to given rank.

   // Gets the size of data assigned to given rank.
@@ -66,9 +64,11 @@ public: 
   static Index

   static Index

   getSizeForRank( Index globalSize, int rank, int partitions )

   getSizeForRank( Index globalSize, int rank, int partitions )

   {

   {

      const Index begin = min( globalSize, rank * globalSize / partitions );

      const Index partSize = globalSize / partitions;

      const Index end = min( globalSize, ( rank + 1 ) * globalSize / partitions );

      const int remainder = globalSize % partitions;

      return end - begin;

      if( rank < remainder )

         return partSize + 1;

      return partSize;

   }

   }





   template< typename Device >

   template< typename Device >