#!/usr/bin/env perl
################################################################################
# exportmap - Generate MVAPICH2 mapping for MPI/Thread Combo
# Mahidhar Tatineni, San Diego Supercomputer Center(SLURM UPDATE) MAY  2015
################################################################################

use strict;
use warnings;
use Sys::Hostname;
use Getopt::Long;
use POSIX;
use Pod::Usage;
use constant CONTACT => 'Contact help@xsede.org for additional assistance.';

our $options; 

$options->{sys_ppn} = 24;
$options->{tpr} = 1;
$options->{rm_ppn} = $ENV{SLURM_NTASKS_PER_NODE};
if ( $options->{rm_ppn} )
{
 $options->{sys_ppn} = $options->{rm_ppn};
}

### Let command-line options override the defaults defined above
$options = get_options( \@ARGV );

my $npernode= $options->{sys_ppn} / $options->{tpr}; 
my $loc1 = trim(`cat /proc/self/cgroup | grep cpuset`);
my @locs = split( m/:/, $loc1 );
my $loc2 = trim($locs[2]);
my $loc3 = trim(`cat /cgroup/cpuset/$loc2/cpuset.cpus`);
my @corelist = split( m/-/, $loc3 );
my $core = $corelist[0];
my $binding_strs = "";
for my $i ( 1 .. $npernode ) {
  $binding_strs .= ":" unless $i == 1;
  $binding_strs .= sprintf( "%d-", $core );
  $core += $options->{tpr};
  $binding_strs .= sprintf( "%d", $core-1 );
  }
print STDOUT "$binding_strs";


#### TRIM ####
sub  trim { my $s = shift; $s =~ s/^\s+|\s+$//g; return $s };

################################################################################
### get_options: Parse command line input and check for errors.  This will be 
### called AFTER the resource manager's paramaters have already been loaded into
### the options hash.
################################################################################
sub get_options
{
    my $argv = shift;
    my $old_options = $options;
    my %new_options;
    use Getopt::Long;
    Getopt::Long::Configure( "require_order", "no_ignore_case", 
                            "auto_version", "auto_help", "no_auto_abbrev" );
    if ( @$argv > 0 ) 
    {
        GetOptions(
            "tpr|tpp|threads-per-rank|threads-per-process=i"  => \$new_options{tpr}, 
        );
    }

    my %options = ( %$old_options, %new_options );
    return \%options;
}

################################################################################
### bailout: terminate with an error and optionally delete the temporary 
### nodefile created by ibrun
################################################################################
sub bailout
{
    my $errno = shift;
    my $errmsg = shift;
    print STDERR "GENMAP ERROR: $errmsg\n" . CONTACT . "\n";

    ### Delete temporary nodefile on error
    if ( $options
    &&   $options->{rm_nodefile}
    &&   $options->{nodefile}
    &&   $options->{nodefile} ne $options->{rm_nodefile}
    &&  -e $options->{nodefile} )
    {
        unlink $options->{nodefile};
    }

    exit $errno;
}
sub genmap_warn
{
    print STDERR "GENMAP WARNING: ";
    printf( @_ );
}

################################################################################
### dprintf: print messages if debugging is enabled
################################################################################
sub dprintf
{
    return unless $options->{verbose};
    print "GENMAP: ";
    printf( @_ );
}