def lapoflap3d(in_field, tmp_field, out_field, num_halo=2): """ Compute Laplacian of the Laplacian in i-, j- and k-direction using 2nd-order centered differences based on a numba stencil kernel. Parameters ---------- in_field : input field (nx x ny x nz). tmp_field : intermediate result (must be of same size as in_field). out_field : result (must be of same size as in_field). num_halo : number of halo points. Returns ------- out_field : Laplacian-of-Laplacian of the input field computed in i-, j- and k- direction. """ # closure inlining def laplacian3d_kernel(in_field): """ laplacian kernel which is passed to the numba stencil function """ return ( -6.0 * in_field[0, 0, 0] + in_field[-1, 0, 0] + in_field[+1, 0, 0] + in_field[0, -1, 0] + in_field[0, +1, 0] + in_field[0, 0, -1] + in_field[0, 0, +1] ) tmp_field = stencil( laplacian3d_kernel, neighborhood=( (-num_halo + 1, num_halo - 1), (-num_halo + 1, num_halo - 1), (-num_halo + 1, num_halo - 1), ), )(in_field, out=tmp_field) out_field = stencil( laplacian3d_kernel, neighborhood=( (-num_halo, num_halo), (-num_halo, num_halo), (-num_halo, num_halo), ), )(tmp_field, out=out_field) return out_field
def laplacian1d(in_field, out_field, num_halo=1): """ Compute the Laplacian of the in_field in i-direction based on a numba stencil kernel for the 2nd-order centered differences. Parameters ---------- in_field : input field (nx x ny x nz). out_field : result (must be of same size as in_field). num_halo : number of halo points. Returns ------- out_field : in_field with Laplacian computed in i-direction. """ # closure inlining def laplacian1d_kernel(in_field): """ laplacian kernel which is passed to the numba stencil function """ return -2.0 * in_field[0, 0, 0] + in_field[-1, 0, 0] + in_field[+1, 0, 0] out_field = stencil( laplacian1d_kernel, neighborhood=( (-num_halo, num_halo), (-num_halo, num_halo), (-num_halo, num_halo), ), )(in_field, out=out_field) return out_field
def FMA(in_field, in_field2, in_field3, out_field, num_halo=0): """ Pointwise stencil to test for fused multiply-add based on a numba stencil kernel Parameters ---------- in_field,in_field2, in_field3 : input field (nx x ny x nz). out_field : result (must be of same size as in_field). num_halo : number of halo points. Returns ------- out_field : fused multiply-add applied to in_field. """ # closure inlining def FMA_kernel(in_field, in_field2, in_field3): """ kernel which is passed to the numba stencil function """ return in_field[0, 0, 0] + in_field2[0, 0, 0] * in_field3[0, 0, 0] out_field = stencil( FMA_kernel, neighborhood=( (-num_halo, num_halo), (-num_halo, num_halo), (-num_halo, num_halo), ), )(in_field, in_field2, in_field3, out=out_field) return out_field
def f(A, B, f): numba.stencil(f)(A, out=B)
def execute_user_code(self, user_args, user_kwargs, tracing): # type: (list, dict, bool) -> (object, COMPSsException) """ Executes the user code. Disables the tracing hook if tracing is enabled. Restores it at the end of the user code execution. :param user_args: Function args. :param user_kwargs: Function kwargs. :param tracing: If tracing enabled. :return: The user function returns and the compss exception (if any). """ # Tracing hook is disabled by default during the user code of the task. # The user can enable it with tracing_hook=True in @task decorator for # specific tasks or globally with the COMPSS_TRACING_HOOK=true # environment variable. restore_hook = False pro_f = None if tracing: global_tracing_hook = False if TRACING_HOOK_ENV_VAR in os.environ: hook_enabled = os.environ[TRACING_HOOK_ENV_VAR] == "true" global_tracing_hook = hook_enabled if self.decorator_arguments['tracing_hook'] or global_tracing_hook: # The user wants to keep the tracing hook pass else: # When Extrae library implements the function to disable, # use it, as: # import pyextrae # pro_f = pyextrae.shutdown() # Since it is not available yet, we manage the tracing hook # by ourselves pro_f = sys.getprofile() sys.setprofile(None) restore_hook = True user_returns = None compss_exception = None if self.decorator_arguments['numba']: # Import all supported functionalities from numba import jit from numba import njit from numba import generated_jit from numba import vectorize from numba import guvectorize from numba import stencil from numba import cfunc numba_mode = self.decorator_arguments['numba'] numba_flags = self.decorator_arguments['numba_flags'] if type(numba_mode) is dict: # Use the flags defined by the user numba_flags['cache'] = True # Always force cache user_returns = jit(self.user_function, **numba_flags)(*user_args, **user_kwargs) elif numba_mode is True or numba_mode == 'jit': numba_flags['cache'] = True # Always force cache user_returns = jit(self.user_function, **numba_flags)(*user_args, **user_kwargs) # Alternative way of calling: # user_returns = jit(cache=True)(self.user_function) \ # (*user_args, **user_kwargs) elif numba_mode == 'generated_jit': user_returns = generated_jit(self.user_function, **numba_flags)(*user_args, **user_kwargs) elif numba_mode == 'njit': numba_flags['cache'] = True # Always force cache user_returns = njit(self.user_function, **numba_flags)(*user_args, **user_kwargs) elif numba_mode == 'vectorize': numba_signature = self.decorator_arguments['numba_signature'] # noqa: E501 user_returns = vectorize( numba_signature, **numba_flags )(self.user_function)(*user_args, **user_kwargs) elif numba_mode == 'guvectorize': numba_signature = self.decorator_arguments['numba_signature'] # noqa: E501 numba_decl = self.decorator_arguments['numba_declaration'] user_returns = guvectorize( numba_signature, numba_decl, **numba_flags )(self.user_function)(*user_args, **user_kwargs) elif numba_mode == 'stencil': user_returns = stencil( **numba_flags )(self.user_function)(*user_args, **user_kwargs) elif numba_mode == 'cfunc': numba_signature = self.decorator_arguments['numba_signature'] # noqa: E501 user_returns = cfunc( numba_signature )(self.user_function).ctypes(*user_args, **user_kwargs) else: raise Exception("Unsupported numba mode.") else: try: # Normal task execution user_returns = self.user_function(*user_args, **user_kwargs) except COMPSsException as ce: compss_exception = ce # Check old targetDirection if 'targetDirection' in self.decorator_arguments: target_label = 'targetDirection' else: target_label = 'target_direction' compss_exception.target_direction = self.decorator_arguments[target_label] # noqa: E501 # Reestablish the hook if it was disabled if restore_hook: sys.setprofile(pro_f) return user_returns, compss_exception
def f(A, B, f): # pragma: no cover numba.stencil(f)(A, out=B)