Numba не может распараллелить цикл - PullRequest
Новая
       16

Numba не может распараллелить цикл

1 голос
/ 06 апреля 2020

Следующий код выполняется, но l oop с prange не удается распараллелить, хотя он явно паралеллизуем: 


import numpy as np
from numba import njit, prange


@njit(parallel=True)
def f1(money, u, v, cost_u, cost_v):

    # task: find index of u iu_opt and index of v iv_opt 
    # such that u[iu_opt] + v[iv_opt] is maximal subject
    # to cost_u[iu_opt] + cost_v[iv_opt] < money


    na = money.size

    ncu = cost_u.size
    ncv = cost_v.size

    iu_opt = np.empty((na,),dtype=np.int16)
    iv_opt = np.empty((na,),dtype=np.int16)

    for ia in prange(na):


        money_i = money[ia]
        ivbest = 0 # initially pick iv = 0

        # find max iu corresponding to iv = 0
        for iu in range(ncu-1,-1,-1):
            if cost_u[iu] + cost_v[0] < money_i: break

        iubest = iu
        # compute initial score
        score_best = u[iu] + v[0]  


        # then try to increase iv
        for iv in range(1,ncv):

            # it not enough money for u_0
            if cost_v[iv]  + cost_u[0] > money_i: break 

            while cost_v[iv] + cost_u[iu] > money_i:
                iu -= 1

            assert iu >= 0            

            score_now = u[iu] + v[iv] 


            if score_now > score_best:
                ivbest = iv
                iubest = iu
                score_best = score_now

        iu_opt[ia] = iubest
        iv_opt[ia] = ivbest

    return iu_opt, iv_opt



na = 50
ncu = 204
ncv = 205


money = np.arange(na)/(na)


cost_u = np.arange(ncu)/ncu
u = np.cumsum(np.random.random_sample(ncu))
cost_v = np.arange(ncv)/ncv
v = np.cumsum(np.random.random_sample(ncv))

iu, iv = f1(money, u, v, cost_u, cost_v)
f1.parallel_diagnostics(level=4)

Если это помогает, установка проблемы следующая: значение u [i] стоит cost_u [i], а значение v [j] стоит cost_v [j] (все эти последовательности строго растут), для каждого значения money [ia] мы хотим найти i и j, максимизирующие u [i] + v [j], учитывая, что cost_u [i] + cost_v [j] не может превышать деньги [ia].

1 Ответ

0 голосов
/ 08 апреля 2020

На случай, если у кого-то возникнет подобная проблема, я наконец заставил ее работать, разделив внутренности большого l oop на другую функцию njit. Вот код: 

import numpy as np
from numba import njit, prange


@njit(parallel=True)
def f1(money, u, v, cost_u, cost_v):

    # task: find index of u iu_opt and index of v iv_opt 
    # such that u[iu_opt] + v[iv_opt] is maximal subject
    # to cost_u[iu_opt] + cost_v[iv_opt] < money


    na = money.size

    iu_opt = np.empty((na,),dtype=np.int16)
    iv_opt = np.empty((na,),dtype=np.int16)

    for ia in prange(na):

        money_i = money[ia]

        iubest, ivbest = f1_int(money_i,u,v,cost_u,cost_v)

        iu_opt[ia] = iubest
        iv_opt[ia] = ivbest

    return iu_opt, iv_opt


@njit
def f1_int(money_i,u,v,cost_u,cost_v):
    ivbest = 0 # initially pick iv = 0

    ncu = cost_u.size
    ncv = cost_v.size

    # find max iu corresponding to iv = 0
    for iu in range(ncu-1,-1,-1):
        if cost_u[iu] + cost_v[0] < money_i: break

    iubest = iu
    # compute initial score
    score_best = u[iu] + v[0]  


    # then try to increase iv
    for iv in range(1,ncv):

        # it not enough money for u_0
        if cost_v[iv]  + cost_u[0] > money_i: break 

        while cost_v[iv] + cost_u[iu] > money_i:
            iu -= 1

        assert iu >= 0            

        score_now = u[iu] + v[iv] 


        if score_now > score_best:
            ivbest = iv
            iubest = iu
            score_best = score_now
    return iubest, ivbest

na = 50
ncu = 204
ncv = 205


money = np.arange(na)/(na)


cost_u = np.arange(ncu)/ncu
u = np.cumsum(np.random.random_sample(ncu))
cost_v = np.arange(ncv)/ncv
v = np.cumsum(np.random.random_sample(ncv))

iu, iv = f1(money, u, v, cost_u, cost_v)
f1.parallel_diagnostics(level=4)

Это на самом деле не отвечает на вопрос, почему проблема возникает, но работает как-то.

...