def add_subscript_if_subscript_nonempty(agg, subscript=()):
from pymbolic.primitives import Subscript, Variable
if len(subscript) == 0:
return Variable(agg)
else:
return Subscript(Variable(agg), tuple(map(Variable,
subscript)))
def test_lpy_array_splitter_f_deep():
from pymbolic.primitives import Subscript, Variable
# create opts
opts = dummy_loopy_opts(depth=8, order='F')
# create array split
asplit = array_splitter(opts)
k = lp.split_iname(_create('F'), 'i', 8)
k = asplit.split_loopy_arrays(k)
# test that it runs
k()
# check dim
a1 = next(x for x in k.args if x.name == 'a1')
assert a1.shape == (8, 10, 2)
# and indexing
assign = next(insn.assignee for insn in k.instructions if insn.id == 'a1')
assert isinstance(assign,
Subscript) and assign.index == (Variable('i_inner'), 0,
Variable('i_outer'))
# now test with evenly sized
a2 = next(x for x in k.args if x.name == 'a2')
assert a2.shape == (8, 16, 2)
assign = next(insn.assignee for insn in k.instructions if insn.id == 'a2')
assert isinstance(assign,
Subscript) and assign.index == (Variable('i_inner'), 0,
Variable('i_outer'))
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 subscript_or_var(agg, subscript=()):
from pymbolic.primitives import Subscript, Variable
if len(subscript) == 0:
return Variable(agg)
else:
return Subscript(Variable(agg), tuple(map(Variable,
subscript)))
def test_forced_iname_deps_and_reduction():
# See https://github.com/inducer/loopy/issues/24
# This is (purposefully) somewhat un-idiomatic, to replicate the conditions
# under which the above bug was found. If assignees were phi[i], then the
# iname propagation heuristic would not assume that dependent instructions
# need to run inside of 'i', and hence the forced_iname_* bits below would not
# be needed.
i1 = lp.CInstruction("i", "doSomethingToGetPhi();", assignees="phi")
from pymbolic.primitives import Subscript, Variable
i2 = lp.Assignment("a",
lp.Reduction("sum", "j",
Subscript(Variable("phi"), Variable("j"))),
forced_iname_deps=frozenset(),
forced_iname_deps_is_final=True)
k = lp.make_kernel(
"{[i,j] : 0<=i,j<n}",
[i1, i2],
[
lp.GlobalArg("a", dtype=np.float32, shape=()),
lp.ValueArg("n", dtype=np.int32),
lp.TemporaryVariable("phi", dtype=np.float32, shape=("n", )),
],
target=lp.CTarget(),
)
k = lp.preprocess_kernel(k)
assert 'i' not in k.insn_inames("insn_0_j_update")
print(k.stringify(with_dependencies=True))
def parse_terminal(self, pstate):
scope = self.tree_walker.scope_stack[-1]
from pymbolic.primitives import Subscript, Call, Variable
from pymbolic.parser import (_identifier, _openpar, _closepar, _float)
next_tag = pstate.next_tag()
if next_tag is _float:
value = pstate.next_str_and_advance().lower()
if "d" in value:
dtype = np.float64
else:
dtype = np.float32
value = value.replace("d", "e")
if value.startswith("."):
prev_value = value
value = "0" + value
print value, prev_value
elif value.startswith("-."):
prev_value = value
value = "-0" + value[1:]
print value, prev_value
return TypedLiteral(value, dtype)
elif next_tag is _identifier:
name = pstate.next_str_and_advance()
if pstate.is_at_end() or pstate.next_tag() is not _openpar:
# not a subscript
scope.use_name(name)
return Variable(name)
left_exp = Variable(name)
pstate.advance()
pstate.expect_not_end()
if scope.is_known(name):
cls = Subscript
else:
cls = Call
if pstate.next_tag is _closepar:
pstate.advance()
left_exp = cls(left_exp, ())
else:
args = self.parse_expression(pstate, self._PREC_FUNC_ARGS)
if not isinstance(args, tuple):
args = (args, )
left_exp = cls(left_exp, args)
pstate.expect(_closepar)
pstate.advance()
return left_exp
else:
return ExpressionParserBase.parse_terminal(self, pstate)
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 knl():
knl = lp.make_kernel("{[i]: 0<=i<n}", [
lp.Assignment(
Variable("out")[i],
Variable(func_name)(Variable("a")[i]))
],
default_offset=lp.auto)
return lp.split_iname(knl,
"i",
128,
outer_tag="g.0",
inner_tag="l.0")
def test_mem_access_counter_nonconsec():
knl = lp.make_kernel(
"[n,m,ell] -> {[i,k,j]: 0<=i<n and 0<=k<m and 0<=j<ell}", [
"""
c[i, j, k] = a[i,j,k]*b[i,j,k]/3.0+a[i,j,k]
e[i, k] = g[i,k]*(2+h[i,k])
"""
],
name="nonconsec",
assumptions="n,m,ell >= 1")
knl = lp.add_and_infer_dtypes(
knl, dict(a=np.float32, b=np.float32, g=np.float64, h=np.float64))
knl = lp.split_iname(knl, "i", 16)
knl = lp.tag_inames(knl, {"i_inner": "l.0", "i_outer": "g.0"})
mem_map = lp.get_mem_access_map(knl, count_redundant_work=True) # noqa
n = 512
m = 256
ell = 128
params = {'n': n, 'm': m, 'ell': ell}
f64nonconsec = mem_map[lp.MemAccess('global',
np.float64,
stride=Variable('m'),
direction='load',
variable='g')].eval_with_dict(params)
f64nonconsec += mem_map[lp.MemAccess('global',
np.float64,
stride=Variable('m'),
direction='load',
variable='h')].eval_with_dict(params)
f32nonconsec = mem_map[lp.MemAccess('global',
np.dtype(np.float32),
stride=Variable('m') * Variable('ell'),
direction='load',
variable='a')].eval_with_dict(params)
f32nonconsec += mem_map[lp.MemAccess('global',
np.dtype(np.float32),
stride=Variable('m') *
Variable('ell'),
direction='load',
variable='b')].eval_with_dict(params)
assert f64nonconsec == 2 * n * m
assert f32nonconsec == 3 * n * m * ell
f64nonconsec = mem_map[lp.MemAccess('global',
np.float64,
stride=Variable('m'),
direction='store',
variable='e')].eval_with_dict(params)
f32nonconsec = mem_map[lp.MemAccess('global',
np.float32,
stride=Variable('m') * Variable('ell'),
direction='store',
variable='c')].eval_with_dict(params)
assert f64nonconsec == n * m
assert f32nonconsec == n * m * ell
def test_subst_into_pwaff():
from pymbolic.primitives import Variable
arg_dict = {"m": 3 * Variable("nx"), "n": 2 * Variable("ny") + 4}
space = isl.Set("[nx, ny, nz] -> { []: }").params().space
poly = isl.PwAff("[m, n] -> { [3 * m + 2 * n] : "
"m > 0 and n > 0; [7* m + 4*n] : m > 0 and n <= 0 }")
from loopy.isl_helpers import subst_into_pwaff
result = subst_into_pwaff(space, poly, arg_dict)
expected = isl.PwAff(
"[nx, ny, nz] -> { [(9nx + 4ny+8)] : nx > 0 and ny > -2;"
" [(21nx + 8ny+16)] : nx > 0 and ny <= -2 }")
assert result == expected
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 get_var_assignees(insn):
try:
return var_assignees_cache[insn]
except KeyError:
result = {Variable(assignee) for assignee in insn.get_assignees()}
var_assignees_cache[insn] = result
return result
def test_buffer_sizes():
wrapper = __test_cases()
for opts in wrapper:
# create a dummy callgen
callgen = CallgenResult(order=opts.order, lang=opts.lang,
dev_mem_type=wrapper.state['dev_mem_type'],
type_map=type_map(opts.lang))
# create a memory manager
mem = get_memory(callgen, host_namer=HostNamer(), device_namer=DeviceNamer())
# test with value arg
a1 = lp.GlobalArg('a1', shape=(arc.problem_size), dtype=np.int32)
assert mem.non_ic_size(a1) == '1'
assert mem.buffer_size(True, a1, num_ics='per_run') == \
'per_run * sizeof(int)'
assert mem.buffer_size(False, a1) == 'problem_size * sizeof(int)'
# test with Variable
from pymbolic.primitives import Variable
a1 = lp.GlobalArg('a1', shape=(Variable(arc.problem_size.name)),
dtype=np.int32)
assert mem.non_ic_size(a1) == '1'
assert mem.buffer_size(True, a1, num_ics='per_run') == \
'per_run * sizeof(int)'
assert mem.buffer_size(False, a1) == 'problem_size * sizeof(int)'
def map_Function(self, expr):
name = self.function_name(expr)
if name in self.functions:
args = tuple(self.rec(arg) for arg in expr.args)
from pymbolic.primitives import Variable, Lookup
if name == "Abs":
call = Lookup(Variable("math"), "fabs")
elif name == "sign":
call = Lookup(Variable("math"), "copysign")
args = (1,)+args
else:
call = Lookup(Variable("math"), name)
return call(*args)
else:
return self.not_supported(expr)
def map_linear_subscript(self, expr, type_context):
from pymbolic.primitives import Variable
if not isinstance(expr.aggregate, Variable):
raise RuntimeError("linear indexing on non-variable: %s" % expr)
if expr.aggregate.name in self.kernel.arg_dict:
arg = self.kernel.arg_dict[expr.aggregate.name]
from loopy.kernel.data import ImageArg
if isinstance(arg, ImageArg):
raise RuntimeError(
"linear indexing is not supported on images: %s" % expr)
else:
# GlobalArg
if arg.offset:
offset = Variable(arg.offset)
else:
offset = 0
return self.make_subscript(arg, var(expr.aggregate.name),
self.rec(offset + expr.index, "i"))
elif expr.aggregate.name in self.kernel.temporary_variables:
raise RuntimeError(
"linear indexing is not supported on temporaries: %s" % expr)
else:
raise RuntimeError("nothing known about variable '%s'" %
expr.aggregate.name)
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 map_linear_subscript(self, expr, enclosing_prec, type_context):
from pymbolic.primitives import Variable
if not isinstance(expr.aggregate, Variable):
raise RuntimeError("linear indexing on non-variable: %s" % expr)
if expr.aggregate.name in self.kernel.arg_dict:
arg = self.kernel.arg_dict[expr.aggregate.name]
from loopy.kernel.data import ImageArg
if isinstance(arg, ImageArg):
raise RuntimeError(
"linear indexing is not supported on images: %s" % expr)
else:
# GlobalArg
if arg.offset:
offset = Variable(arg.offset)
else:
offset = 0
return self.parenthesize_if_needed(
"%s[%s]" % (expr.aggregate.name,
self.rec(offset + expr.index, PREC_NONE, 'i')),
enclosing_prec, PREC_CALL)
elif expr.aggregate.name in self.kernel.temporary_variables:
raise RuntimeError(
"linear indexing is not supported on temporaries: %s" % expr)
else:
raise RuntimeError("nothing known about variable '%s'" %
expr.aggregate.name)
def map_substitution(self, name, tag, arguments, expn_state):
rule = self.rule_mapping_context.old_subst_rules[name]
rec_arguments = self.rec(arguments, expn_state)
if tag is None:
tags = None
else:
tags = (tag,)
new_expn_state = expn_state.copy(
stack=expn_state.stack + ((name, tags),),
arg_context=self.make_new_arg_context(
name, rule.arguments, rec_arguments, expn_state.arg_context))
result = self.rec(rule.expression, new_expn_state)
new_name = self.rule_mapping_context.register_subst_rule(
name, rule.arguments, result)
if tag is None:
sym = Variable(new_name)
else:
sym = TaggedVariable(new_name, tag)
if arguments:
return sym(*rec_arguments)
else:
return sym
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 split_access_axis(expr):
axis_nr, order = array_to_rest[expr.aggregate.name]
idx = expr.index
if not isinstance(idx, tuple):
idx = (idx, )
idx = list(idx)
axis_idx = idx[axis_nr]
if auto_split_inames:
from pymbolic.primitives import Variable
if not isinstance(axis_idx, Variable):
raise RuntimeError(
"found access '%s' in which axis %d is not a "
"single variable--cannot split "
"(Have you tried to do the split yourself, manually, "
"beforehand? If so, you shouldn't.)" % (expr, axis_nr))
split_iname = idx[axis_nr].name
assert split_iname in kernel.all_inames()
try:
outer_iname, inner_iname = split_vars[split_iname]
except KeyError:
outer_iname = var_name_gen(split_iname + "_outer")
inner_iname = var_name_gen(split_iname + "_inner")
split_vars[split_iname] = outer_iname, inner_iname
inner_index = Variable(inner_iname)
outer_index = Variable(outer_iname)
else:
from loopy.symbolic import simplify_using_aff
inner_index = simplify_using_aff(kernel, axis_idx % count)
outer_index = simplify_using_aff(kernel, axis_idx // count)
idx[axis_nr] = inner_index
if order == "F":
idx.insert(axis + 1, outer_index)
elif order == "C":
idx.insert(axis, outer_index)
else:
raise RuntimeError("order '%s' not understood" % order)
return expr.aggregate.index(tuple(idx))
def __init__(self, fft, dk, dx, effective_k):
self.fft = fft
grid_size = fft.grid_shape[0] * fft.grid_shape[1] * fft.grid_shape[2]
queue = self.fft.sub_k["momenta_x"].queue
sub_k = list(x.get().astype("int") for x in self.fft.sub_k.values())
k_names = ("k_x", "k_y", "k_z")
self.momenta = {}
self.momenta = {}
for mu, (name, kk) in enumerate(zip(k_names, sub_k)):
kk_mu = effective_k(dk[mu] * kk.astype(fft.rdtype), dx[mu])
self.momenta[name] = cla.to_device(queue, kk_mu)
args = [
lp.GlobalArg("fk", fft.cdtype, shape="(Nx, Ny, Nz)"),
lp.GlobalArg("k_x", fft.rdtype, shape=("Nx", )),
lp.GlobalArg("k_y", fft.rdtype, shape=("Ny", )),
lp.GlobalArg("k_z", fft.rdtype, shape=("Nz", )),
lp.ValueArg("m_squared", fft.rdtype),
]
from pystella.field import Field
from pymbolic.primitives import Variable, If, Comparison
fk = Field("fk")
indices = fk.indices
rho_tmp = Variable("rho_tmp")
tmp_insns = [(rho_tmp, Field("rhok") * (1 / grid_size))]
mom_vars = tuple(Variable(name) for name in k_names)
minus_k_squared = sum(kk_i[x_i]
for kk_i, x_i in zip(mom_vars, indices))
sol = rho_tmp / (minus_k_squared - Variable("m_squared"))
solution = {
Field("fk"): If(Comparison(minus_k_squared, "<", 0), sol, 0)
}
from pystella.elementwise import ElementWiseMap
options = lp.Options(return_dict=True)
self.knl = ElementWiseMap(solution,
args=args,
halo_shape=0,
options=options,
tmp_instructions=tmp_insns,
lsize=(16, 2, 1))
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 get_code_transformer(self):
from sumpy.codegen import VectorComponentRewriter
vcr = VectorComponentRewriter([self.dir_vec_name])
from pymbolic.primitives import Variable
via = _VectorIndexAdder(self.dir_vec_name, (Variable("itgt"), ))
def transform(expr):
return via(vcr(expr))
return transform
def test_sympy_interop(proc_shape):
if proc_shape != (1, 1, 1):
pytest.skip("test field only on one rank")
from pystella.field.sympy import pymbolic_to_sympy, sympy_to_pymbolic
import sympy as sym
f = ps.Field("f", offset="h")
g = ps.Field("g", offset="h")
expr = f[0]**2 * g + 2 * g[1] * f
sympy_expr = pymbolic_to_sympy(expr)
new_expr = sympy_to_pymbolic(sympy_expr)
sympy_expr_2 = pymbolic_to_sympy(new_expr)
assert sym.simplify(sympy_expr - sympy_expr_2) == 0, \
"sympy <-> pymbolic conversion not invertible"
expr = f + shift_fields(f, (1, 2, 3))
sympy_expr = pymbolic_to_sympy(expr)
new_expr = sympy_to_pymbolic(sympy_expr)
sympy_expr_2 = pymbolic_to_sympy(new_expr)
assert sym.simplify(sympy_expr - sympy_expr_2) == 0, \
"sympy <-> pymbolic conversion not invertible with shifted indices"
# from pymbolic.functions import fabs, exp, exmp1
fabs = parse("math.fabs")
exp = parse("math.exp")
expm1 = parse("math.expm1")
x = sym.Symbol("x")
expr = sym.Abs(x)
assert sympy_to_pymbolic(expr) == fabs(var("x"))
expr = sym.exp(x)
assert sympy_to_pymbolic(expr) == exp(var("x"))
expr = sym.Function("expm1")(x) # pylint: disable=E1102
assert sympy_to_pymbolic(expr) == expm1(var("x"))
expr = sym.Function("aaa")(x) # pylint: disable=E1102
from pymbolic.primitives import Call, Variable
assert sympy_to_pymbolic(expr) == Call(Variable("aaa"), (Variable("x"), ))
def map_int_g(self, expr, name_hint=None):
try:
return self.expr_to_var[expr]
except KeyError:
from pytential.utils import sort_arrays_together
source_kernels, densities = \
sort_arrays_together(expr.source_kernels, expr.densities, key=str)
# make sure operator assignments stand alone and don't get muddled
# up in vector arithmetic
density_vars = [self.assign_to_new_var(self.rec(density)) for
density in densities]
group = self.group_to_operators[self.op_group_features(expr)]
names = [self.get_var_name() for op in group]
sorted_ops = sorted(group, key=lambda op: repr(op.target_kernel))
target_kernels = [op.target_kernel for op in sorted_ops]
target_kernel_to_index = \
{kernel: i for i, kernel in enumerate(target_kernels)}
for op in group:
assert op.qbx_forced_limit in [-2, -1, None, 1, 2]
kernel_arguments = {
arg_name: self.rec(arg_val)
for arg_name, arg_val in expr.kernel_arguments.items()}
outputs = [
PotentialOutput(
name=name,
target_kernel_index=target_kernel_to_index[op.target_kernel],
target_name=op.target,
qbx_forced_limit=op.qbx_forced_limit,
)
for name, op in zip(names, group)
]
self.code.append(
ComputePotentialInstruction(
outputs=outputs,
target_kernels=tuple(target_kernels),
kernel_arguments=kernel_arguments,
source_kernels=source_kernels,
densities=density_vars,
source=expr.source,
priority=max(getattr(op, "priority", 0) for op in group),
dep_mapper_factory=self.dep_mapper_factory))
from pymbolic.primitives import Variable
for name, group_expr in zip(names, group):
self.expr_to_var[group_expr] = Variable(name)
return self.expr_to_var[expr]
def map_variable(self, expr):
name, tag = parse_tagged_name(expr)
new_name = self.renames.get(name)
if new_name is None:
return IdentityMapper.map_variable(self, expr)
if tag is None:
return Variable(new_name)
else:
return TaggedVariable(new_name, tag)
def test_subst_into_pwqpolynomial():
from pymbolic.primitives import Variable
arg_dict = {
"m": 3 * Variable("nx"),
"n": 3 * Variable("ny"),
"nx": Variable("nx"),
"ny": Variable("ny"),
"nz": Variable("nz")
}
space = isl.Set("[nx, ny, nz] -> { []: }").space
poly = isl.PwQPolynomial("[m, n] -> { (256 * m + 256 * m * n) : "
"m > 0 and n > 0; 256 * m : m > 0 and n <= 0 }")
from loopy.isl_helpers import subst_into_pwqpolynomial
result = subst_into_pwqpolynomial(space, poly, arg_dict)
expected_pwqpoly = isl.PwQPolynomial(
"[nx, ny, nz] -> {"
"(768 * nx + 2304 * nx * ny) : nx > 0 and ny > 0;"
"768 * nx : nx > 0 and ny <= 0 }")
assert (result - expected_pwqpoly).is_zero()
def sexpr_to_pymbolic(expr):
from pymbolic.primitives import Variable, Sum, Product, Power
stack = []
def error(token):
raise ParserError("unexpected token at position %d (got '%r')" %
(token.pos, token))
for token in tokenize(expr):
if token.token_type in (TokenType.LPAREN, TokenType.IDENT):
stack.append(token)
elif token.token_type in (TokenType.FLOAT, TokenType.INT,
TokenType.STRING):
stack.append(token.data)
elif token.token_type == TokenType.RPAREN:
args = []
while stack and not (isinstance(stack[-1], Token)
and stack[-1].token_type == TokenType.LPAREN):
args.append(stack.pop())
if not stack:
error(token)
lparen = stack.pop()
args = tuple(reversed(args))
if not args:
error(token)
sym = args[0]
if not isinstance(sym, Token):
error(lparen)
assert sym.token_type == TokenType.IDENT
if sym.data == "Sum":
val = Sum(args[1:])
elif sym.data == "Product":
val = Product(args[1:])
elif sym.data == "Power":
val = Power(*args[1:])
elif sym.data == "Var":
val = Variable(*args[1:])
else:
error(sym)
stack.append(val)
elif token.token_type == TokenType.END:
if len(stack) != 1:
error(token)
if isinstance(stack[0], Token):
error(token)
return stack[0]
def get_dependencies(self):
# arg is include_subscripts
dep_mapper = self.dep_mapper_factory()
from operator import or_
deps = reduce(or_, (dep_mapper(expr) for expr in self.exprs))
from pymbolic.primitives import Variable
deps -= {Variable(name) for name in self.names}
return deps
def aggregate_two_assignments(ass_1, ass_2):
names = ass_1.names + ass_2.names
from pymbolic.primitives import Variable
deps = (ass_1.get_dependencies() | ass_2.get_dependencies()) \
- {Variable(name) for name in names}
return Assign(
names=names, exprs=ass_1.exprs + ass_2.exprs,
_dependencies=deps,
priority=max(ass_1.priority, ass_2.priority))
def __call__(self, operand, *args, **kwargs):
# If the call is handed an object array full of operands,
# return an object array of the operator applied to each of the
# operands.
from pytools.obj_array import is_obj_array, with_object_array_or_scalar
if is_obj_array(operand):
def make_op(operand_i):
return self(operand_i, *args, **kwargs)
return with_object_array_or_scalar(make_op, operand)
else:
return var.__call__(self, operand, *args, **kwargs)
def __init__(self, name, tag):
Variable.__init__(self, name)
self.tag = tag