Python Forum

Hello,

I am a complete beginner in GPU programming.

I have the following piece of code that I have been running in parallel on CPUs.

import cupy
import ot.gpu

with Pool(processes = workers) as p_kappa: 
            W = list(p_kappa.imap(partial(W_comp, mx_all, dist, lamb), G.edges()),\
                          total = len(G.edges))

def W_comp(mx_all, dist, lamb,  e):
    i = e[0]
    j = e[1]

    W = np.zeros(5)
    for it in range(5):

        Nx = np.array(mx_all[i][1][it]).flatten()
        mx = mx_all[i][0][it].toarray().flatten()
        my = mx_all[j][0][it].toarray().flatten()

        dNx = dist[Nx,:].copy(order='C')

        W[it] = ot.sinkhorn(mx, my, dNx, lamb)

Basically, a linear program ot.sinkhorn is in a nested loop.

I have discovered that there is a GPU version of the code ot.gpu.sinkhorn

Hence it would be ideal to execute the whole nested loop on GPU to benefit.

Could someone help how to do this?
Thank you!

This might be helpful: https://weeraman.com/put-that-gpu-to-goo...a437168c01