def interpolate_cached(compiled_expression, target_func):
kernel = compiled_expression.module.tabulate_expression
# Register complex types
cffi_support.register_type(ffi.typeof('double _Complex'),
numba.types.complex128)
cffi_support.register_type(ffi.typeof('float _Complex'),
numba.types.complex64)
reference_geometry = np.asarray(compiled_expression.fiat_element.ref_el.get_vertices())
# Unpack mesh and dofmap data
mesh = target_func.function_space.mesh
geom_dofmap = mesh.geometry.dofmap.array
geom_pos = mesh.geometry.dofmap.offsets
geom = mesh.geometry.x
gdim = mesh.geometry.dim
dofmap = target_func.function_space.dofmap.list.array
# Prepare coefficients and their dofmaps
# Global vectors and dofmaps are prepared here, local are
# fetched inside hot cell-loop
# Number of coefficients in ffcx-processed ufl form
num_coeffs = compiled_expression.module.num_coefficients
# Positions of ffcx-preprocessed coefficients in original form
cpos = compiled_expression.module.original_coefficient_positions
coeffs = ufl.algorithms.analysis.extract_coefficients(compiled_expression.expr)
coeffs_dofmaps = List.empty_list(numba.types.Array(numba.typeof(dofmap[0]), 1, "C", readonly=True))
coeffs_vectors = List.empty_list(numba.types.Array(numba.typeof(PETSc.ScalarType()), 1, "C", readonly=True))
for i in range(num_coeffs):
coeffs_dofmaps.append(coeffs[cpos[i]].function_space.dofmap.list.array)
coeffs_vectors.append(np.asarray(coeffs[cpos[i]].vector))
local_coeffs_sizes = np.asarray([coeff.function_space.element.space_dimension()
for coeff in coeffs], dtype=np.int)
local_coeffs_size = np.sum(local_coeffs_sizes, dtype=np.int)
# Prepare and pack constants
constants = ufl.algorithms.analysis.extract_constants(compiled_expression.expr)
constants_vector = np.array([], dtype=PETSc.ScalarType())
if len(constants) > 0:
constants_vector = np.hstack([c.value.flatten() for c in constants])
# Num DOFs of the target element
local_b_size = target_func.function_space.element.space_dimension()
value_size = int(np.product(compiled_expression.expr.ufl_shape))
num_coeffs = len(coeffs_vectors)
with target_func.vector.localForm() as b:
b.set(0.0)
assemble_vector_ufc(np.asarray(b), kernel, (geom_dofmap, geom_pos, geom), dofmap,
coeffs_vectors, coeffs_dofmaps, constants_vector, reference_geometry,
local_coeffs_sizes, local_coeffs_size, local_b_size, gdim, value_size)
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
f = ufl.as_vector([0.0, 1.0 / 16])
b1 = -ufl.inner(f, v) * ds(1)
# JIT compile individual blocks tabulation kernels
ufc_form00 = dolfinx.jit.ffcx_jit(mesh.mpi_comm(), a00)
kernel00 = ufc_form00.create_cell_integral(-1).tabulate_tensor
ufc_form01 = dolfinx.jit.ffcx_jit(mesh.mpi_comm(), a01)
kernel01 = ufc_form01.create_cell_integral(-1).tabulate_tensor
ufc_form10 = dolfinx.jit.ffcx_jit(mesh.mpi_comm(), a10)
kernel10 = ufc_form10.create_cell_integral(-1).tabulate_tensor
ffi = cffi.FFI()
cffi_support.register_type(ffi.typeof('double _Complex'),
numba.types.complex128)
c_signature = numba.types.void(
numba.types.CPointer(numba.typeof(PETSc.ScalarType())),
numba.types.CPointer(numba.typeof(PETSc.ScalarType())),
numba.types.CPointer(numba.typeof(PETSc.ScalarType())),
numba.types.CPointer(numba.types.double),
numba.types.CPointer(numba.types.int32),
numba.types.CPointer(numba.types.uint8), numba.types.uint32)
@numba.cfunc(c_signature, nopython=True)
def tabulate_condensed_tensor_A(A_,
w_,
c_,
coords_,
print("Could not load PETSc library for CFFI (ABI mode).")
raise
# 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
ADD_VALUES = PETSc.InsertMode.ADD_VALUES
# 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:
petsc_lib_cffi = ffi.dlopen(petsc_lib_name)
else:
try:
petsc_lib_cffi = ffi.dlopen(
os.path.join(petsc_dir, petsc_arch, "lib", "libpetsc.so"))
except OSError:
petsc_lib_cffi = ffi.dlopen(
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)
def interpolate_compiled(compiled_expression, target_func):
"""Compiled interpolation
Interpolates UFL expression into target function using FFCx
code generation and compilation.
Note
----
Works only for affine-mapped point-evaluation finite elements, e.g.
lagrange/discontinuous lagrange of arbitrary order.
"""
kernel = compiled_expression.module[0].tabulate_tensor_float64
# Register complex types
cffi_support.register_type(ffi.typeof('double _Complex'),
numba.types.complex128)
cffi_support.register_type(ffi.typeof('float _Complex'),
numba.types.complex64)
# Unpack mesh and dofmap data
mesh = target_func.function_space.mesh
geom_dofmap = mesh.geometry.dofmap.array
geom_pos = mesh.geometry.dofmap.offsets
geom = mesh.geometry.x
gdim = mesh.geometry.dim
dofmap = target_func.function_space.dofmap.list.array
# Prepare coefficients and their dofmaps
# Global vectors and dofmaps are prepared here, local are
# fetched inside hot cell-loop
# Number of coefficients in ffcx-processed ufl form
num_coeffs = compiled_expression.module[0].num_coefficients
# Positions of ffcx-preprocessed coefficients in original form
cpos = compiled_expression.module[0].original_coefficient_positions
coeffs = ufl.algorithms.analysis.extract_coefficients(compiled_expression.expr)
coeffs_dofmaps = List.empty_list(numba.types.Array(numba.typeof(dofmap[0]), 1, "C", readonly=True))
coeffs_vectors = List.empty_list(numba.types.Array(numba.typeof(PETSc.ScalarType()), 1, "C", readonly=True))
coeffs_bs = List.empty_list(numba.types.int_)
for i in range(num_coeffs):
coeffs_dofmaps.append(coeffs[cpos[i]].function_space.dofmap.list.array)
coeffs_vectors.append(np.asarray(coeffs[cpos[i]].vector))
coeffs_bs.append(coeffs[cpos[i]].function_space.dofmap.bs)
coeffs_sizes = np.asarray([coeff.function_space.element.space_dimension
for coeff in coeffs], dtype=np.int_)
# Prepare and pack constants
constants = ufl.algorithms.analysis.extract_constants(compiled_expression.expr)
constants_vector = np.array([], dtype=PETSc.ScalarType())
if len(constants) > 0:
constants_vector = np.hstack([c.value.flatten() for c in constants])
value_size = int(np.product(compiled_expression.expr.ufl_shape))
basix_space_dim = compiled_expression.ffcx_element.dim
space_dim = target_func.function_space.element.space_dimension
dofmap_bs = target_func.function_space.dofmap.bs
element_bs = target_func.function_space.dofmap.dof_layout.block_size
# Prepare mapping of subelements into the parent finite element
# This mapping stores also how dofs are collapsed when symmetry to a TensorElement
# is applied
if hasattr(compiled_expression.target_el, "flattened_sub_element_mapping"):
subel_map = np.array(compiled_expression.target_el.flattened_sub_element_mapping())
else:
subel_map = np.array(range(value_size))
num_coeffs = len(coeffs_vectors)
with target_func.vector.localForm() as b:
b.set(0.0)
assemble_vector_ufc(np.asarray(b), kernel, (geom_dofmap, geom_pos, geom), dofmap,
coeffs_vectors, coeffs_dofmaps, coeffs_bs, constants_vector,
coeffs_sizes, gdim, basix_space_dim, space_dim,
value_size, subel_map, dofmap_bs, element_bs)
# silly pointer interface
_lk_mkl_spexport_p = clib.lk_mkl_spexport_p
_lk_mkl_spe_free = clib.lk_mkl_spe_free
_lk_mkl_spe_nrows = clib.lk_mkl_spe_nrows
_lk_mkl_spe_ncols = clib.lk_mkl_spe_ncols
_lk_mkl_spe_row_sp = clib.lk_mkl_spe_row_sp
_lk_mkl_spe_row_ep = clib.lk_mkl_spe_row_ep
_lk_mkl_spe_colinds = clib.lk_mkl_spe_colinds
_lk_mkl_spe_values = clib.lk_mkl_spe_values
except OSError:
clib = None
# support intptr_t
if cffi_utils is not None:
# if we don't have working Numba, skip this
cffi_utils.register_type(ffi.typeof('intptr_t'), nt.intp)
# extract sizes
_int_size = ffi.sizeof('int')
_dbl_size = ffi.sizeof('double')
_mkl_errors = [
'SPARSE_STATUS_SUCCESS', 'SPARSE_STATUS_NOT_INITIALIZED',
'SPARSE_STATUS_ALLOC_FAILED', 'SPARSE_STATUS_INVALID_VALUE',
'SPARSE_STATUS_EXECUTION_FAILED', 'SPARSE_STATUS_INTERNAL_ERROR',
'SPARSE_STATUS_NOT_SUPPORTED'
]
def _mkl_check_return(rv, call='<unknown>'):
if rv:
__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)
ERRORS = {}
from ffcx import fiatinterface
from ffcx.codegeneration import jit as ffcx_jit
import cffi
import scipy.sparse
import numpy
import numba
from numba.core.typing import cffi_utils
# Register C complex types
ffi = cffi.FFI()
cffi_utils.register_type(ffi.typeof('double _Complex'), numba.types.complex128)
cffi_utils.register_type(ffi.typeof('float _Complex'), numba.types.complex64)
def c_to_numpy(ctype):
c2numpy = {
'double': numpy.float64,
'float': numpy.float32,
'double complex': numpy.complex128,
'float complex': numpy.complex64,
'long double': numpy.longdouble
}
return c2numpy.get(ctype)
def numpy_to_c(dtype):
numpy2c = {
numpy.float64: 'double',
numpy.float32: 'float',
numpy.complex128: 'double complex',
numpy.complex64: 'float complex',
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)
