petsc_ctype.py

from petsc4py import PETSc
import numpy as np
import coffee.system
from pyop2 import compilation
import ctypes


n = 10

a_petsc = PETSc.Vec().create(PETSc.COMM_WORLD)
b_petsc = PETSc.Vec().create(PETSc.COMM_WORLD)
a_petsc.createWithArray(np.empty(n))
a_petsc.setType('seqviennacl')
a_petsc.set(1)
a_petsc.setValues(range(n), np.arange(n)**2)
# a_petsc.view()

b_petsc.createWithArray(np.empty(n))
b_petsc.setType('seqviennacl')
b_petsc.set(7)
# 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]

code_to_compile = ("""
#include "petsc.h"
#include "petscvec.h"

extern "C" void petsc_try_func(Vec *x, Vec *y) {
    VecAXPY(*y, 6.0, *x);
    return ;
}""")


axpy_func = compilation.load(code_to_compile,
                             extension,
                             "petsc_try_func",
                             cppargs=cppargs,
                             ldargs=ldargs,
                             argtypes=argtypes,
                             restype=ctypes.c_int,
                             compiler='gnu',
                             comm=a_petsc.getComm())
axpy_func(a_petsc.handle, b_petsc.handle)

b_petsc.view()
a_petsc.view()