How to set different No. of MKL threads for each MPICH process

Question

I want to create different number of MKL threads for each MPICH process.

For example, lets say I have 4 MPICH processes. What I want is:

Process 1=4 MKL thread, Process 2=3 MKL thread, Process 3=5 MKL thread, etc.

I know export MKL_NUM_THREADS=4; export MKL_DOMAIN_NUM_THREADS="MKL_ALL=1, MKL_BLAS=4", but I don't get what I have to do in my specific case.

Answer 1

If what you want to achieve is that different ranks in an MPI job use different number of MKL threads, you can do it in two different ways.

You can set the number of threads in code according to the process rank, e.g.

#define NUM_PROCS 4

int threads_per_proc[NUM_PROCS] = { 4, 3, 5, 5 };
int rank;
MPI_Comm_rank(MPI_COMM_WORLD, &rank);
// ...
// Signal an error if rank >= 4
// ...
mkl_set_num_threads(threads_per_proc[rank]);

You can also use the MPMD launch mode of mpiexec (or mpirun), but that would be a bit cumbersome:

mpiexec -env MKL_NUM_THREADS 4 -n 1 ./program : \
        -env MKL_NUM_THREADS 3 -n 1 ./program : \
        -env MKL_NUM_THREADS 5 -n 2 ./program

I have omitted the options to set MKL_DOMAIN_NUM_THREADS for brevity. This launches one copy of the MPI program ./program with MKL_NUM_THREADS set to 4 (which copy becomes rank 0); one copy with MKL_NUM_THREADS set to 3 (which copy becomes rank 1); two copies with MKL_NUM_THREADS set to 5 (these copies become ranks 2 and 3).