def get_matsetvalues_api():
"""Make MatSetValuesLocal from PETSc available via cffi in API mode"""
worker = os.getenv('PYTEST_XDIST_WORKER', None)
import uuid
r = uuid.uuid1().int >> 64
module_name = f"_petsc_cffi_{r}_{worker}"
if MPI.COMM_WORLD.Get_rank() == 0:
ffibuilder = cffi.FFI()
ffibuilder.cdef("""
typedef int... PetscInt;
typedef ... PetscScalar;
typedef int... InsertMode;
int MatSetValuesLocal(void* mat, PetscInt nrow, const PetscInt* irow,
PetscInt ncol, const PetscInt* icol,
const PetscScalar* y, InsertMode addv);
""")
ffibuilder.set_source(
module_name,
"""
# include "petscmat.h"
""",
libraries=['petsc'],
include_dirs=[
os.path.join(petsc_dir, petsc_arch, 'include'),
os.path.join(petsc_dir, 'include')
] + dolfinx_pc["include_dirs"],
library_dirs=[os.path.join(petsc_dir, petsc_arch, 'lib')],
extra_compile_args=[])
# Build module in same directory as test file
# path = pathlib.Path(__file__).parent.absolute()
# Build module in same directory
path = pathlib.Path(os.getcwd())
ffibuilder.compile(tmpdir=path, verbose=True)
MPI.COMM_WORLD.Barrier()
spec = importlib.util.find_spec(module_name, path.as_posix())
if spec is None:
raise ImportError("Failed to find CFFI generated module")
try:
module = importlib.util.module_from_spec(spec)
except:
print("unable to import module from spec =", spec)
import ipdb
ipdb.set_trace()
cffi_support.register_module(module)
cffi_support.register_type(module.ffi.typeof("PetscScalar"),
numba_scalar_t)
return module.lib.MatSetValuesLocal
def get_matsetvalues_api():
"""Make MatSetValuesLocal from PETSc available via cffi in API mode"""
if dolfinx.pkgconfig.exists("dolfinx"):
dolfinx_pc = dolfinx.pkgconfig.parse("dolfinx")
else:
raise RuntimeError("Could not find DOLFINx pkgconfig file")
worker = os.getenv('PYTEST_XDIST_WORKER', None)
module_name = "_petsc_cffi_{}".format(worker)
if MPI.COMM_WORLD.Get_rank() == 0:
ffibuilder = cffi.FFI()
ffibuilder.cdef("""
typedef int... PetscInt;
typedef ... PetscScalar;
typedef int... InsertMode;
int MatSetValuesLocal(void* mat, PetscInt nrow, const PetscInt* irow,
PetscInt ncol, const PetscInt* icol,
const PetscScalar* y, InsertMode addv);
""")
ffibuilder.set_source(
module_name,
"""
#include "petscmat.h"
""",
libraries=['petsc'],
include_dirs=[
os.path.join(petsc_dir, petsc_arch, 'include'),
os.path.join(petsc_dir, 'include')
] + dolfinx_pc["include_dirs"],
library_dirs=[os.path.join(petsc_dir, petsc_arch, 'lib')],
extra_compile_args=[])
# Build module in same directory as test file
path = pathlib.Path(__file__).parent.absolute()
ffibuilder.compile(tmpdir=path, verbose=True)
MPI.COMM_WORLD.Barrier()
spec = importlib.util.find_spec(module_name)
if spec is None:
raise ImportError("Failed to find CFFI generated module")
module = importlib.util.module_from_spec(spec)
cffi_support.register_module(module)
cffi_support.register_type(module.ffi.typeof("PetscScalar"),
numba_scalar_t)
return module.lib.MatSetValuesLocal
moduleauthor:: J. Derek Tucker <*****@*****.**>
"""
import numba
from numba.core.typing import cffi_utils
from _DP import ffi, lib
import _DP
from numpy import linspace, interp, zeros, diff, double, sqrt, arange, float64, int32, int64, trapz, ones
from numpy import zeros, frombuffer, ascontiguousarray, empty, load, roll, dot, eye, arccos, reshape, float32
from numpy import kron, floor
from numpy.linalg import norm, svd, det, solve
DP = lib.DP
cffi_utils.register_module(_DP)
@numba.jit()
def grad(f, binsize):
n = f.shape[0]
g = zeros(n)
h = binsize * arange(1, n + 1)
g[0] = (f[1] - f[0]) / (h[1] - h[0])
g[-1] = (f[-1] - f[(-2)]) / (h[-1] - h[-2])
h = h[2:] - h[0:-2]
g[1:-1] = (f[2:] - f[0:-2]) / h[0]
return g
],
library_dirs=[os.path.join(petsc_dir, petsc_arch, 'lib')],
extra_compile_args=[])
# Build module in same directory as test file
path = pathlib.Path(__file__).parent.absolute()
ffibuilder.compile(tmpdir=path, verbose=False)
MPI.COMM_WORLD.Barrier()
spec = importlib.util.find_spec(module_name)
if spec is None:
raise ImportError("Failed to find CFFI generated module")
module = importlib.util.module_from_spec(spec)
cffi_support.register_module(module)
MatSetValues_api = module.lib.MatSetValuesLocal
cffi_support.register_type(module.ffi.typeof("PetscScalar"), numba_scalar_t)
# See https://github.com/numba/numba/issues/4036 for why we need 'sink'
@numba.njit
def sink(*args):
pass
@numba.njit
def area(x0, x1, x2) -> float:
"""Compute the area of a triangle embedded in 2D from the three vertices"""
a = (x1[0] - x2[0])**2 + (x1[1] - x2[1])**2
b = (x0[0] - x2[0])**2 + (x0[1] - x2[1])**2
def initialize_petsc() -> typing.Tuple[cffi.FFI, typing.Any]:
"""
Initialize petsc and CFFI for usage in numba
"""
# Get details of PETSc install
petsc_dir = PETSc_get_config()['PETSC_DIR']
petsc_arch = petsc4py.lib.getPathArchPETSc()[1]
# Get PETSc int and scalar types
cmplx = True if np.dtype(PETSc.ScalarType).kind == 'c' else False
scalar_size = np.dtype(PETSc.ScalarType).itemsize
index_size = np.dtype(PETSc.IntType).itemsize
if index_size == 8:
c_int_t = "int64_t"
ctypes_index = ctypes.c_int64 # type: ignore
elif index_size == 4:
c_int_t = "int32_t"
ctypes_index = ctypes.c_int32 # type: ignore
else:
raise RuntimeError(
"Cannot translate PETSc index size into a C type, index_size: {}."
.format(index_size))
if cmplx and scalar_size == 16:
c_scalar_t = "double _Complex"
numba_scalar_t = numba.types.complex128
elif cmplx and scalar_size == 8:
c_scalar_t = "float _Complex"
numba_scalar_t = numba.types.complex64
elif not cmplx and scalar_size == 8:
c_scalar_t = "double"
numba_scalar_t = numba.types.float64
elif not cmplx and scalar_size == 4:
c_scalar_t = "float"
numba_scalar_t = numba.types.float32
else:
raise RuntimeError(
"Cannot translate PETSc scalar type to a C type, complex: {} size: {}."
.format(complex, scalar_size))
# Load PETSc library via ctypes
petsc_lib_name = ctypes.util.find_library("petsc")
if petsc_lib_name is not None:
petsc_lib_ctypes = ctypes.CDLL(petsc_lib_name)
else:
try:
petsc_lib_ctypes = ctypes.CDLL(os.path.join(
petsc_dir, petsc_arch, "lib", "libpetsc.so"))
except OSError:
try:
petsc_lib_ctypes = ctypes.CDLL(os.path.join(
petsc_dir, petsc_arch, "lib", "libpetsc.dylib"))
except OSError:
raise RuntimeError("Could not load PETSc library for CFFI (ABI mode).")
# Get the PETSc MatSetValuesLocal function via ctypes
MatSetValues_ctypes = petsc_lib_ctypes.MatSetValuesLocal
MatSetValues_ctypes.argtypes = (ctypes.c_void_p, ctypes_index, ctypes.POINTER(
ctypes_index), ctypes_index, ctypes.POINTER(ctypes_index),
ctypes.c_void_p, ctypes.c_int)
del petsc_lib_ctypes
# CFFI - register complex types
ffi = cffi.FFI()
cffi_support.register_type(ffi.typeof(
'double _Complex'), numba.types.complex128)
cffi_support.register_type(ffi.typeof('float _Complex'), numba.types.complex64)
# Get MatSetValuesLocal from PETSc available via cffi in ABI mode
ffi.cdef("""int MatSetValuesLocal(void* mat, {0} nrow, const {0}* irow,
{0} ncol, const {0}* icol, const {1}* y, int addv);
""".format(c_int_t, c_scalar_t))
if petsc_lib_name is not None:
ffi.dlopen(petsc_lib_name)
else:
try:
ffi.dlopen(os.path.join(petsc_dir, petsc_arch, "lib", "libpetsc.so"))
except OSError:
try:
ffi.dlopen(os.path.join(petsc_dir, petsc_arch, "lib", "libpetsc.dylib"))
except OSError:
raise RuntimeError("Could not load PETSc library for CFFI (ABI mode).")
# Make MatSetValuesLocal from PETSc available via cffi in API mode
worker = os.getenv('ASSEMBLE_XDIST_WORKER', None)
module_name = "_petsc_cffi_{}".format(worker)
if MPI.COMM_WORLD.Get_rank() == 0:
os.environ["CC"] = "mpicc"
ffibuilder = cffi.FFI()
ffibuilder.cdef("""
typedef int... PetscInt;
typedef ... PetscScalar;
typedef int... InsertMode;
int MatSetValuesLocal(void* mat, PetscInt nrow, const PetscInt* irow,
PetscInt ncol, const PetscInt* icol,
const PetscScalar* y, InsertMode addv);
""")
ffibuilder.set_source(module_name, """
# include "petscmat.h"
""",
libraries=['petsc'],
include_dirs=[os.path.join(petsc_dir, petsc_arch,
'include'),
os.path.join(petsc_dir, 'include')],
library_dirs=[os.path.join(
petsc_dir, petsc_arch, 'lib')],
extra_compile_args=[])
# Build module in same directory as python script
ffibuilder.compile(".", verbose=False)
MPI.COMM_WORLD.Barrier()
module = importlib.import_module(module_name, ".")
cffi_support.register_module(module)
MatSetValuesLocal_api = module.lib.MatSetValuesLocal
cffi_support.register_type(module.ffi.typeof("PetscScalar"),
numba_scalar_t)
return ffi, MatSetValuesLocal_api
def load_ool_module():
"""
Compile an out-of-line module, return the corresponding ffi and
module objects.
"""
from cffi import FFI
numba_complex = """
typedef struct _numba_complex {
double real;
double imag;
} numba_complex;
"""
bool_define = """
#ifdef _MSC_VER
#define false 0
#define true 1
#define bool int
#else
#include <stdbool.h>
#endif
"""
defs = numba_complex + """
bool boolean();
double sin(double x);
double cos(double x);
int foo(int a, int b, int c);
void vsSin(int n, float* x, float* y);
void vdSin(int n, double* x, double* y);
void vector_real(numba_complex *c, double *real, int n);
void vector_imag(numba_complex *c, double *imag, int n);
"""
source = numba_complex + bool_define + """
static bool boolean()
{
return true;
}
static int foo(int a, int b, int c)
{
return a + b * c;
}
void vsSin(int n, float* x, float* y) {
int i;
for (i=0; i<n; i++)
y[i] = sin(x[i]);
}
void vdSin(int n, double* x, double* y) {
int i;
for (i=0; i<n; i++)
y[i] = sin(x[i]);
}
static void vector_real(numba_complex *c, double *real, int n) {
int i;
for (i = 0; i < n; i++)
real[i] = c[i].real;
}
static void vector_imag(numba_complex *c, double *imag, int n) {
int i;
for (i = 0; i < n; i++)
imag[i] = c[i].imag;
}
"""
ffi = FFI()
ffi.set_source('cffi_usecases_ool', source)
ffi.cdef(defs, override=True)
tmpdir = temp_directory('test_cffi')
ffi.compile(tmpdir=tmpdir)
sys.path.append(tmpdir)
try:
mod = import_dynamic('cffi_usecases_ool')
cffi_support.register_module(mod)
cffi_support.register_type(mod.ffi.typeof('struct _numba_complex'),
complex128)
return mod.ffi, mod
finally:
sys.path.remove(tmpdir)
A[2] = l2;
}
""", extra_compile_args=["-march=skylake", "-Ofast"])
ffibuilder.cdef("""
void tabulate_expression(double* restrict A,
const double* w,
const double* c,
const double* restrict coordinate_dofs);
""")
ffibuilder.compile(verbose=False)
from _cffi_kernelA import ffi, lib # noqa
import _cffi_kernelA # noqa
cffi_support.register_module(_cffi_kernelA)
@numba.njit
def run_pure(kernel):
warr = numpy.array([1.0, 2.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 1.0])
Aarr = numpy.zeros(3, dtype=numpy.double)
geometry = numpy.zeros((3, 2))
constants = numpy.zeros(1, dtype=numpy.double)
for i in range(num_cells):
kernel(ffi.from_buffer(Aarr), ffi.from_buffer(warr), ffi.from_buffer(constants), ffi.from_buffer(geometry))
run_pure(cexpr.module.tabulate_expression)
pure_ffcx_times = []
import numpy as np import numba import scipy.optimize as sopt import scipy.sparse as sspa from scipy.sparse.linalg import spsolve, spilu, splu from numba import cuda import cffi import numba.core.typing.cffi_utils as cffi_support ffi = cffi.FFI() import pmsm_mech_ctrl_cffi as jacs cffi_support.register_module(jacs) f_ini_eval = jacs.lib.f_ini_eval g_ini_eval = jacs.lib.g_ini_eval f_run_eval = jacs.lib.f_run_eval g_run_eval = jacs.lib.g_run_eval h_eval = jacs.lib.h_eval de_jac_ini_xy_eval = jacs.lib.de_jac_ini_xy_eval de_jac_ini_up_eval = jacs.lib.de_jac_ini_up_eval de_jac_ini_num_eval = jacs.lib.de_jac_ini_num_eval sp_jac_ini_xy_eval = jacs.lib.sp_jac_ini_xy_eval sp_jac_ini_up_eval = jacs.lib.sp_jac_ini_up_eval sp_jac_ini_num_eval = jacs.lib.sp_jac_ini_num_eval de_jac_run_xy_eval = jacs.lib.de_jac_run_xy_eval de_jac_run_up_eval = jacs.lib.de_jac_run_up_eval de_jac_run_num_eval = jacs.lib.de_jac_run_num_eval
from numba import njit, vectorize
from numba.core.typing.cffi_utils import register_module
from numba.types import float64
import _erfa_cffi
from .api import eraEpb as _eraEpb
register_module(_erfa_cffi)
@njit
def eraEpb(dj1, dj2):
return _eraEpb(dj1, dj2)
@vectorize([float64(float64, float64)])
def eraEpb_v(dj1, dj2):
return _eraEpb(dj1, dj2)
from sunode import _cvodes
__all__ = [
"lib", "ffi", "ERRORS", "Borrows", "notnull",
"check", "check_ptr", "check_code", "as_numpy"
]
logger = logging.getLogger("sunode.basic")
lib: Any = _cvodes.lib
ffi: Any = _cvodes.ffi
cffi_utils.register_module(_cvodes)
cffi_utils.register_type(
ffi.typeof("N_Vector").item, numba.types.Opaque("N_Vector")
)
cffi_utils.register_type(
ffi.typeof("SUNMatrix").item, numba.types.Opaque("SUNMatrix")
)
_data_dtype = cffi_utils.map_type(ffi.typeof("realtype"))
_index_dtype = cffi_utils.map_type(ffi.typeof("sunindextype"))
data_dtype: Any = np.dtype(_data_dtype.name)
index_dtype: Any = np.dtype(_index_dtype.name)
CPointer = NewType("CPointer", int)
from types import BuiltinFunctionType
import numba.types as nt
from numba.core.typing.cffi_utils import register_module, register_type
from numba import njit
from . import _mkl_ops
register_type(_mkl_ops.ffi.typeof('intptr_t'), nt.intp)
register_module(_mkl_ops)
ffi = _mkl_ops.ffi
# export all the LK names
__all__ = ['ffi']
for name in dir(_mkl_ops.lib):
f = getattr(_mkl_ops.lib, name)
if isinstance(f, BuiltinFunctionType):
globals()[name] = f
__all__.append(name)
