def argument(self, shape, dtype=np.float64):
"""
Return an instance of :class:`numloopy.ArraySymbol` which the loop
kernel expects as an input argument.
"""
if isinstance(shape, int):
shape = (shape, )
assert isinstance(shape, tuple)
inames = tuple(self.name_generator(based_on='i') for _ in shape)
arg_name = self.name_generator(based_on='arr')
rhs = Subscript(Variable(arg_name),
tuple(Variable(iname) for iname in inames))
subst_name = self.name_generator(based_on='subst')
self.register_substitution(lp.SubstitutionRule(subst_name, inames,
rhs))
self.substs_to_arrays[subst_name] = arg_name
self.data.append(lp.GlobalArg(name=arg_name, shape=shape, dtype=dtype))
return ArraySymbol(stack=self,
name=subst_name,
dtype=dtype,
shape=shape)
def arange(self, stop):
"""
Registers a substitution rule on to the stack with an array whose values
are filled equivalent to ``numpy.arange``.
:arg stop: An instance of :class:`int` denoting the extent of the
array.
:return: An instance of :class:`numloopy.ArraySymbol` of shape
``(stop,)``, corresponding to the substitution rule which was
registered.
"""
assert isinstance(stop, int)
subst_name = self.name_generator(based_on="subst")
arg = ArraySymbol(stack=self,
name=subst_name,
shape=(stop, ),
dtype=np.int)
iname = self.name_generator(based_on="i")
rhs = Variable(iname)
rule = lp.SubstitutionRule(subst_name, (iname, ), rhs)
self.register_substitution(rule)
return arg
def cumsum(self, arg):
"""
Registers a substitution rule in order to cumulatively sum the
elements of array ``arg`` along ``axis``. Mimics :func:`numpy.cumsum`.
:return: An instance of :class:`numloopy.ArraySymbol` which is
which is registered as the cumulative summed-substitution rule.
"""
# Note: this can remain as a substitution but loopy does not have
# support for translating inames for substitutions to the kernel
# domains
assert len(arg.shape) == 1
i_iname = self.name_generator(based_on="i")
j_iname = self.name_generator(based_on="i")
space = isl.Space.create_from_names(isl.DEFAULT_CONTEXT,
[i_iname, j_iname])
domain = isl.BasicSet.universe(space)
arg_name = self.name_generator(based_on="arr")
subst_name = self.name_generator(based_on="subst")
domain = domain & make_slab(space, i_iname, 0, arg.shape[0])
domain = domain.add_constraint(
isl.Constraint.ineq_from_names(space, {j_iname: 1}))
domain = domain.add_constraint(
isl.Constraint.ineq_from_names(space, {
j_iname: -1,
i_iname: 1,
1: -1
}))
cumsummed_arg = ArraySymbol(stack=self,
name=arg_name,
shape=arg.shape,
dtype=arg.dtype)
cumsummed_subst = ArraySymbol(stack=self,
name=subst_name,
shape=arg.shape,
dtype=arg.dtype)
subst_iname = self.name_generator(based_on="i")
rule = lp.SubstitutionRule(
subst_name, (subst_iname, ),
Subscript(Variable(arg_name), (Variable(subst_iname), )))
from loopy.library.reduction import SumReductionOperation
insn = lp.Assignment(assignee=Subscript(Variable(arg_name),
(Variable(i_iname), )),
expression=lp.Reduction(
SumReductionOperation(), (j_iname, ),
parse('{}({})'.format(arg.name, j_iname))))
self.data.append(cumsummed_arg)
self.substs_to_arrays[subst_name] = arg_name
self.register_implicit_assignment(insn)
self.domains.append(domain)
self.register_substitution(rule)
return cumsummed_subst
def sum(self, arg, axis=None):
"""
Registers a substitution rule in order to sum the elements of array
``arg`` along ``axis``.
:return: An instance of :class:`numloopy.ArraySymbol` which is
which is registered as the sum-substitution rule.
"""
if isinstance(axis, int):
axis = (axis, )
if not axis:
axis = tuple(range(len(arg.shape)))
inames = [self.name_generator(based_on="i") for _ in arg.shape]
space = isl.Space.create_from_names(isl.DEFAULT_CONTEXT, inames)
domain = isl.BasicSet.universe(space)
for axis_len, iname in zip(arg.shape, inames):
domain &= make_slab(space, iname, 0, axis_len)
self.domains.append(domain)
reduction_inames = tuple(iname for i, iname in enumerate(inames)
if i in axis)
left_inames = tuple(iname for i, iname in enumerate(inames)
if i not in axis)
def _one_if_empty(t):
if t:
return t
else:
return (1, )
subst_name = self.name_generator(based_on="subst")
summed_arg = ArraySymbol(
stack=self,
name=subst_name,
shape=_one_if_empty(
tuple(axis_len for i, axis_len in enumerate(arg.shape)
if i not in axis)),
dtype=arg.dtype)
from loopy.library.reduction import SumReductionOperation
rule = lp.SubstitutionRule(
subst_name, left_inames,
lp.Reduction(SumReductionOperation(), reduction_inames,
parse('{}({})'.format(arg.name, ', '.join(inames)))))
self.register_substitution(rule)
return summed_arg
def __getitem__(self, index):
"""
Registers a substitution rule pointing to the indices through `index`.
:arg index: An instance of :class:`tuple` of int, slices or
:class:`ArraySymbol`.
"""
if isinstance(index, Number):
index = (index, )
assert isinstance(index, tuple)
right_inames = []
left_inames = []
shape = []
for axis_len, idx in zip(self.shape, index):
if isinstance(idx, int):
right_inames.append(idx)
elif isinstance(idx, slice):
# right now only support complete slices
# future plan is to make it diverse by adding it more support
# its easy, anyone intereseted can do it, with a small change
# to make_slab, or wait for kernel_callables_v3 to go through
# into master
assert idx.start is None
assert idx.stop is None
assert idx.step is None
iname = self.stack.name_generator(based_on='i')
right_inames.append(Variable(iname))
left_inames.append(iname)
shape.append(axis_len)
else:
raise TypeError('can be subscripted only with slices or '
'integers')
rhs = Call(Variable(self.name), tuple(right_inames))
subst_name = self.stack.name_generator(based_on='subst')
self.stack.register_substitution(
lp.SubstitutionRule(subst_name, tuple(left_inames), rhs))
def _one_if_empty(t):
if t:
return t
else:
return (1, )
return ArraySymbol(stack=self.stack,
name=subst_name,
dtype=self.dtype,
shape=_one_if_empty(tuple(shape)))
def reshape(self, new_shape, order='C'):
"""
Registers a substitution rule to reshape array with the shape
``new_shape``. Mimics :func:`numpy.ndarray.reshape`.
:arg new_shape: An instance of :class:`tuple` of :class:`int`.
:arg order: Either 'C' or 'F'
"""
# need an error here complain if there is a shape mismatch
# how to do this:
# look at how loopy sets its dim tags, from shape and order.
subst_name = self.stack.name_generator(based_on='subst')
inames = tuple(
self.stack.name_generator(based_on="i") for _ in new_shape)
new_arg = self.copy(stack=self.stack,
name=subst_name,
shape=new_shape,
dim_tags=None,
order=order)
linearized_idx = sum(
Variable(iname) * dim_tag.stride
for iname, dim_tag in zip(inames, new_arg.dim_tags))
strides = tuple(dim_tag.stride for dim_tag in self.dim_tags)
if self.dim_tags[0].stride == 1:
pass
elif self.dim_tags[-1].stride == 1:
strides = strides[::-1]
else:
raise ValueError()
indices = []
for stride in strides[::-1]:
current_idx = linearized_idx // stride
indices.append(simplify_via_aff(current_idx))
linearized_idx -= current_idx * stride
# should we assert that the linearized index is 0?
rule = lp.SubstitutionRule(subst_name,
inames,
expression=Variable(
self.name)(*tuple(indices)))
self.stack.register_substitution(rule)
return ArraySymbol(stack=self.stack, name=subst_name, shape=new_shape)
def fill_array(shape, value, name_generator):
"""
Helper function to fill an array of shape ``shape`` with ``value`` .
:arg shape: An instance of :class:`tuple` denoting the shape of the array.
:arg value: The value with the array should be filled.
:arg name_generator: An instance of :class:`pytools.UniqueNameGenerator`,
for generating the name of the new array.
:return: A tuple of the name of the subtitution, the substitution rule and
the modified name generator.
"""
inames = tuple(name_generator(based_on='i') for _ in shape)
rhs = value
subst_name = name_generator(based_on='subst')
rule = lp.SubstitutionRule(subst_name, inames, rhs)
return subst_name, rule, name_generator
def _arithmetic_op(self, other, op):
"""
Registers a substitution rule that performs ``(self) op (other)``
element-wise for the arrays. :mod:`numpy` broadcasting rules are
followed while performing these operations.
:return: An instance of :class:`ArraySymbol` corresponding to the
registered substitution for the element-wise operation.
"""
assert op in ['+', '-', '*', '/', '<', '<=', '>', '>=']
assert self.stack == other.stack
def _apply_op(var1, var2):
if op == '+':
return var1 + var2, self.dtype
if op == '-':
return var1 - var2, self.dtype
if op == '*':
return var1 * var2, self.dtype
if op == '/':
return var1 / var2, self.dtype
if op == '<':
return var1.lt(var2), np.int
if op == '<=':
return var1.le(var2), np.int
if op == '>':
return var1.gt(var2), np.int
if op == '>=':
return var1.ge(var2), np.int
raise RuntimeError()
if isinstance(other, Number):
inames = tuple(
self.stack.name_generator(based_on='i') for _ in self.shape)
rhs, dtype = _apply_op(
Call(function=Variable(self.name),
parameters=tuple(Variable(iname) for iname in inames)),
other)
subst_name = self.stack.name_generator(based_on='subst')
self.stack.register_substitution(
lp.SubstitutionRule(subst_name, inames, rhs))
return self.copy(name=subst_name, dtype=dtype)
elif isinstance(other, ArraySymbol):
if self.shape == other.shape:
inames = tuple(
self.stack.name_generator(based_on='i')
for _ in self.shape)
rhs, dtype = _apply_op(
Variable(self.name)(*tuple(
Variable(iname) for iname in inames)),
Variable(other.name)(*tuple(
Variable(iname) for iname in inames)))
subst_name = self.stack.name_generator(based_on='subst')
self.stack.register_substitution(
lp.SubstitutionRule(subst_name, inames, rhs))
return self.copy(name=subst_name, dtype=dtype)
else:
left = self
right = other
new_left_shape = self.shape[:]
new_right_shape = other.shape[:]
if len(left.shape) > len(right.shape):
new_right_shape = ((1, ) *
(len(left.shape) - len(right.shape)) +
new_right_shape)
if len(left.shape) < len(right.shape):
new_left_shape = ((1, ) *
(len(right.shape) - len(left.shape)) +
new_left_shape)
# now both have the same shapes
# let compute the shape of these guys
new_shape = [1] * len(new_left_shape)
for i, (left_len, right_len) in enumerate(
zip(new_left_shape, new_right_shape)):
if left_len == right_len:
new_shape[i] = left_len
else:
if left_len != 1 and right_len != 1:
raise ValueError("Cannot be broadcasted.")
else:
new_shape[i] = left_len * right_len
new_shape = tuple(new_shape)
if new_shape != left.shape:
subst_name = self.stack.name_generator(based_on="subst")
inames = tuple(
self.stack.name_generator(based_on="i")
for _ in new_left_shape)
def _empty_if_zero(_idx):
if _idx == (0, ):
return ()
return _idx
indices = _empty_if_zero(
tuple(
Variable(iname) if axis_len != 1 else 0
for iname, axis_len in zip(
inames[-len(left.shape):], left.shape)))
rule = lp.SubstitutionRule(subst_name, inames,
Variable(left.name)(*indices))
self.stack.register_substitution(rule)
new_left = left.copy(name=subst_name,
shape=new_shape,
dim_tags=None,
order=left.order)
else:
new_left = left
if new_shape != right.shape:
subst_name = self.stack.name_generator(based_on="subst")
inames = tuple(
self.stack.name_generator(based_on="i")
for _ in new_right_shape)
def _empty_if_zero(_idx):
if _idx == (0, ):
return ()
return _idx
indices = _empty_if_zero(
tuple(
Variable(iname) if axis_len != 1 else 0
for iname, axis_len in zip(
inames[-len(right.shape):], right.shape)))
rule = lp.SubstitutionRule(subst_name, inames,
Variable(right.name)(*indices))
self.stack.register_substitution(rule)
new_right = right.copy(name=subst_name,
shape=new_shape,
dim_tags=None,
order=right.order)
else:
new_right = right
return new_left._arithmetic_op(new_right, op)
else:
raise NotImplementedError('__mul__ for', type(other))
def __setitem__(self, index, value):
"""
Registers an assignment in order to make the indices represented by
``index`` to ``value``.
:arg index: An instance of :class:`int`, or :class:`slice` or
:class:`numloopy.ArraySymbol`.
:arg value: An instance of :class:`numloopy.ArraySymbol`, with the same
shape as represented by ``index``.
"""
if isinstance(index, (Number, slice, ArraySymbol)):
index = (index, )
assert isinstance(index, tuple)
# current heuristic: assumes that the dereferenced guys are
# always arguments and not temporary variables, maybe we need to fix
# this later?
try:
arg_name = self.stack.substs_to_arrays[self.name]
except KeyError:
inames = tuple(
self.stack.name_generator(based_on="i") for _ in self.shape)
arg_name = self.stack.name_generator(based_on="arr")
insn = lp.Assignment(
assignee=parse('{}[{}]'.format(arg_name, ', '.join(inames))),
expression=parse('{}({})'.format(self.name,
', '.join(inames))))
self.stack.register_implicit_assignment(insn)
space = isl.Space.create_from_names(isl.DEFAULT_CONTEXT, inames)
domain = isl.BasicSet.universe(space)
for iname_name, axis_length in zip(inames, self.shape):
domain &= make_slab(space, iname_name, 0, axis_length)
self.stack.domains.append(domain)
# now handling the second assignment
try:
inames, iname_lens = zip(*tuple(
(self.stack.name_generator(based_on="i"), axis_len)
for idx, axis_len in zip(index, self.shape)
if isinstance(idx, slice) or isinstance(idx, ArraySymbol)))
space = isl.Space.create_from_names(isl.DEFAULT_CONTEXT, inames)
domain = isl.BasicSet.universe(space)
for iname_name, axis_length in zip(inames, iname_lens):
domain &= make_slab(space, iname_name, 0, axis_length)
self.stack.domains.append(domain)
except ValueError:
inames = ()
iname_lens = ()
indices = []
_k = 0
for idx in index:
if isinstance(idx, slice):
indices.append(Variable(inames[_k]))
_k += 1
elif isinstance(idx, ArraySymbol):
indices.append(Variable(idx.name)(Variable(inames[_k])))
_k += 1
else:
indices.append(idx)
assert _k == len(inames)
indices = tuple(indices)
if isinstance(value, ArraySymbol):
insn = lp.Assignment(assignee=Subscript(Variable(arg_name),
indices),
expression='{}({})'.format(
value.name,
', '.join(str(iname)
for iname in inames)))
elif isinstance(value, Number):
insn = lp.Assignment(assignee=Subscript(Variable(arg_name),
indices),
expression=value)
else:
raise TypeError("arrays can be only assigned with number or other "
"arrays")
self.stack.register_implicit_assignment(insn)
if self.name not in self.stack.substs_to_arrays:
subst_name = self.stack.name_generator(based_on="subst")
inames = tuple(
self.stack.name_generator(based_on='i') for _ in self.shape)
rule = lp.SubstitutionRule(
subst_name,
inames,
expression=Subscript(
Variable(arg_name),
tuple(Variable(iname) for iname in inames)))
self.stack.register_substitution(rule)
self.stack.data.append(self.copy(name=arg_name))
self.stack.substs_to_arrays[subst_name] = arg_name
self.name = subst_name