Az első példaprogramunk a vektor összeadás szekvenciális kódja. Ehhez egy függvényben adjuk össze az elempárokat egy ciklusban. A függvényt kívülről hívjuk meg.

import numpy as np
from timeit import default_timer as timer

def VectorAdd(a, b, c):
    for i in range(a.size):
        c[i] = a[i] + b[i]

def main():
    N = 32000000
    
    A = np.ones(N, dtype=np.float32)
    B = np.ones(N, dtype=np.float32)
    C = np.zeros(N, dtype=np.float32)
    
    start = timer()
    
    VectorAdd(A, B, C)
    
    vectoradd_time = timer() - start
    
    print("C[:5] = " + str(C[:5]))
    print("C[-5:] = " + str(C[-5:]))
    
    print("VectorAdd took {} seconds".format(vectoradd_time))

if __name__== "__main__":
    main()

A fenti kódot nem részletezzük túlságosan. Elég annyit látni, hogy a main függvényben előkészítünk két véletlen vektort, amiket a VectorAdd(...) függvényben összeadunk. Az elempárokon egy ciklus iterál végig.

Egy kicsit érdekesebb a program, ha az adat feldolgozást vektorizáljuk, és ezzel párhuzamosítjuk.

Ehhez egy sort kell beszúrnunk az összeadást végző függvény neve elé.

@vectorize(["float32(float32, float32)"], target='cuda')

A vektorizálás hatására numba elemekre bontja sazámítást, így már nincs szükségünk a cillusra sem. Egészen pontosan a teljes kód így néz ki.

import numpy as np
from timeit import default_timer as timer

from numba import vectorize

@vectorize(["float32(float32, float32)"], target='cuda')
def VectorAdd(a, b):
    return a + b

def main():
    N = 32000000
    
    A = np.ones(N, dtype=np.float32)
    B = np.ones(N, dtype=np.float32)
    C = np.zeros(N, dtype=np.float32)
    
    start = timer()
    
    C = VectorAdd(A, B)
    
    vectoradd_time = timer() - start
    
    print("C[:5] = " + str(C[:5]))
    print("C[-5:] = " + str(C[-5:]))
    
    print("VectorAdd took {} seconds".format(vectoradd_time))

if __name__== "__main__":
    main()

Ez eredémyn egyszerűbb, gyorsabb, és a " target='cuda' " megadása miatt automatikusan CUDA GPU-n fog futni.

Az eddigi programunk már önmagában is gyors. Vegyük észre viszont, hogy a GPU-n való számítás sokszor több lépésben történik. A mi adatain pedig a fenti esetekben a CPU-ról indulnak, és oda is ér vissza az eredmény, ami felesleges adat-másolásokat jelenthet.

A Numba-val van viszont lehetőségünk, hogy aa GPU memóriát kézzel kezeljük.

• Adatokat a grafikus memóriába a cuda.to_device(...) függvénnyel tudunk másolni;
• A grafikus memóriában adatot foglalni a cuda.device_array(...) függvénnyel lehet;
• Az eredményeket a központi memóriába a eredmény.copy_to_host() függvénnyel lehet vissza másolni;
• És ha már foglaltunk helyet az eredménynek, akkor azt a vektorizált függvényünk "out=" paraméterével specifikálni is tudjuk.

Valahogy így:

import numpy as np
from timeit import default_timer as timer

from numba import vectorize
from numba import cuda

@vectorize(["float32(float32, float32)"], target='cuda')
def VectorAdd(a, b):
    return a + b

def main():
    N = 32000000

    d_A = cuda.to_device(np.ones(N, dtype=np.float32))
    d_B = cuda.to_device(np.ones(N, dtype=np.float32))
    d_C = cuda.device_array(shape=(N,), dtype=np.float32)

    start = timer()

    VectorAdd(d_A, d_B, out=d_C)

    vectoradd_time = timer() - start

    C = d_C.copy_to_host();

    print("C[:5] = " + str(C[:5]))
    print("C[-5:] = " + str(C[-5:]))

    print("VectorAdd took {} seconds".format(vectoradd_time))


if __name__ == "__main__":
    main()