def map_subscript(self, expr, *args):
dc = self.diff_context
if expr.aggregate.name == dc.by_name:
index = expr.index
if not isinstance(expr.index, tuple):
index = (expr.index, )
assert len(self.diff_inames) == len(index)
conds = [
p.Comparison(var(ti), "==", ei)
for ti, ei in zip(self.diff_inames, index)
]
if len(conds) == 1:
and_conds, = conds
elif len(conds) > 1:
and_conds = p.LogicalAnd(tuple(conds))
else:
assert False
return p.If(and_conds, 1, 0)
else:
dvar_dby = dc.get_diff_var(expr.aggregate.name)
if dvar_dby is None:
return 0
idx = expr.index
if not isinstance(idx, tuple):
idx = (idx, )
return type(expr)(var(dvar_dby), expr.index + self.diff_inames)
def nasa7_conditional(t, poly, part_gen):
# FIXME: Should check minTemp, maxTemp
return p.If(
p.Comparison(t, ">", poly.coeffs[0]),
part_gen(poly.coeffs[1:8], t),
part_gen(poly.coeffs[8:15], t),
)
def knl(sym_then, sym_else):
return make_loopy_program(
"{[iel, idof]: 0<=iel<nelements and 0<=idof<nunit_dofs}", [
lp.Assignment(
var("out")[iel, idof],
p.If(var("crit")[iel, idof], sym_then, sym_else))
])
def upper_half_square_root(x):
return p.If(
p.Comparison(
(x**0.5).a.imag,
"<", 0),
1j*(-x)**0.5,
x**0.5)
def knl():
knl = lp.make_kernel(
"{[i]: 0<=i<n}",
[
lp.Assignment(var("out")[i],
p.If(var("crit")[i], sym_then, sym_else))
])
return lp.split_iname(knl, "i", 128, outer_tag="g.0", inner_tag="l.0")
def test_sympy_if_condition():
pytest.importorskip("sympy")
from pymbolic.interop.sympy import PymbolicToSympyMapper, SympyToPymbolicMapper
forward = PymbolicToSympyMapper()
backward = SympyToPymbolicMapper()
# Test round trip to sympy and back
expr = prim.If(prim.Comparison(x_, "<=", y_), 1, 0)
assert backward(forward(expr)) == expr
def map_Piecewise(self, expr): # noqa
# We only handle piecewises with 2 arguments!
assert len(expr.args) == 2
# We only handle if/else cases
assert expr.args[1][1].is_Boolean and bool(expr.args[1][1]) is True
then = self.rec(expr.args[0][0])
else_ = self.rec(expr.args[1][0])
cond = self.rec(expr.args[0][1])
return prim.If(cond, then, else_)
def test_multi_arg_array_call(ctx_factory):
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
import pymbolic.primitives as p
n = 10
acc_i = p.Variable("acc_i")
i = p.Variable("i")
index = p.Variable("index")
a_i = p.Subscript(p.Variable("a"), p.Variable("i"))
argmin_kernel = lp.make_function("{[i]: 0 <= i < n}", [
lp.Assignment(id="init2", assignee=index, expression=0),
lp.Assignment(id="init1", assignee=acc_i, expression="214748367"),
lp.Assignment(id="insn",
assignee=index,
expression=p.If(p.Expression.eq(acc_i, a_i), i, index),
depends_on="update"),
lp.Assignment(id="update",
assignee=acc_i,
expression=p.Variable("min")(acc_i, a_i),
depends_on="init1,init2")
], [
lp.GlobalArg("a"),
lp.GlobalArg(
"acc_i, index", is_input=False, is_output=True, shape=lp.auto), ...
],
name="custom_argmin")
argmin_kernel = lp.fix_parameters(argmin_kernel, n=n)
knl = lp.make_kernel(
"{[i]:0<=i<n}", """
[]: min_val[()], []: min_index[()] = custom_argmin([i]:b[i])
""")
knl = lp.fix_parameters(knl, n=n)
knl = lp.set_options(knl, return_dict=True)
knl = lp.merge([knl, argmin_kernel])
b = np.random.randn(n)
evt, out_dict = knl(queue, b=b)
tol = 1e-15
from numpy.linalg import norm
assert (norm(out_dict["min_val"] - np.min(b)) < tol)
assert (norm(out_dict["min_index"] - np.argmin(b)) < tol)
def expression_conditional(expr, parameters):
return pym.If(*(expression(c, parameters) for c in expr.children))
def _expression_conditional(expr, ctx):
return p.If(*[expression(c, ctx) for c in expr.children])
def map_IfExp(self, expr): # noqa
# (expr test, expr body, expr orelse)
return p.If(self.rec(expr.test), self.rec(expr.body),
self.rec(expr.orelse))
def loopy_inst_conditional(expr, context):
children = [loopy_instructions(o, context) for o in expr.ufl_operands]
return p.If(*children)
