#!/usr/local/bin/perl
#
#   mpirun  for PowerMPI
#
#   Written by Ilya Evseev, Apr 1999.  Contact e-mail: evseev@csa.ru
#
#   Note for most advanced users of own supercomputing resources:
#
#   On the PowerMouse (telnet pink.csa.ru) you need to select count of nodes
#   ("-np N" switch) strictly equal to actual size of any existing physical
#   partition because NRM configuration mistake. Valid N values are: 1,2,4 or 8
#   For example, "-np 3" or "-np 5" will be treated as invalid. Sorry.
#

sub Syntax
{
    if( $_[0] ne "" )
      { print "mpirun: Error: @_\n"; }
    print "Syntax: mpirun -np N [-show] [--help] [px_switches] program.px args...\n";
    print "..Note: All px_switches are passed to \"px run\" without changes\n";
    print "..Note: for all other details simply type: more \`which $0\`\n";
    exit 1;
}

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
#   NOTE: name of partition is used   #
#   as key for all hash arrays        #
#_____________________________________#
%sizes = {};    # count of nodes in partitions
%owners = {};   # name of current owner, "*" = free


#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
#   Find free partition of given size     #
#   Return name, if found; "" otherwise   #
#_________________________________________#
sub tryPart
{
  # print "DEBUG: tryPart( $_[0] )\n";
    while ( ($key,$value) = each %owners )
    {
        if(( $value eq "*" ) && ( $sizes{$key} == $_[0] ))
          { return $key; }
    }
    "";
}


#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
#   params: N_nodes, i, j, alpha               #
#   return: (N_nodes+j-1) div j + (i-j+alpha)  #
#______________________________________________#
sub calcCPUsCount
{
    my( $temp ) = ( $_[0] + $_[2] - 1 );
    (( $temp - ( $temp % $_[2] )) / $_[2] ) + ( $_[1] - $_[2] + $_[3] );
}


#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
#   Parse command line, extract "-np N"   #
#_________________________________________#
@Args = ();
$NP = 0;
$MaxCPUs = 0;
$Show = "";
for( $argc=0 ; $argc <= $#ARGV ; $argc++ )
{
    $_ = $ARGV[$argc];
    if( $_ eq "-np" )
      { $NP = $ARGV[++$argc]; }
    elsif( $_ eq "-show" )
      { $Show="/bin/echo"; }
    elsif( $_ eq "--help" )
      { Syntax; }
    else
      { $Args[$argc] = $_; }
}
Syntax if( $argc == 0 );
Syntax "invalid or missing \"-np N\" value" if( $NP == 0 );


#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
#   Read static description of partitions: names and sizes   #
#   Each line must be in form  "name   |  pNN  | ..."        #
#____________________________________________________________#
open(NRM_PC,"px nrm -pc|") || die "mpirun: Error: cannot execute 'px nrm -pc'";
while( <NRM_PC> )
{
    if( /^([^ |]*) *\| *p+([0-9]+) *\|/ )
    {
        $sizes{$1} = $2;
        if( $MaxCPUs < $2 )
          { $MaxCPUs = $2; }
    }
}
close( NRM_PC );


#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
#   Read allocation state of partitions                   #
#   Each line must be in form  "name   |  owner  | ..."   #
#_________________________________________________________#
open(NRM_PA,"px nrm -pa|") || die "mpirun: Error: cannot execute 'px nrm -pa'";
while( <NRM_PA> )
{
    if( /^([^ |]*) *\| *([^ |]+) *\|/ )
      { $owners{$1} = $2; }
}
close( NRM_PA );


#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
#   Find optimal partition for running   #
#________________________________________#
if(( $part = tryPart( $NP ) ) eq "" )
{
    for( $i = 1; $part eq ""; $i++ )
    {
        if(( $i > $NP ) && ( $NP + $i > $MaxCPUs ))
          { die "mpirun: Error: cannot find free partition for running\n"; }
     #  $part = tryPart( $NP + $i ) unless($part ne "");
        $part = tryPart( calcCPUsCount( $NP, $i, $i, 0 ) ) unless($part ne "");
        for( $j = $i; ( $j > 0 ) && ( $part eq "" ); --$j )
        {
            $part = tryPart( calcCPUsCount( $NP, $i, $j, 1 ) );
        }
    }
    $NP = "$NP 1";
} else
  { $NP = ""; }

## Simple do it!..
exit exec "$Show px run -a $part $NP @Args";

## EOF ##