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)