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)