numba_dlcm.py

# Use with https://github.com/gmarkall/numba/tree/cuda-linker-options

from numba import cuda, float32, void


def axpy(r, a, x, y):
    start = cuda.grid(1)
    step = cuda.gridsize(1)

    for i in range(start, len(r), step):
        r[i] = a * x[i] + y[i]


sig = void(float32[::1], float32, float32[::1], float32[::1])
with_dlcm_cg = cuda.jit(sig, dlcm='cg')(axpy)
with_dlcm_ca = cuda.jit(sig, dlcm='ca')(axpy)
with_dlcm_xx = cuda.jit(sig)(axpy)

with open('dlcm_cg.sass', 'w') as f:
    f.write(with_dlcm_cg.inspect_sass()[sig.args])

with open('dlcm_ca.sass', 'w') as f:
    f.write(with_dlcm_ca.inspect_sass()[sig.args])

with open('dlcm_xx.sass', 'w') as f:
    f.write(with_dlcm_xx.inspect_sass()[sig.args])