Skip to content

Draft: cupy-compatible DielectricFunctionCalculator

Ask Hjorth Larsen requested to merge cupy-compatible-dfc into master

This is a working implementation of DielectricFunctionCalculator with cupy.

It can be used like this, where lib is either numpy (the module) or cupy (the module):

def test_dfc(lib):
    import sys
    nG = int(sys.argv[1])
    rng = np.random.RandomState(42)

    sqrV_G = lib.asarray(rng.rand(nG))
    chi0_GG = lib.asarray(rng.rand(nG, nG))

    t1 = time()
    dfc = DielectricFunctionCalculator(sqrV_G, chi0_GG, 'GW', lib=lib)
    einv_GG = dfc.get_einv_GG()
    e_GG = dfc.get_e_GG()
    t2 = time()

    T = t2 - t1
    print('time', t2 - t1)

    err = abs(einv_GG @ e_GG - dfc.I_GG).max()
    print('err', err)
    return einv_GG, T

Thus, any class or function with pure (and simple) numpy-based arithmetic can be transparently replaced by cupy.

This means we can isolate "bubbles" of processing, extract those bubbles to functions or small classes, and then ensure (via parametrized testing) that they individually function correctly with both numpy and cupy.

The code connecting bubbles (where there might be non-trivial work depending on whether we're doing GPU or not) will still become somewhat messy, but in my opinion this initial attempt shows that things are not too bad. Also, a decorator for transparent replacement of certain BLAS calls might come in handy.

Edited by Ask Hjorth Larsen

Merge request reports