1# Copyright (c) 2017, Lawrence Livermore National Security, LLC. Produced at 2# the Lawrence Livermore National Laboratory. LLNL-CODE-734707. All Rights 3# reserved. See files LICENSE and NOTICE for details. 4# 5# This file is part of CEED, a collection of benchmarks, miniapps, software 6# libraries and APIs for efficient high-order finite element and spectral 7# element discretizations for exascale applications. For more information and 8# source code availability see http://github.com/ceed. 9# 10# The CEED research is supported by the Exascale Computing Project 17-SC-20-SC, 11# a collaborative effort of two U.S. Department of Energy organizations (Office 12# of Science and the National Nuclear Security Administration) responsible for 13# the planning and preparation of a capable exascale ecosystem, including 14# software, applications, hardware, advanced system engineering and early 15# testbed platforms, in support of the nation's exascale computing imperative. 16 17from _ceed_cffi import ffi, lib 18import tempfile 19import numpy as np 20import contextlib 21from .ceed_constants import MEM_HOST, USE_POINTER, COPY_VALUES 22 23# ------------------------------------------------------------------------------ 24 25 26class QFunctionContext(): 27 """Ceed QFunction Context: stores Ceed QFunction user context data.""" 28 29 # Attributes 30 _ceed = None 31 _pointer = ffi.NULL 32 _array_reference = None 33 34 # Constructor 35 def __init__(self, ceed): 36 # CeedQFunctionContext object 37 self._pointer = ffi.new("CeedQFunctionContext *") 38 39 # Reference to Ceed 40 self._ceed = ceed 41 42 # libCEED call 43 err_code = lib.CeedQFunctionContextCreate( 44 self._ceed._pointer[0], self._pointer) 45 self._ceed._check_error(err_code) 46 47 # Destructor 48 def __del__(self): 49 # libCEED call 50 err_code = lib.CeedQFunctionContextDestroy(self._pointer) 51 self._ceed._check_error(err_code) 52 53 # Representation 54 def __repr__(self): 55 return "<CeedQFunctionContext instance at " + hex(id(self)) + ">" 56 57 # String conversion for print() to stdout 58 def __str__(self): 59 """View a QFunction Context via print().""" 60 61 # libCEED call 62 fmt = ffi.new("char[]", "%f".encode('ascii')) 63 with tempfile.NamedTemporaryFile() as key_file: 64 with open(key_file.name, 'r+') as stream_file: 65 stream = ffi.cast("FILE *", stream_file) 66 67 err_code = lib.CeedQFunctionContextView( 68 self._pointer[0], stream) 69 self._ceed._check_error(err_code) 70 71 stream_file.seek(0) 72 out_string = stream_file.read() 73 74 return out_string 75 76 # Set QFunction Context's data 77 def set_data(self, data, memtype=MEM_HOST, cmode=COPY_VALUES): 78 """Set the data used by a QFunction Context, freeing any previously allocated 79 data if applicable. 80 81 Args: 82 *data: Numpy or Numba array to be used 83 **memtype: memory type of the array being passed, default CEED_MEM_HOST 84 **cmode: copy mode for the array, default CEED_COPY_VALUES""" 85 86 # Store array reference if needed 87 if cmode == USE_POINTER: 88 self._array_reference = data 89 else: 90 self._array_reference = None 91 92 # Setup the numpy array for the libCEED call 93 data_pointer = ffi.new("CeedScalar *") 94 if memtype == MEM_HOST: 95 data_pointer = ffi.cast( 96 "void *", 97 data.__array_interface__['data'][0]) 98 else: 99 array_pointer = ffi.cast( 100 "void *", 101 data.__cuda_array_interface__['data'][0]) 102 103 # libCEED call 104 err_code = lib.CeedQFunctionContextSetData( 105 self._pointer[0], 106 memtype, 107 cmode, 108 len(data) * ffi.sizeof("CeedScalar"), 109 data_pointer) 110 self._ceed._check_error(err_code) 111 112 # Get QFunction Context's data 113 def get_data(self, memtype=MEM_HOST): 114 """Get read/write access to a QFunction Context via the specified memory type. 115 116 Args: 117 **memtype: memory type of the array being passed, default CEED_MEM_HOST 118 119 Returns: 120 *data: Numpy or Numba array""" 121 122 # Retrieve the length of the array 123 size_pointer = ffi.new("CeedInt *") 124 err_code = lib.CeedQFunctionContextGetContextSize( 125 self._pointer[0], size_pointer) 126 self._ceed._check_error(err_code) 127 128 # Setup the pointer's pointer 129 data_pointer = ffi.new("CeedScalar **") 130 131 # libCEED call 132 err_code = lib.CeedQFunctionContextGetData( 133 self._pointer[0], memtype, data_pointer) 134 self._ceed._check_error(err_code) 135 136 # Return array created from buffer 137 if memtype == MEM_HOST: 138 # Create buffer object from returned pointer 139 buff = ffi.buffer( 140 data_pointer[0], 141 ffi.sizeof("CeedScalar") * 142 length_pointer[0]) 143 # return Numpy array 144 return np.frombuffer(buff, dtype="float64") 145 else: 146 # CUDA array interface 147 # https://numba.pydata.org/numba-doc/latest/cuda/cuda_array_interface.html 148 import numba.cuda as nbcuda 149 desc = { 150 'shape': (length_pointer[0]), 151 'typestr': '>f8', 152 'data': (int(ffi.cast("intptr_t", data_pointer[0])), False), 153 'version': 2 154 } 155 # return Numba array 156 return nbcuda.from_cuda_array_interface(desc) 157 158 # Restore the QFunction Context's data 159 def restore_data(self): 160 """Restore an array obtained using get_data().""" 161 162 # Setup the pointer's pointer 163 data_pointer = ffi.new("CeedScalar **") 164 165 # libCEED call 166 err_code = lib.CeedQFunctionDataRestoreData( 167 self._pointer[0], data_pointer) 168 self._ceed._check_error(err_code) 169 170 @contextlib.contextmanager 171 def data(self, *shape, memtype=MEM_HOST): 172 """Context manager for array access. 173 174 Args: 175 shape (tuple): shape of returned numpy.array 176 **memtype: memory type of the data being passed, default CEED_MEM_HOST 177 178 179 Returns: 180 np.array: writable view of QFunction Context 181 """ 182 x = self.get_data(memtype=memtype) 183 if shape: 184 x = x.reshape(shape) 185 yield x 186 self.restore_data() 187 188# ------------------------------------------------------------------------------ 189