def map_basic_index(self, expr: BasicIndex) -> IndexLambda:
vng = UniqueNameGenerator()
indices = []
in_ary = vng("in")
bindings = {in_ary: self.rec(expr.array)}
islice_idx = 0
for idx, axis_len in zip(expr.indices, expr.array.shape):
if isinstance(idx, INT_CLASSES):
if isinstance(axis_len, INT_CLASSES):
indices.append(idx % axis_len)
else:
bnd_name = vng("in")
bindings[bnd_name] = axis_len
indices.append(idx % prim.Variable(bnd_name))
elif isinstance(idx, NormalizedSlice):
indices.append(idx.start
+ idx.step * prim.Variable(f"_{islice_idx}"))
islice_idx += 1
else:
raise NotImplementedError
return IndexLambda(expr=prim.Subscript(prim.Variable(in_ary),
tuple(indices)),
bindings=bindings,
shape=expr.shape,
dtype=expr.dtype,
axes=expr.axes,
tags=expr.tags,
)
def _expression_mathfunction(expr, ctx):
if expr.name.startswith('cyl_bessel_'):
# Bessel functions
if is_complex(ctx.scalar_type):
raise NotImplementedError("Bessel functions for complex numbers: "
"missing implementation")
nu, arg = expr.children
nu_ = expression(nu, ctx)
arg_ = expression(arg, ctx)
# Modified Bessel functions (C++ only)
#
# These mappings work for FEniCS only, and fail with Firedrake
# since no Boost available.
if expr.name in {'cyl_bessel_i', 'cyl_bessel_k'}:
name = 'boost::math::' + expr.name
return p.Variable(name)(nu_, arg_)
else:
# cyl_bessel_{jy} -> {jy}
name = expr.name[-1:]
if nu == gem.Zero():
return p.Variable(f"{name}0")(arg_)
elif nu == gem.one:
return p.Variable(f"{name}1")(arg_)
else:
return p.Variable(f"{name}n")(nu_, arg_)
else:
if expr.name == "ln":
name = "log"
else:
name = expr.name
# Not all mathfunctions apply to complex numbers, but this
# will be picked up in loopy. This way we allow erf(real(...))
# in complex mode (say).
return p.Variable(name)(*(expression(c, ctx) for c in expr.children))
def _test_to_pymbolic(mapper, sym, use_symengine):
x, y = sym.symbols("x,y")
assert mapper(sym.Rational(3, 4)) == prim.Quotient(3, 4)
assert mapper(sym.Integer(6)) == 6
if not use_symengine:
assert mapper(sym.Subs(x**2, (x,), (y,))) == \
prim.Substitution(x_**2, ("x",), (y_,))
deriv = sym.Derivative(x**2, x)
assert mapper(deriv) == prim.Derivative(x_**2, ("x", ))
else:
assert mapper(sym.Subs(x**2, (x,), (y,))) == \
y_**2
deriv = sym.Derivative(x**2, x)
assert mapper(deriv) == 2 * x_
# functions
assert mapper(sym.Function("f")(x)) == prim.Variable("f")(x_)
assert mapper(sym.exp(x)) == prim.Variable("exp")(x_)
# indexed accesses
if not use_symengine:
i, j = sym.symbols("i,j")
assert mapper(sym.Indexed(x, i, j)) == x_[i_, j_]
# constants
import math
# FIXME: Why isn't this exact?
assert abs(mapper(sym.pi) - math.pi) < 1e-14
assert abs(mapper(sym.E) - math.e) < 1e-14
assert mapper(sym.I) == 1j
def facet_integral_predicates(self, mesh, integral_type, kinfo):
self.bag.needs_cell_facets = True
# Number of recerence cell facets
if mesh.cell_set._extruded:
self.num_facets = mesh._base_mesh.ufl_cell().num_facets()
else:
self.num_facets = mesh.ufl_cell().num_facets()
# Index for loop over cell faces of reference cell
fidx = self.bag.index_creator((self.num_facets,))
# Cell is interior or exterior
select = 1 if integral_type.startswith("interior_facet") else 0
i = self.bag.index_creator((1,))
predicates = [pym.Comparison(pym.Subscript(pym.Variable(self.cell_facets_arg), (fidx[0], 0)), "==", select)]
# TODO subdomain boundary integrals, this does the wrong thing for integrals like f*ds + g*ds(1)
# "otherwise" is treated incorrectly as "everywhere"
# However, this replicates an existing slate bug.
if kinfo.subdomain_id != "otherwise":
predicates.append(pym.Comparison(pym.Subscript(pym.Variable(self.cell_facets_arg), (fidx[0], 1)), "==", kinfo.subdomain_id))
# Additional facet array argument to be fed into tsfc loopy kernel
subscript = pym.Subscript(pym.Variable(self.local_facet_array_arg),
(pym.Sum((i[0], fidx[0]))))
facet_arg = SubArrayRef(i, subscript)
return predicates, fidx, facet_arg
def slate_call(self, kernel):
# Slate kernel call
call = pym.Call(pym.Variable(kernel.name), tuple())
output_var = pym.Variable(kernel.args[0].name)
slate_kernel_call_output = self.generate_lhs(self.expression, output_var)
insn = loopy.CallInstruction((slate_kernel_call_output,), call, id="slate_kernel_call")
return [insn]
def test_parser():
from pymbolic import parse
parse("(2*a[1]*b[1]+2*a[0]*b[0])*(hankel_1(-1,sqrt(a[1]**2+a[0]**2)*k) "
"-hankel_1(1,sqrt(a[1]**2+a[0]**2)*k))*k /(4*sqrt(a[1]**2+a[0]**2)) "
"+hankel_1(0,sqrt(a[1]**2+a[0]**2)*k)")
print repr(parse("d4knl0"))
print repr(parse("0."))
print repr(parse("0.e1"))
print repr(parse("0.e1"))
print repr(parse("a >= 1"))
print repr(parse("a <= 1"))
print repr(parse(":"))
print repr(parse("1:"))
print repr(parse(":2"))
print repr(parse("1:2"))
print repr(parse("::"))
print repr(parse("1::"))
print repr(parse(":1:"))
print repr(parse("::1"))
print repr(parse("3::1"))
print repr(parse(":5:1"))
print repr(parse("3:5:1"))
print parse("3::1")
assert parse("e1") == prim.Variable("e1")
assert parse("d1") == prim.Variable("d1")
def map_power(self, expr, type_context):
tgt_dtype = self.infer_type(expr)
base_dtype = self.infer_type(expr.base)
exponent_dtype = self.infer_type(expr.exponent)
if not self.allow_complex or (not tgt_dtype.is_complex()):
return super().map_power(expr, type_context)
if expr.exponent in [2, 3, 4]:
value = expr.base
for _i in range(expr.exponent - 1):
value = value * expr.base
return self.rec(value, type_context)
else:
b_complex = base_dtype.is_complex()
e_complex = exponent_dtype.is_complex()
if b_complex and not e_complex:
return p.Variable(
"%s_powr" % self.complex_type_name(tgt_dtype))(self.rec(
expr.base, type_context,
tgt_dtype), self.rec(expr.exponent, type_context))
else:
return p.Variable(
"%s_pow" % self.complex_type_name(tgt_dtype))(
self.rec(expr.base, type_context, tgt_dtype),
self.rec(expr.exponent, type_context, tgt_dtype))
def initialise_terminals(self, var2terminal, coefficients):
""" Initilisation of the variables in which coefficients
and the Tensors coming from TSFC are saved.
:arg var2terminal: dictionary that maps Slate Tensors to gem Variables
"""
tensor2temp = OrderedDict()
inits = []
for gem_tensor, slate_tensor in var2terminal.items():
assert slate_tensor.terminal, "Only terminal tensors need to be initialised in Slate kernels."
(_, dtype), = assign_dtypes([gem_tensor],
self.tsfc_parameters["scalar_type"])
loopy_tensor = loopy.TemporaryVariable(
gem_tensor.name,
dtype=dtype,
shape=gem_tensor.shape,
address_space=loopy.AddressSpace.LOCAL)
tensor2temp[slate_tensor] = loopy_tensor
if not slate_tensor.assembled:
indices = self.bag.index_creator(self.shape(slate_tensor))
inames = {var.name for var in indices}
var = pym.Subscript(pym.Variable(loopy_tensor.name), indices)
inits.append(
loopy.Assignment(var,
"0.",
id="init%d" % len(inits),
within_inames=frozenset(inames)))
else:
f = slate_tensor.form if isinstance(
slate_tensor.form, tuple) else (slate_tensor.form, )
coeff = tuple(coefficients[c] for c in f)
offset = 0
ismixed = tuple(
(type(c.ufl_element()) == MixedElement) for c in f)
names = []
for (im, c) in zip(ismixed, coeff):
names += [name
for (name, ext) in c.values()] if im else [c[0]]
# Mixed coefficients come as seperate parameter (one per space)
for i, shp in enumerate(*slate_tensor.shapes.values()):
indices = self.bag.index_creator((shp, ))
inames = {var.name for var in indices}
offset_index = (pym.Sum((offset, indices[0])), )
name = names[i] if ismixed else names
var = pym.Subscript(pym.Variable(loopy_tensor.name),
offset_index)
c = pym.Subscript(pym.Variable(name), indices)
inits.append(
loopy.Assignment(var,
c,
id="init%d" % len(inits),
within_inames=frozenset(inames)))
offset += shp
return inits, tensor2temp
def map_non_contiguous_advanced_index(self,
expr: AdvancedIndexInNoncontiguousAxes
) -> IndexLambda:
from pytato.utils import (get_shape_after_broadcasting,
get_indexing_expression)
i_adv_indices = tuple(i
for i, idx_expr in enumerate(expr.indices)
if isinstance(idx_expr, (Array, INT_CLASSES)))
adv_idx_shape = get_shape_after_broadcasting([expr.indices[i_idx]
for i_idx in i_adv_indices])
vng = UniqueNameGenerator()
indices = []
in_ary = vng("in")
bindings = {in_ary: self.rec(expr.array)}
islice_idx = len(adv_idx_shape)
for idx, axis_len in zip(expr.indices, expr.array.shape):
if isinstance(idx, INT_CLASSES):
if isinstance(axis_len, INT_CLASSES):
indices.append(idx % axis_len)
else:
bnd_name = vng("in")
bindings[bnd_name] = self.rec(axis_len)
indices.append(idx % prim.Variable(bnd_name))
elif isinstance(idx, NormalizedSlice):
indices.append(idx.start
+ idx.step * prim.Variable(f"_{islice_idx}"))
islice_idx += 1
elif isinstance(idx, Array):
if isinstance(axis_len, INT_CLASSES):
bnd_name = vng("in")
bindings[bnd_name] = self.rec(idx)
indirect_idx_expr = prim.Subscript(prim.Variable(bnd_name),
get_indexing_expression(
idx.shape,
adv_idx_shape))
if not idx.tags_of_type(AssumeNonNegative):
indirect_idx_expr = indirect_idx_expr % axis_len
indices.append(indirect_idx_expr)
else:
raise NotImplementedError("Advanced indexing over"
" parametric axis lengths.")
else:
raise NotImplementedError(f"Indices of type {type(idx)}.")
return IndexLambda(expr=prim.Subscript(prim.Variable(in_ary),
tuple(indices)),
bindings=bindings,
shape=expr.shape,
dtype=expr.dtype,
axes=expr.axes,
tags=expr.tags,
)
def statement_evaluate(leaf, ctx):
expr = leaf.expression
if isinstance(expr, gem.ListTensor):
ops = []
var, index = ctx.pymbolic_variable_and_destruct(expr)
for multiindex, value in numpy.ndenumerate(expr.array):
ops.append(
lp.Assignment(p.Subscript(var, index + multiindex),
expression(value, ctx),
within_inames=ctx.active_inames()))
return ops
elif isinstance(expr, gem.Constant):
return []
elif isinstance(expr, gem.ComponentTensor):
idx = ctx.gem_to_pym_multiindex(expr.multiindex)
var, sub_idx = ctx.pymbolic_variable_and_destruct(expr)
lhs = p.Subscript(var, idx + sub_idx)
with active_indices(dict(zip(expr.multiindex, idx)),
ctx) as ctx_active:
return [
lp.Assignment(lhs,
expression(expr.children[0], ctx_active),
within_inames=ctx_active.active_inames())
]
elif isinstance(expr, gem.Inverse):
idx = ctx.pymbolic_multiindex(expr.shape)
var = ctx.pymbolic_variable(expr)
lhs = (SubArrayRef(idx, p.Subscript(var, idx)), )
idx_reads = ctx.pymbolic_multiindex(expr.children[0].shape)
var_reads = ctx.pymbolic_variable(expr.children[0])
reads = (SubArrayRef(idx_reads, p.Subscript(var_reads, idx_reads)), )
rhs = p.Call(p.Variable("inverse"), reads)
return [
lp.CallInstruction(lhs, rhs, within_inames=ctx.active_inames())
]
elif isinstance(expr, gem.Solve):
idx = ctx.pymbolic_multiindex(expr.shape)
var = ctx.pymbolic_variable(expr)
lhs = (SubArrayRef(idx, p.Subscript(var, idx)), )
reads = []
for child in expr.children:
idx_reads = ctx.pymbolic_multiindex(child.shape)
var_reads = ctx.pymbolic_variable(child)
reads.append(
SubArrayRef(idx_reads, p.Subscript(var_reads, idx_reads)))
rhs = p.Call(p.Variable("solve"), tuple(reads))
return [
lp.CallInstruction(lhs, rhs, within_inames=ctx.active_inames())
]
else:
return [
lp.Assignment(ctx.pymbolic_variable(expr),
expression(expr, ctx, top=True),
within_inames=ctx.active_inames())
]
def map_constant(self, expr, type_context):
from loopy.symbolic import Literal
if isinstance(expr, (complex, np.complexfloating)):
real = self.rec(expr.real)
imag = self.rec(expr.imag)
iota = p.Variable("I" if "I" not in self.kernel.all_variable_names(
) else "_Complex_I")
return real + imag * iota
elif np.isnan(expr):
return p.Variable("NAN")
elif np.isneginf(expr):
return -p.Variable("INFINITY")
elif np.isinf(expr):
return p.Variable("INFINITY")
elif isinstance(expr, np.generic):
# Explicitly typed: Generated code must reflect type exactly.
# FIXME: This assumes a 32-bit architecture.
if isinstance(expr, np.float32):
return Literal(repr(expr) + "f")
elif isinstance(expr, np.float64):
return Literal(repr(expr))
# Disabled for now, possibly should be a subtarget.
# elif isinstance(expr, np.float128):
# return Literal(repr(expr)+"l")
elif isinstance(expr, np.integer):
suffix = ""
iinfo = np.iinfo(expr)
if iinfo.min == 0:
suffix += "u"
if iinfo.max > (2**31 - 1):
suffix += "l"
return Literal(repr(expr) + suffix)
elif isinstance(expr, np.bool_):
return Literal("true") if expr else Literal("false")
else:
raise LoopyError("do not know how to generate code for "
"constant of numpy type '%s'" %
type(expr).__name__)
elif np.isfinite(expr):
if type_context == "f":
return Literal(repr(np.float32(expr)) + "f")
elif type_context == "d":
return Literal(repr(float(expr)))
elif type_context in ["i", "b"]:
return int(expr)
else:
if is_integer(expr):
return int(expr)
raise RuntimeError("don't know how to generate code "
"for constant '%s'" % expr)
else:
raise LoopyError("don't know how to generate code "
"for constant '%s'" % expr)
def generate_tsfc_calls(self, terminal, loopy_tensor):
"""A setup method to initialize all the local assembly
kernels generated by TSFC. This function also collects any
information regarding orientations and extra include directories.
"""
cxt_kernels = self.tsfc_cxt_kernels(terminal)
for cxt_kernel in cxt_kernels:
for tsfc_kernel in cxt_kernel.tsfc_kernels:
integral_type = cxt_kernel.original_integral_type
slate_tensor = cxt_kernel.tensor
mesh = slate_tensor.ufl_domain()
kinfo = tsfc_kernel.kinfo
reads = []
inames_dep = []
if integral_type not in self.supported_integral_types:
raise ValueError("Integral type '%s' not recognized" %
integral_type)
# Prepare lhs and args for call to tsfc kernel
output_var = pym.Variable(loopy_tensor.name)
reads.append(output_var)
output = self.generate_lhs(slate_tensor, output_var)
kernel_data = self.collect_tsfc_kernel_data(
mesh, cxt_kernel.coefficients, self.bag.coefficients,
kinfo)
reads.extend(self.loopify_tsfc_kernel_data(kernel_data))
# Generate predicates for different integral types
if self.is_integral_type(integral_type, "cell_integral"):
predicates = None
if kinfo.subdomain_id != "otherwise":
raise NotImplementedError(
"No subdomain markers for cells yet")
elif self.is_integral_type(integral_type, "facet_integral"):
predicates, fidx, facet_arg = self.facet_integral_predicates(
mesh, integral_type, kinfo)
reads.append(facet_arg)
inames_dep.append(fidx[0].name)
elif self.is_integral_type(integral_type, "layer_integral"):
predicates = self.layer_integral_predicates(
slate_tensor, integral_type)
else:
raise ValueError(
"Unhandled integral type {}".format(integral_type))
# TSFC kernel call
key = self.bag.call_name_generator(integral_type)
call = pym.Call(pym.Variable(kinfo.kernel.name), tuple(reads))
insn = loopy.CallInstruction(
(output, ),
call,
within_inames=frozenset(inames_dep),
predicates=predicates,
id=key)
yield insn, kinfo.kernel.code
def parse_terminal(self, pstate):
scope = self.tree_walker.scope_stack[-1]
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 p.Variable(name)
left_exp = p.Variable(name)
pstate.advance()
pstate.expect_not_end()
if scope.is_known(name):
cls = p.Subscript
else:
cls = p.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 _(poly: ct.NasaPoly2, arg_name):
log = p.Variable("log")
def gen(c, t):
assert len(c) == 7
return (c[0] * log(t) + c[1] * t + c[2] / 2 * t**2 + c[3] / 3 * t**3 +
c[4] / 4 * t**4 + c[6])
return nasa7_conditional(p.Variable(arg_name), poly, gen)
def to_loopy_expression(
self, indices: SymbolicIndex,
expr_context: LoopyExpressionContext) -> ScalarExpression:
assert len(indices) == self.num_indices
expr_context.update_depends_on(self.depends_on)
if indices == ():
return prim.Variable(self.name)
else:
return prim.Variable(self.name)[indices]
def wrap_in_typecast(self, actual_type, needed_dtype, s):
if (actual_type.is_complex() and needed_dtype.is_complex()
and actual_type != needed_dtype):
return p.Variable("%s_cast" %
self.complex_type_name(needed_dtype))(s)
elif not actual_type.is_complex() and needed_dtype.is_complex():
return p.Variable("%s_fromreal" %
self.complex_type_name(needed_dtype))(s)
else:
return super().wrap_in_typecast_lazy(actual_type, needed_dtype, s)
def layer_integral_predicates(self, tensor, integral_type):
self.bag.needs_mesh_layers = True
layer = pym.Variable(self.layer_arg)
# TODO: Variable layers
nlayer = pym.Variable(self.layer_count)
which = {"interior_facet_horiz_top": pym.Comparison(layer, "<", nlayer),
"interior_facet_horiz_bottom": pym.Comparison(layer, ">", 0),
"exterior_facet_top": pym.Comparison(layer, "==", nlayer),
"exterior_facet_bottom": pym.Comparison(layer, "==", 0)}[integral_type]
return [which]
def test_flop_counter():
x = prim.Variable("x")
y = prim.Variable("y")
z = prim.Variable("z")
subexpr = prim.CommonSubexpression(3 * (x**2 + y + z))
expr = 3 * subexpr + subexpr
from pymbolic.mapper.flop_counter import FlopCounter, CSEAwareFlopCounter
assert FlopCounter()(expr) == 4 * 2 + 2
assert CSEAwareFlopCounter()(expr) == 4 + 2
def test_parser():
from pymbolic import parse
parse("(2*a[1]*b[1]+2*a[0]*b[0])*(hankel_1(-1,sqrt(a[1]**2+a[0]**2)*k) "
"-hankel_1(1,sqrt(a[1]**2+a[0]**2)*k))*k /(4*sqrt(a[1]**2+a[0]**2)) "
"+hankel_1(0,sqrt(a[1]**2+a[0]**2)*k)")
print(repr(parse("d4knl0")))
print(repr(parse("0.")))
print(repr(parse("0.e1")))
assert parse("0.e1") == 0
assert parse("1e-12") == 1e-12
print(repr(parse("a >= 1")))
print(repr(parse("a <= 1")))
print(repr(parse(":")))
print(repr(parse("1:")))
print(repr(parse(":2")))
print(repr(parse("1:2")))
print(repr(parse("::")))
print(repr(parse("1::")))
print(repr(parse(":1:")))
print(repr(parse("::1")))
print(repr(parse("3::1")))
print(repr(parse(":5:1")))
print(repr(parse("3:5:1")))
print(repr(parse("g[i,k]+2.0*h[i,k]")))
print(repr(parse("g[i,k]+(+2.0)*h[i,k]")))
print(repr(parse("a - b - c")))
print(repr(parse("-a - -b - -c")))
print(repr(parse("- - - a - - - - b - - - - - c")))
print(repr(parse("~(a ^ b)")))
print(repr(parse("(a | b) | ~(~a & ~b)")))
print(repr(parse("3 << 1")))
print(repr(parse("1 >> 3")))
print(parse("3::1"))
assert parse("e1") == prim.Variable("e1")
assert parse("d1") == prim.Variable("d1")
from pymbolic import variables
f, x, y, z = variables("f x y z")
assert parse("f((x,y),z)") == f((x, y), z)
assert parse("f((x,),z)") == f((x, ), z)
assert parse("f(x,(y,z),z)") == f(x, (y, z), z)
assert parse("f(x,(y,z),z, name=15)") == f(x, (y, z), z, name=15)
assert parse("f(x,(y,z),z, name=15, name2=17)") == f(x, (y, z),
z,
name=15,
name2=17)
def slate_call(self, kernel, temporaries):
# Slate kernel call
reads = []
for t in temporaries:
shape = t.shape
name = t.name
idx = self.bag.index_creator(shape)
reads.append(SubArrayRef(idx, pym.Subscript(pym.Variable(name), idx)))
call = pym.Call(pym.Variable(kernel.name), tuple(reads))
output_var = pym.Variable(kernel.args[0].name)
slate_kernel_call_output = self.generate_lhs(self.expression, output_var)
insn = loopy.CallInstruction((slate_kernel_call_output,), call, id="slate_kernel_call")
return insn
def test_parser():
from pymbolic import var
from dagrt.expression import parse
parse("(2*a[1]*b[1]+2*a[0]*b[0])*(hankel_1(-1,sqrt(a[1]**2+a[0]**2)*k) "
"-hankel_1(1,sqrt(a[1]**2+a[0]**2)*k))*k /(4*sqrt(a[1]**2+a[0]**2)) "
"+hankel_1(0,sqrt(a[1]**2+a[0]**2)*k)")
print(repr(parse("d4knl0")))
print(repr(parse("0.")))
print(repr(parse("0.e1")))
assert parse("0.e1") == 0
assert parse("1e-12") == 1e-12
print(repr(parse("a >= 1")))
print(repr(parse("a <= 1")))
print(repr(parse("g[i,k]+2.0*h[i,k]")))
print(repr(parse("g[i,k]+(+2.0)*h[i,k]")))
print(repr(parse("a - b - c")))
print(repr(parse("-a - -b - -c")))
print(repr(parse("- - - a - - - - b - - - - - c")))
print(repr(parse("~(a ^ b)")))
print(repr(parse("(a | b) | ~(~a & ~b)")))
print(repr(parse("3 << 1")))
print(repr(parse("1 >> 3")))
print(parse("3::1"))
import pymbolic.primitives as prim
assert parse("e1") == prim.Variable("e1")
assert parse("d1") == prim.Variable("d1")
from pymbolic import variables
f, x, y, z = variables("f x y z")
assert parse("f((x,y),z)") == f((x, y), z)
assert parse("f((x,),z)") == f((x, ), z)
assert parse("f(x,(y,z),z)") == f(x, (y, z), z)
assert parse("f(x,(y,z),z, name=15)") == f(x, (y, z), z, name=15)
assert parse("f(x,(y,z),z, name=15, name2=17)") == f(x, (y, z),
z,
name=15,
name2=17)
assert parse("<func>yoink") == var("<func>yoink")
assert parse("-< func > yoink") == -var("<func>yoink")
print(repr(parse("<func>yoink < <p>x")))
print(repr(parse("<func>yoink < - <p>x")))
def loopy_inst_mathfunc(expr, context):
children = [loopy_instructions(o, context) for o in expr.ufl_operands]
if expr._name == "ln":
name = "log"
else:
name = expr._name
return p.Variable(name)(*children)
def binary_tree_mul(start, end):
if start + 1 == end:
return c_applied[start]
mid = (start + end) // 2
lsum = binary_tree_mul(start, mid)
rsum = binary_tree_mul(mid, end)
return p.Variable("%s_mul" % tgt_name)(lsum, rsum)
def _indices_for_reshape(self, expr: Reshape) -> SymbolicIndex:
if expr.array.shape == ():
# RHS must be a scalar i.e. RHS' indices are empty
assert expr.size == 1
return ()
newstrides = [1] # reshaped array strides
for new_axis_len in reversed(expr.shape[1:]):
assert isinstance(new_axis_len, INT_CLASSES)
newstrides.insert(0, newstrides[0]*new_axis_len)
flattened_idx = sum(prim.Variable(f"_{i}")*stride
for i, stride in enumerate(newstrides))
oldstrides: List[IntegralT] = [1] # input array strides
for axis_len in reversed(expr.array.shape[1:]):
assert isinstance(axis_len, INT_CLASSES)
oldstrides.insert(0, oldstrides[0]*axis_len)
assert isinstance(expr.array.shape[-1], INT_CLASSES)
oldsizetills = [expr.array.shape[-1]] # input array size till for axes idx
for old_axis_len in reversed(expr.array.shape[:-1]):
assert isinstance(old_axis_len, INT_CLASSES)
oldsizetills.insert(0, oldsizetills[0]*old_axis_len)
return tuple(((flattened_idx % sizetill) // stride)
for stride, sizetill in zip(oldstrides, oldsizetills))
def map_Pow(self, expr): # noqa
# SymEngine likes to use as e**a to express exp(a); we undo that here.
base, exp = expr.args
if base == symengine.E:
return prim.Variable("exp")(self.rec(exp))
else:
return prim.Power(self.rec(base), self.rec(exp))
def expression_kernel(expr, args):
r"""Produce a :class:`pyop2.Kernel` from the processed UFL expression
expr and the corresponding args."""
# Empty slot indicating assignment to indexed LHS, so don't do anything
if type(expr) is Zero:
return
fs = args[0].function.function_space()
import islpy as isl
inames = isl.make_zero_and_vars(["d"])
domain = (inames[0].le_set(
inames["d"])) & (inames["d"].lt_set(inames[0] + fs.dof_dset.cdim))
context = Bag()
context.within_inames = frozenset(["d"])
context.indices = (p.Variable("d"), )
insn = loopy_instructions(expr, context)
data = [arg.arg for arg in args]
knl = loopy.make_function([domain], [insn],
data,
name="expression",
silenced_warnings=["summing_if_branches_ops"])
return op2.Kernel(knl, "expression")
def make_random_expression(var_values, size):
from random import randrange
import pymbolic.primitives as p
v = randrange(1500)
size[0] += 1
if v < 500 and size[0] < 40:
term_count = randrange(2, 5)
if randrange(2) < 1:
cls = p.Sum
else:
cls = p.Product
return cls(
tuple(
make_random_expression(var_values, size)
for i in range(term_count)))
elif v < 750:
return make_random_value()
elif v < 1000:
var_name = "var_%d" % len(var_values)
assert var_name not in var_values
var_values[var_name] = make_random_value()
return p.Variable(var_name)
elif v < 1250:
# Cannot use '-' because that destroys numpy constants.
return p.Sum((make_random_expression(var_values, size),
-make_random_expression(var_values, size)))
elif v < 1500:
# Cannot use '/' because that destroys numpy constants.
return p.Quotient(make_random_expression(var_values, size),
make_random_expression(var_values, size))
def statement_for(tree, ctx):
extent = tree.index.extent
assert extent
idx = ctx.name_gen(ctx.index_names[tree.index])
ctx.index_extent[idx] = extent
with active_indices({tree.index: p.Variable(idx)}, ctx) as ctx_active:
return statement(tree.children[0], ctx_active)
def pymbolic_multiindex(self, shape):
indices = []
for extent in shape:
name = self.name_gen(self.index_names[extent])
self.index_extent[name] = extent
indices.append(p.Variable(name))
return tuple(indices)
def _(poly: ct.NasaPoly2, arg_name):
def gen(c, t):
assert len(c) == 7
return (c[0] + c[1] / 2 * t + c[2] / 3 * t**2 + c[3] / 4 * t**3 +
c[4] / 5 * t**4 + c[5] / t)
return nasa7_conditional(p.Variable(arg_name), poly, gen)
