def set_openmp_cores(cores):
# PYCLIBRARY ISSUE
# TODO: Would be really nice to be able to directly pass
# a list here instead of needing to specify _array+type.
# Should be able to do it safely for any const* argument
# Note that the base type of the pointer type could be
# derived via a reverse lookup table.
# E.g., Inverse of POINTER(c_int)-->LP_c_int
_check(_bf.bfAffinitySetOpenMPCores(len(cores), _array(cores, 'int')))
def map(func_string, data, axis_names=None, shape=None,
func_name=None, extra_code=None,
block_shape=None, block_axes=None):
"""Apply a function to a set of ndarrays.
Args:
func_string (str): The function to apply to the arrays, as a string (see
below for examples).
data (dict): Map of string names to ndarrays or scalars.
axis_names (list): List of string names by which each axis is referenced
in func_string.
shape: The shape of the computation. If None, the broadcast shape
of all data arrays is used.
func_name (str): Name of the function, for debugging purposes.
extra_code (str): Additional code to be included at global scope.
block_shape: The 2D shape of the thread block (y,x) with which the kernel
is launched.
This is a performance tuning parameter.
If NULL, a heuristic is used to select the block shape.
Changes to this parameter do _not_ require re-compilation of
the kernel.
block_axes: List of axis indices (or names) specifying the 2 computation
axes to which the thread block (y,x) is mapped.
This is a performance tuning parameter.
If NULL, a heuristic is used to select the block axes.
Values may be negative for reverse indexing.
Changes to this parameter _do_ require re-compilation of the
kernel.
Note:
Only GPU computation is currently supported.
Examples::
# Add two arrays together
bf.map("c = a + b", {'c': c, 'a': a, 'b': b})
# Compute outer product of two arrays
bf.map("c(i,j) = a(i) * b(j)",
{'c': c, 'a': a, 'b': b},
axis_names=('i','j'))
# Split the components of a complex array
bf.map("a = c.real; b = c.imag", {'c': c, 'a': a, 'b': b})
# Raise an array to a scalar power
bf.map("c = pow(a, p)", {'c': c, 'a': a, 'p': 2.0})
# Slice an array with a scalar index
bf.map("c(i) = a(i,k)", {'c': c, 'a': a, 'k': 7}, ['i'], shape=c.shape)
"""
try:
func_string = func_string.encode()
if func_name is not None:
func_name = func_name.encode()
if extra_code is not None:
extra_code = extra_code.encode()
except AttributeError:
# Python2 catch
pass
narg = len(data)
ndim = len(shape) if shape is not None else 0
arg_arrays = []
args = []
arg_names = []
if block_axes is not None:
# Allow referencing axes by name
block_axes = [axis_names.index(bax) if isinstance(bax, str)
else bax
for bax in block_axes]
if block_axes is not None and len(block_axes) != 2:
raise ValueError("block_axes must contain exactly 2 entries")
if block_shape is not None and len(block_shape) != 2:
raise ValueError("block_shape must contain exactly 2 entries")
for key, arg in data.items():
arg = _convert_to_array(arg)
# Note: We must keep a reference to each array lest they be garbage
# collected before their corresponding BFarray is used.
arg_arrays.append(arg)
args.append(arg.as_BFarray())
arg_names.append(key)
_check(_bf.bfMap(ndim, _array(shape, dtype=ctypes.c_long),
_array(axis_names),
narg, _array(args), _array(arg_names),
func_name, func_string, extra_code,
_array(block_shape), _array(block_axes)))