toa_benchmark.py

import warnings
warnings.filterwarnings('ignore')
import sys
import timeit
import numba
import numpy as np
import skimage.io
import pyopencl as cl

# testar:
#   [X] Usar OpenCL c/ Buffer
#   [X] Usar OpenCL c/ Buffer em CPU Intel
#   [ ] Usar Image2d na implementação com OpenCL pura
#   [X] Usar arrayfire
#   [ ] Usar PyTorch
#   [ ] Usar Eigen (com Python)


def read_image(filepath):
    return skimage.io.imread(filepath)


@numba.jit(nopython=True, parallel=False, nogil=True, error_model='numpy', fastmath=True)
def _numba_toa_overwrite(data, m, a, sun_elevation, out):
    for y in range(data.shape[0]):
        for x in range(data.shape[1]):
            out[y, x] = (data[y, x] * m + a) / np.sin(sun_elevation)
        

@numba.jit(nopython=True, parallel=True, nogil=True, error_model='numpy', fastmath=True)
def _numba_toa(data, m, a, sun_elevation):
    return (data * m + a) / np.sin(sun_elevation)
    

class numba_impl:
    name = 'numba'

    def __init__(self, data):
        self.data = data
        self.out = np.empty_like(data)
        # _numba_toa_overwrite(self.data, np.float32(1), np.float32(1), np.float32(1), self.out)
        _numba_toa(self.data, np.float32(1), np.float32(1), np.float32(1))
        

    def __call__(self, m, a, sun_elevation):
        # _numba_toa_overwrite(self.data, m, a, sun_elevation, self.out)
        # return self.out
        return _numba_toa(self.data, m, a, sun_elevation)
        


class numpy_impl:
    name = 'numpy'

    def __init__(self, data):
        self.data = data

    def __call__(self, m, a, sun_elevation):
        return (self.data * m + a) / np.sin(sun_elevation)


class arrayfire_impl:
    name = 'arrayfire'

    def __init__(self, data, backend='cuda'):
        import arrayfire as af

        af.set_backend(backend)
        self.data = data
        self.arr = af.to_array(self.data)
        self.output = np.empty_like(self.data)

    def __call__(self, m, a, sun_elevation):
        return (self.arr * m + a) / np.sin(sun_elevation)     # TODO: remove np.sin!


class opencl_impl:
    name = 'pyopencl'

    def __init__(self, data):
        self.data = data
        type_size = np.dtype(data.dtype).itemsize
        total_work = data.size
        self.local_work_size = (np.int32(8), )
        self.global_work_size = (np.int32(total_work + (self.local_work_size[0] - total_work % self.local_work_size[0])), )

        self.output = np.empty_like(self.data)

        platforms = cl.get_platforms()
        self.ctx = cl.Context(dev_type=cl.device_type.ALL,
                              properties=[(cl.context_properties.PLATFORM, platforms[0])])

        self.program = cl.Program(self.ctx, """
        __kernel void toa(const float m, const float a, const float sun_elevation,
                          __global const float* input, __global float* output, const int N) {
            int i = get_global_id(0);
            if (i < N) {
                output[i] = (input[i] * m + a) / sin(sun_elevation);
            }
        }""").build()

    def __call__(self, m, a, sun_elevation):
        with cl.CommandQueue(self.ctx) as queue:
            input_buffer = cl.Buffer(self.ctx, cl.mem_flags.READ_ONLY | cl.mem_flags.USE_HOST_PTR, hostbuf=self.data)

            output_buffer = cl.Buffer(self.ctx, cl.mem_flags.WRITE_ONLY | cl.mem_flags.USE_HOST_PTR, hostbuf=self.output)

            self.program.toa(queue, self.global_work_size, self.local_work_size,
                m, a, sun_elevation, input_buffer, output_buffer, np.int32(self.data.size))

            return self.output



def run_impl(impl, data, m, a, s, ref, N=20, **kwargs):
    fn = impl(data, **kwargs)

    start = timeit.default_timer()
    for _ in range(N):
        result = fn(m, a, s)
    end = timeit.default_timer()
    print(f'{impl.name} took {end - start}s (close enough? {np.allclose(result, ref)})')

    return result


if __name__ == '__main__':
    input_filepath = '/mnt/media/data/satms/LC82200752017038LGN00_B4.TIF'
    imdata = read_image(input_filepath).astype(np.float32)
    # imdata = np.random.rand(7000, 7000).astype(np.float32)

    mul_band = np.float32(1.0269e-02)
    add_band = np.float32(-51.34694)
    sun_elevation = np.float32(57.44187420)

    toa_ref = numpy_impl(imdata)(mul_band, add_band, sun_elevation)

    impls = [
        numpy_impl,
        numba_impl,
        # arrayfire_impl,
        opencl_impl,
    ]

    for impl_class in impls:
        run_impl(impl_class, imdata, mul_band, add_band, sun_elevation, ref=toa_ref)