loopy_to_petsc_viennacl.py

from petsc4py import PETSc
import loopy as lp
import numpy as np
import pyopencl as cl
import coffee.system
from pyop2 import compilation
import ctypes
from mako.template import Template
import re

ctx = cl.create_some_context()


def get_grid_sizes(kernel):
    parameters = {}
    for arg in kernel.args:
        if isinstance(arg, lp.ValueArg) and arg.approximately is not None:
            parameters[arg.name] = arg.approximately

    glens, llens = kernel.get_grid_size_upper_bounds_as_exprs()

    from pymbolic import evaluate
    from pymbolic.mapper.evaluator import UnknownVariableError
    try:
        glens = evaluate(glens, parameters)
        llens = evaluate(llens, parameters)
    except UnknownVariableError as name:
        from warnings import warn
        warn("could not check axis bounds because no value "
                "for variable '%s' was passed to check_kernels()"
                % name)

    return llens, glens


def loopy_to_viennacl(kernel):
    lsize, gsize = get_grid_sizes(kernel)
    kernel = kernel.copy(
            target=lp.PyOpenCLTarget(ctx.devices[0]))
    # code instructions ahead.
    c_code_str = r'''#include <CL/cl.h>
            #include "petsc.h"
            #include "petscvec.h"
            #include "petscviennacl.h"

            // ViennaCL Headers
            #include "viennacl/ocl/backend.hpp"
            #include "viennacl/vector.hpp"
            #include "viennacl/backend/memory.hpp"

            char kernel_source[] = "${kernel_src}";

            extern "C" void ${kernel_name}_viennacl(${", ". join(
                                ['Vec ' + arg.name for arg in args])})
            {
                // viennacl vector declarations
                % for arg in args:
                viennacl::vector<PetscScalar> *${arg.name}_viennacl;
                % endfor

                // getting the array from the petsc vecs
                % for arg in args:
                VecViennaCLGetArrayReadWrite(${arg.name}, &${arg.name}_viennacl);
                % endfor

                // defining the kernel
                viennacl::ocl::context ctx =
                        ${args[0].name}_viennacl->handle().opencl_handle().context();
                viennacl::ocl::program & my_prog =
                            ctx.add_program(kernel_source, "kernel_program");
                viennacl::ocl::kernel &viennacl_kernel =
                        my_prog.get_kernel("${kernel_name}");

                // set work group sizes
                % for i, ls in enumerate(lsize):
                viennacl_kernel.local_work_size(${i}, ${ls});
                % endfor
                % for i, gs in enumerate(gsize):
                viennacl_kernel.global_work_size(${i}, ${gs});
                % endfor

                // enqueueing the kernel
                viennacl::ocl::enqueue(viennacl_kernel(${", ". join(
                                ['*' + arg.name + '_viennacl' for arg in
                                args])}));


                // restoring the arrays to the petsc vecs
                % for arg in args:
                VecViennaCLRestoreArrayReadWrite(${arg.name}, &${arg.name}_viennacl);
                % endfor
            }
            '''
    c_code_str = re.sub("\\n            ", "\n", c_code_str)

    c_code = Template(c_code_str)

    return c_code.render(
        kernel_src=lp.generate_code_v2(kernel).device_code().replace('\n',
            '\\n"\n"'),
        kernel_name=kernel.name,
        args=kernel.args,
        lsize=lsize,
        gsize=gsize)


n = 10


knl = lp.make_kernel(
        "{[i]: 0 <= i < %d}" % n,
        """
        b[i] = 2.0*a[i]
        """,
        [
            lp.GlobalArg('a', dtype=np.float64),
            ...],
        lang_version=(2018, 2))
knl = lp.split_iname(knl, 'i', 1, outer_tag="g.0", inner_tag="l.0")

code_to_compile = loopy_to_viennacl(knl)

a_petsc = PETSc.Vec().create(PETSc.COMM_WORLD)
a_petsc.createWithArray(np.empty(n))
a_petsc.setType('seqviennacl')
a_petsc.set(10.0)
a_petsc.view()

b_petsc = PETSc.Vec().create(PETSc.COMM_WORLD)
b_petsc.createWithArray(np.empty(n))
b_petsc.setType('seqviennacl')
b_petsc.set(0)
b_petsc.view()

compiler = coffee.system.compiler
extension = 'cpp'
cppargs = ['-I/home/kgk2/pack/petsc/include',
        '-I/home/kgk2/pack/petsc/arch-linux2-c-debug/include',
        '-I/home/kgk2/pack/firedrake/src/PyOP2/pyop2', '-msse']
ldargs = ['-L/home/kgk2/pack/petsc/lib',
        '-L/home/kgk2/pack/petsc/arch-linux2-c-debug/lib',
        '-Wl,-rpath,/home/kgk2/pack/petsc/lib',
        '-Wl,-rpath,/home/kgk2/pack/petsc/arch-linux2-c-debug/lib', '-lpetsc',
        '-lm', '-lOpenCL']
argtypes = [ctypes.c_void_p, ctypes.c_void_p]

func = compilation.load(code_to_compile,
                             extension,
                             knl.name+'_viennacl',
                             cppargs=cppargs,
                             ldargs=ldargs,
                             argtypes=argtypes,
                             restype=ctypes.c_int,
                             compiler='mpi',
                             comm=a_petsc.getComm())
func(a_petsc.handle, b_petsc.handle)
b_petsc.assemble()
a_petsc.assemble()

b_petsc.view()
a_petsc.view()