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 expression_indexed(expr, parameters):
aggregate, multiindex = (expression(c, parameters) for c in expr.children)
return pym.Subscript(aggregate, multiindex)
extents = [int(numpy.prod(expr.aggregate.shape[i+1:])) for i in range(len(multiindex))]
make_sum = lambda x, y: pym.Sum((x, y))
index = reduce(make_sum, [pym.Product((e, m)) for e, m in zip(extents, multiindex)])
return pym.Subscript(aggregate, (index,))
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 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_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 build_ass():
# A_T[i,j] = sum(k, A0[i,j,k] * G_T[k]);
# Get variable symbols for all required variables
i, j, k = inames["i"], inames["j"], inames["k"]
A_T, A0, G_T = args["A_T"], args["A0"], args["G_T"]
# The target of the assignment
target = pb.Subscript(A_T, (i, j))
# The rhs expression: Frobenius inner product <A0[i,j],G_T>
reduce_op = lp.library.reduction.SumReductionOperation()
reduce_expr = pb.Subscript(A0, (i, j, k)) * pb.Subscript(G_T, (k))
expr = lp.Reduction(reduce_op, k, reduce_expr)
return lp.Assignment(target, expr)
def pymbolic_variable(self, node):
pym = self._gem_to_pym_var(node)
if node in self.indices:
indices = self.fetch_multiindex(self.indices[node])
if indices:
return p.Subscript(pym, indices)
return pym
def generate_lhs(self, tensor, temp):
""" Generation of an lhs for the loopy kernel,
which contains the TSFC assembly of the tensor.
"""
idx = self.bag.index_creator(self.shape(tensor))
lhs = pym.Subscript(temp, idx)
return SubArrayRef(idx, lhs)
def _expression_indexed(expr, ctx):
rank = ctx.pym_multiindex(expr.multiindex)
var = expression(expr.children[0], ctx)
if isinstance(var, p.Subscript):
rank = var.index + rank
var = var.aggregate
return p.Subscript(var, rank)
def _apply_elem_wise_func(x: Array, func_name: str) -> IndexLambda:
if x.dtype.kind != "f":
raise ValueError(f"'{func_name}' does not support '{x.dtype}' arrays.")
expr = prim.Call(var(f"pytato.c99.{func_name}"), (prim.Subscript(
var("in"), tuple(var(f"_{i}") for i in range(len(x.shape)))), ))
return IndexLambda(x.namespace, expr, x.shape, x.dtype, {"in": x})
def map_SympyField(self, expr):
f = expr.field
if expr.subscript is not None:
subscript = tuple(self.rec(i) for i in expr.subscript)
return pp.Subscript(f, subscript)
else:
return f
def build_ass():
"""
A[i,j] = c*sum(k, B[k,i]*B[k,j])
"""
# The target of the assignment
target = pb.Subscript(args["A"], (inames["i"], inames["j"]))
# The rhs expression: A reduce operation of the matrix columns
# Maybe replace with manual increment?
reduce_op = lp.library.reduction.SumReductionOperation()
reduce_expr = pb.Subscript(args["B"],
(inames["k"], inames["i"])) * pb.Subscript(
args["B"], (inames["k"], inames["j"]))
expr = args["c"] * lp.Reduction(reduce_op, inames["k"], reduce_expr)
return lp.Assignment(target, expr)
def test_math_function(target, tp):
# Test correct maths functions are generated for C and OpenCL
# backend instead for different data type
data_type = {"f32": np.float32, "f64": np.float64}[tp]
import pymbolic.primitives as p
i = p.Variable("i")
xi = p.Subscript(p.Variable("x"), i)
yi = p.Subscript(p.Variable("y"), i)
zi = p.Subscript(p.Variable("z"), i)
n = 100
domain = "{[i]: 0<=i<%d}" % n
data = [
lp.GlobalArg("x", data_type, shape=(n, )),
lp.GlobalArg("y", data_type, shape=(n, )),
lp.GlobalArg("z", data_type, shape=(n, ))
]
inst = [lp.Assignment(xi, p.Variable("min")(yi, zi))]
knl = lp.make_kernel(domain, inst, data, target=target())
code = lp.generate_code_v2(knl).device_code()
assert "fmin" in code
if tp == "f32" and target == CTarget:
assert "fminf" in code
else:
assert "fminf" not in code
inst = [lp.Assignment(xi, p.Variable("max")(yi, zi))]
knl = lp.make_kernel(domain, inst, data, target=target())
code = lp.generate_code_v2(knl).device_code()
assert "fmax" in code
if tp == "f32" and target == CTarget:
assert "fmaxf" in code
else:
assert "fmaxf" not in code
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():
loopy_tensor = loopy.TemporaryVariable(gem_tensor.name,
shape=gem_tensor.shape,
address_space=loopy.AddressSpace.LOCAL)
tensor2temp[slate_tensor] = loopy_tensor
if isinstance(slate_tensor, slate.Tensor):
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)))
elif isinstance(slate_tensor, slate.AssembledVector):
f = slate_tensor._function
coeff = coefficients[f]
offset = 0
ismixed = (type(f.ufl_element()) == MixedElement)
names = [name for (name, ext) in coeff.values()] if ismixed else coeff[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_field(self, expr, *args, **kwargs):
if expr.ignore_prepends:
pre_index = ()
else:
prepend = kwargs.get("prepend_with") or ()
pre_index = tuple(parse_if_str(x) for x in prepend)
pre_index = pre_index + kwargs.pop("outer_subscript", ())
full_index = pre_index + expr.index_tuple
if full_index == tuple():
x = expr.child
else:
if isinstance(expr.child, pp.Subscript):
full_index = pre_index + expr.child.index_tuple + expr.index_tuple
x = pp.Subscript(expr.child.aggregate, self.rec(full_index))
else:
x = pp.Subscript(expr.child, self.rec(full_index))
return self.rec(x, *args, **kwargs)
def loopify_tsfc_kernel_data(self, kernel_data):
""" This method generates loopy arguments from the kernel data,
which are then fed to the TSFC loopy kernel. The arguments
are arrays and have to be fed element by element to loopy
aka they have to be subarrayrefed.
"""
arguments = []
for c, name in kernel_data:
extent = self.extent(c)
idx = self.bag.index_creator(extent)
arguments.append(SubArrayRef(idx, pym.Subscript(pym.Variable(name), idx)))
return arguments
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 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 _get_sub_array_ref(array: Array,
name: str) -> "lp.symbolic.SubArrayRef":
inames = tuple(
state.var_name_gen(f"_{name}_dim{d}")
for d in range(array.ndim))
domains.append(
domain_for_shape(
inames, shape_to_scalar_expression(array.shape, self,
state), {}))
inames_as_vars = tuple(var(iname) for iname in inames)
return SubArrayRef(inames_as_vars,
prim.Subscript(var(name), inames_as_vars))
def _expression_flexiblyindexed(expr, ctx):
var = expression(expr.children[0], ctx)
rank = []
for off, idxs in expr.dim2idxs:
for index, stride in idxs:
assert isinstance(index, gem.Index)
rank_ = [off]
for index, stride in idxs:
rank_.append(p.Product((ctx.active_indices[index], stride)))
rank.append(p.Sum(tuple(rank_)))
return p.Subscript(var, tuple(rank))
def pymbolic_variable(self, node):
try:
pym = self.gem_to_pymbolic[node]
except KeyError:
name = self.name_gen(node.name)
pym = p.Variable(name)
self.gem_to_pymbolic[node] = pym
if node in self.indices:
indices = self.pym_multiindex(self.indices[node])
if indices:
return p.Subscript(pym, indices)
else:
return pym
else:
return pym
def statement_evaluate(leaf, ctx):
expr = leaf.expression
if isinstance(expr, gem.ListTensor):
ops = []
var = ctx.pymbolic_variable(expr)
index = ()
if isinstance(var, p.Subscript):
var, index = var.aggregate, var.index_tuple
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 []
else:
return [lp.Assignment(ctx.pymbolic_variable(expr), expression(expr, ctx, top=True), within_inames=ctx.active_inames())]
def _apply_elem_wise_func(inputs: Tuple[ArrayOrScalar, ...],
func_name: str,
ret_dtype: Optional[_dtype_any] = None,
np_func_name: Optional[str] = None
) -> ArrayOrScalar:
if all(isinstance(x, SCALAR_CLASSES) for x in inputs):
if np_func_name is None:
np_func_name = func_name
np_func = getattr(np, np_func_name)
return np_func(*inputs) # type: ignore
if not inputs:
raise ValueError("at least one argument must be present")
shape = None
sym_args = []
bindings = {}
for index, inp in enumerate(inputs):
if isinstance(inp, Array):
if inp.dtype.kind not in ["f", "c"]:
raise ValueError("only floating-point or complex "
"arguments supported")
if shape is None:
shape = inp.shape
elif inp.shape != shape:
# FIXME: merge this logic with arithmetic, so that broadcasting
# is implemented properly
raise NotImplementedError("broadcasting in function application")
if ret_dtype is None:
ret_dtype = inp.dtype
bindings[f"in_{index}"] = inp
sym_args.append(
prim.Subscript(var(f"in_{index}"),
tuple(var(f"_{i}") for i in range(len(shape)))))
else:
sym_args.append(inp)
assert shape is not None
assert ret_dtype is not None
return IndexLambda(
prim.Call(var(f"pytato.c99.{func_name}"), tuple(sym_args)),
shape, ret_dtype, bindings, axes=_get_default_axes(len(shape)))
def map_Subscript(self, expr): # noqa
# (expr value, slice slice, expr_context ctx)
def none_or_rec(x):
if x is None:
return x
else:
return self.rec(x)
if isinstance(expr.slice, slice):
index = slice(none_or_rec(expr.slice.start),
none_or_rec(expr.slice.stop),
none_or_rec(expr.slice.step))
else:
index = none_or_rec(expr.slice)
return p.Subscript(self.rec(expr.value), index)
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 _make_reduction_lambda(op: ReductionOperation, a: Array,
axis: Optional[Union[int, Tuple[int]]],
initial: Any) -> Array:
"""
Return a :class:`IndexLambda` that performs reduction over the *axis* axes
of *a* with the reduction op *op*.
:arg op: The reduction operation to perform.
:arg a: The :class:`pytato.Array` on which to perform the reduction.
:arg axis: The axes over which the reduction is to be performed. If axis is
*None*, perform reduction over all of *a*'s axes.
"""
new_shape, reduction_axes = _normalize_reduction_axes(a.shape, axis)
del axis
indices, redn_bounds = _get_reduction_indices_bounds(
a.shape, reduction_axes)
if initial is _NoValue:
for iax in reduction_axes:
shape_iax = a.shape[iax]
from pytato.utils import are_shape_components_equal
if are_shape_components_equal(shape_iax, 0):
raise ValueError(
"zero-size reduction operation with no supplied "
"'initial' value")
if isinstance(iax, Array):
raise NotImplementedError(
"cannot statically determine emptiness of "
f"reduction axis {iax} (0-based)")
elif initial != op.neutral_element(a.dtype):
raise NotImplementedError("reduction with 'initial' not equal to the "
"neutral element")
return make_index_lambda(
Reduce(prim.Subscript(prim.Variable("in"), tuple(indices)), op,
redn_bounds), {"in": a}, new_shape, a.dtype)
def _as_array_or_scalar(exprs: Sequence[ScalarExpression],
bindings: Mapping[str, Array],
out_shape: ShapeType
) -> Tuple[ArrayOrScalar, ...]:
"""
Helper routine invoked in :func:`index_lambda_to_high_level_op`. For every
expression in *exprs* either infers (and returns) it as a scalar or infers
(and returns) the array corresponding to one of *bindings* while defining
an :class:`~pytato.array.IndexLambda` of *out_shape* shape.
"""
result: List[ArrayOrScalar] = []
if out_shape != get_shape_after_broadcasting(bindings.values()):
raise UnknownIndexLambdaExpr()
binding_to_subscript = {bnd_name: p.Subscript(
p.Variable(bnd_name),
get_indexing_expression(bnd.shape,
out_shape))
for bnd_name, bnd in bindings.items()}
for expr in exprs:
if isinstance(expr, SCALAR_CLASSES):
result.append(expr)
elif (isinstance(expr, p.Variable)
and bindings[expr.name].shape == ()):
result.append(bindings[expr.name])
elif (isinstance(expr, p.Subscript)
and isinstance(expr.aggregate, p.Variable)
and (binding_to_subscript[expr.aggregate.name]
== expr)):
result.append(bindings[expr.aggregate.name])
else:
raise UnknownIndexLambdaExpr()
return tuple(result)
def map_einsum(self, expr: Einsum) -> Array:
import operator
from functools import reduce
from pytato.scalar_expr import Reduce
from pytato.utils import (dim_to_index_lambda_components,
are_shape_components_equal)
from pytato.array import ElementwiseAxis, ReductionAxis
bindings = {f"in{k}": self.rec(arg) for k, arg in enumerate(expr.args)}
redn_bounds: Dict[str, Tuple[ScalarExpression, ScalarExpression]] = {}
args_as_pym_expr: List[prim.Subscript] = []
namegen = UniqueNameGenerator(set(bindings))
# {{{ add bindings coming from the shape expressions
for access_descr, (iarg, arg) in zip(expr.access_descriptors,
enumerate(expr.args)):
subscript_indices = []
for iaxis, axis in enumerate(access_descr):
if not are_shape_components_equal(
arg.shape[iaxis],
expr._access_descr_to_axis_len()[axis]):
# axis is broadcasted
assert are_shape_components_equal(arg.shape[iaxis], 1)
subscript_indices.append(0)
continue
if isinstance(axis, ElementwiseAxis):
subscript_indices.append(prim.Variable(f"_{axis.dim}"))
else:
assert isinstance(axis, ReductionAxis)
redn_idx_name = f"_r{axis.dim}"
if redn_idx_name not in redn_bounds:
# convert the ShapeComponent to a ScalarExpression
redn_bound, redn_bound_bindings = (
dim_to_index_lambda_components(
arg.shape[iaxis], namegen))
redn_bounds[redn_idx_name] = (0, redn_bound)
bindings.update({k: self.rec(v)
for k, v in redn_bound_bindings.items()})
subscript_indices.append(prim.Variable(redn_idx_name))
args_as_pym_expr.append(prim.Subscript(prim.Variable(f"in{iarg}"),
tuple(subscript_indices)))
# }}}
inner_expr = reduce(operator.mul, args_as_pym_expr[1:],
args_as_pym_expr[0])
if redn_bounds:
from pytato.reductions import SumReductionOperation
inner_expr = Reduce(inner_expr,
SumReductionOperation(),
redn_bounds)
return IndexLambda(expr=inner_expr,
shape=tuple(self.rec(s) if isinstance(s, Array) else s
for s in expr.shape),
dtype=expr.dtype,
bindings=bindings,
axes=expr.axes,
tags=expr.tags)
def parse_postfix(self, pstate, min_precedence, left_exp):
import pymbolic.primitives as primitives
import pymbolic.parser as p
did_something = False
next_tag = pstate.next_tag()
if next_tag is p._openpar and p._PREC_CALL > min_precedence:
pstate.advance()
pstate.expect_not_end()
if next_tag is p._closepar:
pstate.advance()
left_exp = primitives.Call(left_exp, ())
else:
args = self.parse_expression(pstate)
if not isinstance(args, tuple):
args = (args, )
pstate.expect(p._closepar)
pstate.advance()
if left_exp == primitives.Variable("matrix"):
left_exp = np.array(list(list(row) for row in args))
else:
left_exp = primitives.Call(left_exp, args)
did_something = True
elif next_tag is p._openbracket and p._PREC_CALL > min_precedence:
pstate.advance()
pstate.expect_not_end()
left_exp = primitives.Subscript(left_exp,
self.parse_expression(pstate))
pstate.expect(p._closebracket)
pstate.advance()
did_something = True
elif next_tag is p._dot and p._PREC_CALL > min_precedence:
pstate.advance()
pstate.expect(p._identifier)
left_exp = primitives.Lookup(left_exp, pstate.next_str())
pstate.advance()
did_something = True
elif next_tag is p._plus and p._PREC_PLUS > min_precedence:
pstate.advance()
left_exp += self.parse_expression(pstate, p._PREC_PLUS)
did_something = True
elif next_tag is p._minus and p._PREC_PLUS > min_precedence:
pstate.advance()
left_exp -= self.parse_expression(pstate, p._PREC_PLUS)
did_something = True
elif next_tag is p._times and p._PREC_TIMES > min_precedence:
pstate.advance()
left_exp *= self.parse_expression(pstate, p._PREC_TIMES)
did_something = True
elif next_tag is p._over and p._PREC_TIMES > min_precedence:
pstate.advance()
from pymbolic.primitives import Quotient
left_exp = Quotient(left_exp,
self.parse_expression(pstate, p._PREC_TIMES))
did_something = True
elif next_tag is self.power_sym and p._PREC_POWER > min_precedence:
pstate.advance()
exponent = self.parse_expression(pstate, p._PREC_POWER)
if left_exp == np.e:
from pymbolic.primitives import Call, Variable
left_exp = Call(Variable("exp"), (exponent, ))
else:
left_exp **= exponent
did_something = True
elif next_tag is p._comma and p._PREC_COMMA > min_precedence:
# The precedence makes the comma left-associative.
pstate.advance()
if pstate.is_at_end() or pstate.next_tag() is p._closepar:
left_exp = (left_exp, )
else:
new_el = self.parse_expression(pstate, p._PREC_COMMA)
if isinstance(left_exp, tuple) \
and not isinstance(left_exp, FinalizedTuple):
left_exp = left_exp + (new_el, )
else:
left_exp = (left_exp, new_el)
did_something = True
return left_exp, did_something
def map_Indexed(self, expr): # noqa
return prim.Subscript(
self.rec(expr.args[0].args[0]),
tuple(self.rec(i) for i in expr.args[1:])
)
