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)
