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 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 get_subscript(array_index: int,
offset: ScalarExpression) -> Subscript:
aggregate = var(f"_in{array_index}")
index = [
var(f"_{i}") if i != expr.axis else (var(f"_{i}") - offset)
for i in range(len(expr.shape))
]
return Subscript(aggregate, tuple(index))
def copy_and_rename(expr):
if isinstance(expr, Field):
return expr.copy(child=copy_and_rename(expr.child))
elif isinstance(expr, Subscript):
return Subscript(copy_and_rename(expr.aggregate), expr.index)
elif isinstance(expr, Variable):
return Variable(gen_tmp_name(expr))
elif isinstance(expr, str):
return gen_tmp_name(expr)
def map_subscript(self, expr, expn_state):
from loopy.symbolic import simplify_using_aff
from pymbolic.primitives import Subscript
new_indices = tuple(
simplify_using_aff(self.kernel, self.rec(idx, expn_state))
for idx in expr.index_tuple)
return Subscript(self.rec(expr.aggregate, expn_state), new_indices)
def _math_func(self, op_name, x, res_dtype):
from pymbolic import var
from pymbolic.primitives import Subscript
import pytato as pt
indices = tuple(var(f'_{i}') for i in range(len(x.shape)))
expr = Subscript(var('_in0'), indices)
return pt.IndexLambda(self.ns,
expr,
shape=x.shape,
dtype=res_dtype,
bindings={'_in0': x})
def map_call(self, expr):
if expr.function.name in self.substs_to_args:
def _zero_if_none(_t):
if _t == ():
return (0, )
return _t
return Subscript(
Variable(self.substs_to_args[expr.function.name]),
_zero_if_none(tuple(self.rec(par) for par in expr.parameters)))
return super(SubstToArrayExapander, self).map_call(expr)
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