code4.py 1.52 KB
Newer Older
1 2 3
#!/usr/bin/env python
import numpy as np
import argparse
RODIONOV Sergey's avatar
RODIONOV Sergey committed
4
from multiprocessing import Process, Queue, cpu_count
5 6 7 8 9
import mkl

# To be sure that we use 1 thread in MKL
mkl.set_num_threads(1)

RODIONOV Sergey's avatar
RODIONOV Sergey committed
10 11 12
# function wait order from orderq
# stop if order < 0, otherwise make computation and put result in the rezq
def fun(orderq, rezq):
13 14 15 16 17 18 19
    while True:
        order = orderq.get();
        if (order < 0):
            return
        np.random.seed(order)
        A = np.random.rand (100, 100)
        B = np.random.rand (100, 100)
RODIONOV Sergey's avatar
RODIONOV Sergey committed
20 21
        rezq.put(np.sum(np.dot(A, B)))

22 23 24 25 26 27 28 29 30 31 32 33

# Parse command line arguments                                                                                                
parser = argparse.ArgumentParser()
parser.add_argument("Nm", help="Number of matrices", type=int);
parser.add_argument("-ncores", help="Number of cores to use", type=int);
args = parser.parse_args()

Nm      = args.Nm
ncores  = args.ncores

if (ncores == None):
    ncores = cpu_count()
RODIONOV Sergey's avatar
RODIONOV Sergey committed
34
        
35 36 37
#Adjust number of cores if necessary
if (ncores > Nm):
    ncores = Nm
RODIONOV Sergey's avatar
RODIONOV Sergey committed
38 39 40 41

# Make order and result queues
orderq = Queue()
rezq   = Queue()    
42 43 44 45

# List of all workers
procs  = []

RODIONOV Sergey's avatar
RODIONOV Sergey committed
46 47 48 49

#Initialize all workers
for i in range(ncores):
    p      = Process(target=fun, args=(orderq, rezq))    
50 51 52
    p.start()
    procs.append(p)

RODIONOV Sergey's avatar
RODIONOV Sergey committed
53 54 55
# Put all jobs in orderq
for i in range(Nm):
    orderq.put(i);
56

RODIONOV Sergey's avatar
RODIONOV Sergey committed
57 58 59 60
rez = 0
# Recive all results
for i in range(Nm):
    rez += rezq.get()
61 62 63

#stop all workers
for i in range(ncores):
RODIONOV Sergey's avatar
RODIONOV Sergey committed
64 65 66 67 68
    orderq.put(-1)

#Join all workers
for p in procs:
    p.join()
69

RODIONOV Sergey's avatar
RODIONOV Sergey committed
70
print (rez)