def _generic_product(self, other, product_class):
"""
:arg product_class: A subclass of :class:`_GAProduct`.
"""
if self.space.is_orthogonal:
bpw = product_class.orthogonal_blade_product_weight
else:
bpw = product_class.generic_blade_product_weight
if self.space is not other.space:
raise ValueError("can only compute products of multivectors "
"from identical spaces")
from pymbolic.primitives import is_zero
new_data = {}
for sbits, scoeff in six.iteritems(self.data):
for obits, ocoeff in six.iteritems(other.data):
new_bits = sbits ^ obits
weight = bpw(sbits, obits, self.space)
if not is_zero(weight):
# These are nonzero by definition.
coeff = (weight * canonical_reordering_sign(sbits, obits) *
scoeff * ocoeff)
new_coeff = new_data.setdefault(new_bits, 0) + coeff
if is_zero(new_coeff):
del new_data[new_bits]
else:
new_data[new_bits] = new_coeff
return MultiVector(new_data, self.space)
def map_power(self, expr, type_context):
tgt_dtype = self.infer_type(expr)
exponent_dtype = self.infer_type(expr.exponent)
from pymbolic.primitives import is_constant, is_zero
if is_constant(expr.exponent):
if is_zero(expr.exponent):
return 1
elif is_zero(expr.exponent - 1):
return self.rec(expr.base, type_context)
elif is_zero(expr.exponent - 2):
return self.rec(expr.base * expr.base, type_context)
if exponent_dtype.is_integral():
from loopy.codegen import SeenFunction
func_name = (
"loopy_pow_"
f"{tgt_dtype.numpy_dtype}_{exponent_dtype.numpy_dtype}")
self.codegen_state.seen_functions.add(
SeenFunction("int_pow", func_name, (tgt_dtype, exponent_dtype),
(tgt_dtype, )))
return var(func_name)(self.rec(expr.base, type_context),
self.rec(expr.exponent, type_context))
else:
return self.rec(var("pow")(expr.base, expr.exponent), type_context)
def _generic_product(self, other, product_class):
"""
:arg product_class: A subclass of :class:`_GAProduct`.
"""
if self.space.is_orthogonal:
bpw = product_class.orthogonal_blade_product_weight
else:
bpw = product_class.generic_blade_product_weight
if self.space is not other.space:
raise ValueError("can only compute products of multivectors "
"from identical spaces")
from pymbolic.primitives import is_zero
new_data = {}
for sbits, scoeff in six.iteritems(self.data):
for obits, ocoeff in six.iteritems(other.data):
new_bits = sbits ^ obits
weight = bpw(sbits, obits, self.space)
if not is_zero(weight):
# These are nonzero by definition.
coeff = (weight
* canonical_reordering_sign(sbits, obits)
* scoeff * ocoeff)
new_coeff = new_data.setdefault(new_bits, 0) + coeff
if is_zero(new_coeff):
del new_data[new_bits]
else:
new_data[new_bits] = new_coeff
return MultiVector(new_data, self.space)
def base_impl(expr, type_context):
from pymbolic.primitives import is_constant, is_zero
if is_constant(expr.exponent):
if is_zero(expr.exponent):
return 1
elif is_zero(expr.exponent - 1):
return self.rec(expr.base, type_context)
elif is_zero(expr.exponent - 2):
return self.rec(expr.base * expr.base, type_context)
return type(expr)(self.rec(expr.base, type_context),
self.rec(expr.exponent, type_context))
def map_power(self, expr, enclosing_prec):
from pymbolic.mapper.stringifier import PREC_NONE
from pymbolic.primitives import is_constant, is_zero
if is_constant(expr.exponent):
if is_zero(expr.exponent):
return "1"
elif is_zero(expr.exponent - 1):
return self.rec(expr.base, enclosing_prec)
elif is_zero(expr.exponent - 2):
return self.rec(expr.base * expr.base, enclosing_prec)
return self.format("pow(%s, %s)", self.rec(expr.base, PREC_NONE),
self.rec(expr.exponent, PREC_NONE))
def map_power(self, expr, enclosing_prec):
from pymbolic.mapper.stringifier import PREC_NONE
from pymbolic.primitives import is_constant, is_zero
if is_constant(expr.exponent):
if is_zero(expr.exponent):
return "1"
elif is_zero(expr.exponent - 1):
return self.rec(expr.base, enclosing_prec)
elif is_zero(expr.exponent - 2):
return self.rec(expr.base*expr.base, enclosing_prec)
return self.format("pow(%s, %s)",
self.rec(expr.base, PREC_NONE),
self.rec(expr.exponent, PREC_NONE))
def base_impl(expr, enclosing_prec, type_context):
from pymbolic.mapper.stringifier import PREC_NONE
from pymbolic.primitives import is_constant, is_zero
if is_constant(expr.exponent):
if is_zero(expr.exponent):
return "1"
elif is_zero(expr.exponent - 1):
return self.rec(expr.base, enclosing_prec, type_context)
elif is_zero(expr.exponent - 2):
return self.rec(
expr.base*expr.base, enclosing_prec, type_context)
return "pow(%s, %s)" % (
self.rec(expr.base, PREC_NONE, type_context),
self.rec(expr.exponent, PREC_NONE, type_context))
def base_impl(expr, enclosing_prec, type_context):
from pymbolic.mapper.stringifier import PREC_NONE
from pymbolic.primitives import is_constant, is_zero
if is_constant(expr.exponent):
if is_zero(expr.exponent):
return "1"
elif is_zero(expr.exponent - 1):
return self.rec(expr.base, enclosing_prec, type_context)
elif is_zero(expr.exponent - 2):
return self.rec(expr.base * expr.base, enclosing_prec,
type_context)
return "pow(%s, %s)" % (
self.rec(expr.base, PREC_NONE, type_context),
self.rec(expr.exponent, PREC_NONE, type_context))
def map_ComplexDouble(self, expr): # noqa
r = self.rec(expr.real_part())
i = self.rec(expr.imaginary_part())
if prim.is_zero(i):
return r
else:
return r + 1j * i
def map_Rational(self, expr):
num = self.rec(expr.p)
denom = self.rec(expr.q)
if prim.is_zero(denom - 1):
return num
return prim.Quotient(num, denom)
def map_ComplexDouble(self, expr): # noqa
r = self.rec(expr.real_part())
i = self.rec(expr.imaginary_part())
if prim.is_zero(i):
return r
else:
return r + 1j * i
def map_common_subexpression(self, expr, *args, **kwargs):
from pymbolic.primitives import is_zero
result = self.rec(expr.child, *args, **kwargs)
if is_zero(result):
return 0
return type(expr)(result, expr.prefix, **expr.get_extra_properties())
def map_Rational(self, expr):
num = self.rec(expr.p)
denom = self.rec(expr.q)
if prim.is_zero(denom-1):
return num
return prim.Quotient(num, denom)
def map_product(self, expr, enclosing_prec, *args, **kwargs):
entries = []
i = 0
from pymbolic.primitives import is_zero
while i < len(expr.children):
child = expr.children[i]
if False and is_zero(child + 1) and i + 1 < len(expr.children):
# NOTE: That space needs to be there.
# Otherwise two unary minus signs merge into a pre-decrement.
entries.append(
self.format(
"- %s",
self.rec(expr.children[i + 1], PREC_UNARY, *args,
**kwargs)))
i += 2
else:
entries.append(self.rec(child, PREC_PRODUCT, *args, **kwargs))
i += 1
entries.sort(reverse=self.reverse)
result = "*".join(entries)
return self.parenthesize_if_needed(result, enclosing_prec,
PREC_PRODUCT)
def map_quotient(self, expr):
if is_zero(expr.numerator - 1):
return expr
else:
# not the smartest thing we can do, but at least *something*
return pymbolic.flattened_product([
type(expr)(1, self.rec(expr.denominator)),
self.rec(expr.numerator)])
def map_product(self, expr):
from pymbolic.primitives import is_zero
assert expr.children
return sum(ToCountMap({(self.type_inf(expr), 'mul'): 1})
+ self.rec(child)
for child in expr.children
if not is_zero(child + 1)) + \
ToCountMap({(self.type_inf(expr), 'mul'): -1})
def map_product(self, expr):
from pymbolic.primitives import is_zero
assert expr.children
return sum(ToCountMap({(self.type_inf(expr), 'mul'): 1})
+ self.rec(child)
for child in expr.children
if not is_zero(child + 1)) + \
ToCountMap({(self.type_inf(expr), 'mul'): -1})
def map_Rational(self, expr): # noqa
p, q = expr.p, expr.q
num = self.rec(p)
denom = self.rec(q)
if prim.is_zero(denom - 1):
return num
return prim.Quotient(num, denom)
def map_quotient(self, expr):
if is_zero(expr.numerator - 1):
return expr
else:
# not the smartest thing we can do, but at least *something*
return pymbolic.flattened_product([
type(expr)(1, self.rec(expr.denominator)),
self.rec(expr.numerator)
])
def map_Rational(self, expr): # noqa
p, q = expr.p, expr.q
num = self.rec(p)
denom = self.rec(q)
if prim.is_zero(denom-1):
return num
return prim.Quotient(num, denom)
def eval_arg(arg_spec):
arg_expr, is_int = arg_spec
arg = self.rec(arg_expr)
if is_zero(arg):
if insn.is_boundary and not is_int:
return BoundaryZeros()
else:
return VolumeZeros()
else:
return arg
def map_common_subexpression(self, expr, *args, **kwargs):
from pymbolic.primitives import is_zero
result = self.rec(expr.child, *args, **kwargs)
if is_zero(result):
return 0
return type(expr)(
result,
expr.prefix,
**expr.get_extra_properties())
def map_elementwise_linear(self, op, field_expr):
field = self.rec(field_expr)
from hedge.tools import is_zero
if is_zero(field):
return 0
out = self.discr.volume_zeros()
self.executor.do_elementwise_linear(op, field, out)
return out
def eval_arg(arg_spec):
arg_expr, is_int = arg_spec
arg = self.rec(arg_expr)
if is_zero(arg):
if insn.is_boundary and not is_int:
return BoundaryZeros()
else:
return VolumeZeros()
else:
return arg
def map_elementwise_linear(self, op, field_expr):
field = self.rec(field_expr)
from hedge.tools import is_zero
if is_zero(field):
return 0
out = self.discr.volume_zeros()
self.executor.do_elementwise_linear(op, field, out)
return out
def get_neg_product(expr):
from pymbolic.primitives import is_zero, Product
if isinstance(expr, Product) \
and len(expr.children) and is_zero(expr.children[0]+1):
if len(expr.children) == 2:
# only the minus sign and the other child
return expr.children[1]
else:
return Product(expr.children[1:])
else:
return None
def get_neg_product(expr):
from pymbolic.primitives import is_zero, Product
if isinstance(expr, Product) \
and len(expr.children) and is_zero(expr.children[0]+1):
if len(expr.children) == 2:
# only the minus sign and the other child
return expr.children[1]
else:
return Product(expr.children[1:])
else:
return None
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)
from pymbolic.primitives import is_constant, is_zero
if is_constant(expr.exponent):
if is_zero(expr.exponent):
return 1
elif is_zero(expr.exponent - 1):
return self.rec(expr.base, type_context)
elif is_zero(expr.exponent - 2):
return self.rec(expr.base * expr.base, type_context)
if exponent_dtype.is_integral():
from loopy.codegen import SeenFunction
func_name = (
"loopy_pow_"
f"{tgt_dtype.numpy_dtype}_{exponent_dtype.numpy_dtype}")
self.codegen_state.seen_functions.add(
SeenFunction("int_pow", func_name, (tgt_dtype, exponent_dtype),
(tgt_dtype, )))
# FIXME: This need some more callables to be registered.
return var(func_name)(self.rec(expr.base, type_context),
self.rec(expr.exponent, type_context))
else:
from loopy.codegen import SeenFunction
clbl = self.codegen_state.ast_builder.known_callables["pow"]
clbl = clbl.with_types({
0: tgt_dtype,
1: exponent_dtype
}, self.codegen_state.callables_table)[0]
self.codegen_state.seen_functions.add(
SeenFunction(clbl.name, clbl.name_in_target,
(base_dtype, exponent_dtype), (tgt_dtype, )))
return var(clbl.name_in_target)(self.rec(expr.base, type_context),
self.rec(expr.exponent,
type_context))
def __add__(self, other):
if not isinstance(other, MultiVector):
other = MultiVector(other, self.space)
if self.space is not other.space:
raise ValueError("can only add multivectors from identical spaces")
all_bits = set(self.data.iterkeys()) | set(other.data.iterkeys())
new_data = {}
for bits in all_bits:
new_coeff = self.data.get(bits, 0) + other.data.get(bits, 0)
if not is_zero(new_coeff):
new_data[bits] = new_coeff
return MultiVector(new_data, self.space)
def __add__(self, other):
other = _cast_or_ni(other, self.space)
if other is NotImplemented:
return NotImplemented
if self.space is not other.space:
raise ValueError("can only add multivectors from identical spaces")
all_bits = set(six.iterkeys(self.data)) | set(six.iterkeys(other.data))
from pymbolic.primitives import is_zero
new_data = {}
for bits in all_bits:
new_coeff = self.data.get(bits, 0) + other.data.get(bits, 0)
if not is_zero(new_coeff):
new_data[bits] = new_coeff
return MultiVector(new_data, self.space)
def __add__(self, other):
other = _cast_or_ni(other, self.space)
if other is NotImplemented:
return NotImplemented
if self.space is not other.space:
raise ValueError("can only add multivectors from identical spaces")
all_bits = set(six.iterkeys(self.data)) | set(six.iterkeys(other.data))
from pymbolic.primitives import is_zero
new_data = {}
for bits in all_bits:
new_coeff = self.data.get(bits, 0) + other.data.get(bits, 0)
if not is_zero(new_coeff):
new_data[bits] = new_coeff
return MultiVector(new_data, self.space)
def map_ref_quad_mass(self, op, field_expr):
field = self.rec(field_expr)
from hedge.tools import is_zero
if is_zero(field):
return 0
qtag = op.quadrature_tag
from hedge._internal import perform_elwise_operator
out = self.discr.volume_zeros()
for eg in self.discr.element_groups:
eg_quad_info = eg.quadrature_info[qtag]
perform_elwise_operator(eg_quad_info.ranges, eg.ranges,
eg_quad_info.ldis_quad_info.mass_matrix(),
field, out)
return out
def map_ref_quad_mass(self, op, field_expr):
field = self.rec(field_expr)
from hedge.tools import is_zero
if is_zero(field):
return 0
qtag = op.quadrature_tag
from hedge._internal import perform_elwise_operator
out = self.discr.volume_zeros()
for eg in self.discr.element_groups:
eg_quad_info = eg.quadrature_info[qtag]
perform_elwise_operator(eg_quad_info.ranges, eg.ranges,
eg_quad_info.ldis_quad_info.mass_matrix(),
field, out)
return out
def map_quad_int_faces_grid_upsampler(self, op, field_expr):
field = self.rec(field_expr)
from hedge.tools import is_zero
if is_zero(field):
return 0
qtag = op.quadrature_tag
from hedge._internal import perform_elwise_operator
quad_info = self.discr.get_quadrature_info(qtag)
out = np.zeros(quad_info.int_faces_node_count, field.dtype)
for eg in self.discr.element_groups:
eg_quad_info = eg.quadrature_info[qtag]
perform_elwise_operator(eg.ranges, eg_quad_info.el_faces_ranges,
eg_quad_info.ldis_quad_info.volume_to_face_up_interpolation_matrix(),
field, out)
return out
def map_product(self, expr, enclosing_prec):
entries = []
i = 0
from pymbolic.primitives import is_zero
while i < len(expr.children):
child = expr.children[i]
if False and is_zero(child+1) and i+1 < len(expr.children):
# NOTE: That space needs to be there.
# Otherwise two unary minus signs merge into a pre-decrement.
entries.append(
self.format("- %s", self.rec(expr.children[i+1], PREC_UNARY)))
i += 2
else:
entries.append(self.rec(child, PREC_PRODUCT))
i += 1
entries.sort(reverse=self.reverse)
result = "*".join(entries)
return self.parenthesize_if_needed(result, enclosing_prec, PREC_PRODUCT)
def map_quad_bdry_grid_upsampler(self, op, field_expr):
field = self.rec(field_expr)
from hedge.tools import is_zero
if is_zero(field):
return 0
bdry = self.discr.get_boundary(op.boundary_tag)
bdry_q_info = bdry.get_quadrature_info(op.quadrature_tag)
out = np.zeros(bdry_q_info.node_count, field.dtype)
from hedge._internal import perform_elwise_operator
for fg, from_ranges, to_ranges, ldis_quad_info in zip(
bdry.face_groups, bdry.fg_ranges, bdry_q_info.fg_ranges,
bdry_q_info.fg_ldis_quad_infos):
perform_elwise_operator(
from_ranges, to_ranges,
ldis_quad_info.face_up_interpolation_matrix(), field, out)
return out
def map_quad_int_faces_grid_upsampler(self, op, field_expr):
field = self.rec(field_expr)
from hedge.tools import is_zero
if is_zero(field):
return 0
qtag = op.quadrature_tag
from hedge._internal import perform_elwise_operator
quad_info = self.discr.get_quadrature_info(qtag)
out = np.zeros(quad_info.int_faces_node_count, field.dtype)
for eg in self.discr.element_groups:
eg_quad_info = eg.quadrature_info[qtag]
perform_elwise_operator(
eg.ranges, eg_quad_info.el_faces_ranges,
eg_quad_info.ldis_quad_info.
volume_to_face_up_interpolation_matrix(), field, out)
return out
def map_quad_bdry_grid_upsampler(self, op, field_expr):
field = self.rec(field_expr)
from hedge.tools import is_zero
if is_zero(field):
return 0
bdry = self.discr.get_boundary(op.boundary_tag)
bdry_q_info = bdry.get_quadrature_info(op.quadrature_tag)
out = np.zeros(bdry_q_info.node_count, field.dtype)
from hedge._internal import perform_elwise_operator
for fg, from_ranges, to_ranges, ldis_quad_info in zip(
bdry.face_groups,
bdry.fg_ranges,
bdry_q_info.fg_ranges,
bdry_q_info.fg_ldis_quad_infos):
perform_elwise_operator(from_ranges, to_ranges,
ldis_quad_info.face_up_interpolation_matrix(),
field, out)
return out
def get_flux_var_info(fluxes):
from pytools import Record
class FluxVariableInfo(Record):
pass
scalar_parameters = set()
fvi = FluxVariableInfo(
scalar_parameters=None,
arg_specs=[],
arg_names=[],
flux_idx_and_dep_to_arg_name={}, # or 0 if zero
)
field_expr_to_arg_name = {}
from hedge.flux import \
FieldComponent, FluxDependencyMapper, \
FluxScalarParameter
from hedge.optemplate import BoundaryPair
for flux_idx, flux_binding in enumerate(fluxes):
for dep in FluxDependencyMapper(include_calls=False)(flux_binding.op.flux):
if isinstance(dep, FluxScalarParameter):
scalar_parameters.add(dep)
elif isinstance(dep, FieldComponent):
is_bdry = isinstance(flux_binding.field, BoundaryPair)
if is_bdry:
if dep.is_interior:
this_field_expr = flux_binding.field.field
else:
this_field_expr = flux_binding.field.bfield
else:
this_field_expr = flux_binding.field
from hedge.tools import is_obj_array
if is_obj_array(this_field_expr):
fc_field_expr = this_field_expr[dep.index]
else:
assert dep.index == 0
fc_field_expr = this_field_expr
def set_or_check(dict_instance, key, value):
try:
existing_value = dict_instance[key]
except KeyError:
dict_instance[key] = value
else:
assert existing_value == value
from pymbolic.primitives import is_zero
if is_zero(fc_field_expr):
fvi.flux_idx_and_dep_to_arg_name[flux_idx, dep] = 0
else:
if fc_field_expr not in field_expr_to_arg_name:
arg_name = "arg%d" % len(fvi.arg_specs)
field_expr_to_arg_name[fc_field_expr] = arg_name
fvi.arg_names.append(arg_name)
fvi.arg_specs.append((fc_field_expr, dep.is_interior))
else:
arg_name = field_expr_to_arg_name[fc_field_expr]
set_or_check(
fvi.flux_idx_and_dep_to_arg_name,
(flux_idx, dep),
arg_name)
if not is_bdry:
# Interior fluxes are used flipped as well.
# Make sure we have assigned arg names for the
# flipped case as well.
set_or_check(
fvi.flux_idx_and_dep_to_arg_name,
(flux_idx,
FieldComponent(dep.index, not dep.is_interior)),
arg_name)
else:
raise ValueError("unknown flux dependency type: %s" % dep)
fvi.scalar_parameters = list(scalar_parameters)
return fvi
def bind_one(subexpr):
if p.is_zero(subexpr):
return subexpr
else:
return OperatorBinding(self, subexpr)
def __init__(self, data, space=None):
"""
:arg data: This may be one of the following:
* a :class:`numpy.ndarray`, which will be turned into a grade-1
multivector,
* a mapping from tuples of basis indices (together indicating a blade,
order matters and will be mapped to 'normalized' blades) to
coefficients,
* an array as described in :attr:`data`,
* a scalar--where everything that doesn't fall into the above cases
is viewed as a scalar.
:arg space: A :class:`Space` instance. If *None* or an integer,
:func:`get_euclidean_space` is called to obtain a default space with
the right number of dimensions for *data*. Note: dimension guessing only
works when a :class:`numpy.ndarray` is being passed for *data*.
"""
dimensions = None
if isinstance(data, np.ndarray):
if len(data.shape) != 1:
raise ValueError("only numpy vectors (not higher-rank objects) "
"are supported for 'data'")
dimensions, = data.shape
data = dict(
((i,), xi) for i, xi in enumerate(data))
elif isinstance(data, dict):
pass
else:
data = {0: data}
if space is None:
space = get_euclidean_space(dimensions)
else:
if dimensions is not None and space.dimensions != dimensions:
raise ValueError(
"dimension count of 'space' does not match that of 'data'")
# {{{ normalize data to bitmaps, if needed
from pytools import single_valued
from pymbolic.primitives import is_zero
if data and single_valued(isinstance(k, tuple) for k in six.iterkeys(data)):
# data is in non-normalized non-bits tuple form
new_data = {}
for basis_indices, coeff in six.iteritems(data):
bits, sign = space.bits_and_sign(basis_indices)
new_coeff = new_data.setdefault(bits, 0) + sign*coeff
if is_zero(new_coeff):
del new_data[bits]
else:
new_data[bits] = new_coeff
data = new_data
# }}}
# assert that multivectors don't get nested
from pytools import any
assert not any(isinstance(coeff, MultiVector)
for coeff in six.itervalues(data))
self.space = space
self.data = data
def get_flux_var_info(fluxes):
from pytools import Record
class FluxVariableInfo(Record):
pass
scalar_parameters = set()
fvi = FluxVariableInfo(
scalar_parameters=None,
arg_specs=[],
arg_names=[],
flux_idx_and_dep_to_arg_name={}, # or 0 if zero
)
field_expr_to_arg_name = {}
from hedge.flux import \
FieldComponent, FluxDependencyMapper, \
FluxScalarParameter
from hedge.optemplate import BoundaryPair
for flux_idx, flux_binding in enumerate(fluxes):
from hedge.optemplate.primitives import CFunction
for dep in FluxDependencyMapper(include_calls=False)(flux_binding.op.flux):
if isinstance(dep, FluxScalarParameter):
scalar_parameters.add(dep)
elif isinstance(dep, FieldComponent):
is_bdry = isinstance(flux_binding.field, BoundaryPair)
if is_bdry:
if dep.is_interior:
this_field_expr = flux_binding.field.field
else:
this_field_expr = flux_binding.field.bfield
else:
this_field_expr = flux_binding.field
from hedge.tools import is_obj_array
if is_obj_array(this_field_expr):
fc_field_expr = this_field_expr[dep.index]
else:
assert dep.index == 0
fc_field_expr = this_field_expr
def set_or_check(dict_instance, key, value):
try:
existing_value = dict_instance[key]
except KeyError:
dict_instance[key] = value
else:
assert existing_value == value
from pymbolic.primitives import is_zero
if is_zero(fc_field_expr):
fvi.flux_idx_and_dep_to_arg_name[flux_idx, dep] = 0
else:
if fc_field_expr not in field_expr_to_arg_name:
arg_name = "arg%d" % len(fvi.arg_specs)
field_expr_to_arg_name[fc_field_expr] = arg_name
fvi.arg_names.append(arg_name)
fvi.arg_specs.append((fc_field_expr, dep.is_interior))
else:
arg_name = field_expr_to_arg_name[fc_field_expr]
set_or_check(
fvi.flux_idx_and_dep_to_arg_name,
(flux_idx, dep),
arg_name)
if not is_bdry:
# Interior fluxes are used flipped as well.
# Make sure we have assigned arg names for the
# flipped case as well.
set_or_check(
fvi.flux_idx_and_dep_to_arg_name,
(flux_idx,
FieldComponent(dep.index, not dep.is_interior)),
arg_name)
elif isinstance(dep, CFunction):
pass
else:
raise ValueError("unknown flux dependency type: %s" % dep)
fvi.scalar_parameters = list(scalar_parameters)
return fvi
def collect_common_factors_on_increment(kernel, var_name, vary_by_axes=()):
# FIXME: Does not understand subst rules for now
if kernel.substitutions:
from loopy.transform.subst import expand_subst
kernel = expand_subst(kernel)
if var_name in kernel.temporary_variables:
var_descr = kernel.temporary_variables[var_name]
elif var_name in kernel.arg_dict:
var_descr = kernel.arg_dict[var_name]
else:
raise NameError("array '%s' was not found" % var_name)
# {{{ check/normalize vary_by_axes
if isinstance(vary_by_axes, str):
vary_by_axes = vary_by_axes.split(",")
from loopy.kernel.array import ArrayBase
if isinstance(var_descr, ArrayBase):
if var_descr.dim_names is not None:
name_to_index = dict(
(name, idx)
for idx, name in enumerate(var_descr.dim_names))
else:
name_to_index = {}
def map_ax_name_to_index(ax):
if isinstance(ax, str):
try:
return name_to_index[ax]
except KeyError:
raise LoopyError("axis name '%s' not understood " % ax)
else:
return ax
vary_by_axes = [map_ax_name_to_index(ax) for ax in vary_by_axes]
if (
vary_by_axes
and
(min(vary_by_axes) < 0
or
max(vary_by_axes) > var_descr.num_user_axes())):
raise LoopyError("vary_by_axes refers to out-of-bounds axis index")
# }}}
from pymbolic.mapper.substitutor import make_subst_func
from pymbolic.primitives import (Sum, Product, is_zero,
flattened_sum, flattened_product, Subscript, Variable)
from loopy.symbolic import (get_dependencies, SubstitutionMapper,
UnidirectionalUnifier)
# {{{ common factor key list maintenance
# list of (index_key, common factors found)
common_factors = []
def find_unifiable_cf_index(index_key):
for i, (key, val) in enumerate(common_factors):
unif = UnidirectionalUnifier(
lhs_mapping_candidates=get_dependencies(key))
unif_result = unif(key, index_key)
if unif_result:
assert len(unif_result) == 1
return i, unif_result[0]
return None, None
def extract_index_key(access_expr):
if isinstance(access_expr, Variable):
return ()
elif isinstance(access_expr, Subscript):
index = access_expr.index_tuple
return tuple(index[ax] for ax in vary_by_axes)
else:
raise ValueError("unexpected type of access_expr")
def is_assignee(insn):
return any(
lhs == var_name
for lhs, sbscript in insn.assignees_and_indices())
def iterate_as(cls, expr):
if isinstance(expr, cls):
for ch in expr.children:
yield ch
else:
yield expr
# }}}
# {{{ find common factors
from loopy.kernel.data import Assignment
for insn in kernel.instructions:
if not is_assignee(insn):
continue
if not isinstance(insn, Assignment):
raise LoopyError("'%s' modified by non-expression instruction"
% var_name)
lhs = insn.assignee
rhs = insn.expression
if is_zero(rhs):
continue
index_key = extract_index_key(lhs)
cf_index, unif_result = find_unifiable_cf_index(index_key)
if cf_index is None:
# {{{ doesn't exist yet
assert unif_result is None
my_common_factors = None
for term in iterate_as(Sum, rhs):
if term == lhs:
continue
for part in iterate_as(Product, term):
if var_name in get_dependencies(part):
raise LoopyError("unexpected dependency on '%s' "
"in RHS of instruction '%s'"
% (var_name, insn.id))
product_parts = set(iterate_as(Product, term))
if my_common_factors is None:
my_common_factors = product_parts
else:
my_common_factors = my_common_factors & product_parts
if my_common_factors is not None:
common_factors.append((index_key, my_common_factors))
# }}}
else:
# {{{ match, filter existing common factors
_, my_common_factors = common_factors[cf_index]
unif_subst_map = SubstitutionMapper(
make_subst_func(unif_result.lmap))
for term in iterate_as(Sum, rhs):
if term == lhs:
continue
for part in iterate_as(Product, term):
if var_name in get_dependencies(part):
raise LoopyError("unexpected dependency on '%s' "
"in RHS of instruction '%s'"
% (var_name, insn.id))
product_parts = set(iterate_as(Product, term))
my_common_factors = set(
cf for cf in my_common_factors
if unif_subst_map(cf) in product_parts)
common_factors[cf_index] = (index_key, my_common_factors)
# }}}
# }}}
# {{{ remove common factors
new_insns = []
for insn in kernel.instructions:
if not isinstance(insn, Assignment) or not is_assignee(insn):
new_insns.append(insn)
continue
(_, index_key), = insn.assignees_and_indices()
lhs = insn.assignee
rhs = insn.expression
if is_zero(rhs):
new_insns.append(insn)
continue
index_key = extract_index_key(lhs)
cf_index, unif_result = find_unifiable_cf_index(index_key)
if cf_index is None:
new_insns.append(insn)
continue
_, my_common_factors = common_factors[cf_index]
unif_subst_map = SubstitutionMapper(
make_subst_func(unif_result.lmap))
mapped_my_common_factors = set(
unif_subst_map(cf)
for cf in my_common_factors)
new_sum_terms = []
for term in iterate_as(Sum, rhs):
if term == lhs:
new_sum_terms.append(term)
continue
new_sum_terms.append(
flattened_product([
part
for part in iterate_as(Product, term)
if part not in mapped_my_common_factors
]))
new_insns.append(
insn.copy(expression=flattened_sum(new_sum_terms)))
# }}}
# {{{ substitute common factors into usage sites
def find_substitution(expr):
if isinstance(expr, Subscript):
v = expr.aggregate.name
elif isinstance(expr, Variable):
v = expr.name
else:
return expr
if v != var_name:
return expr
index_key = extract_index_key(expr)
cf_index, unif_result = find_unifiable_cf_index(index_key)
unif_subst_map = SubstitutionMapper(
make_subst_func(unif_result.lmap))
_, my_common_factors = common_factors[cf_index]
if my_common_factors is not None:
return flattened_product(
[unif_subst_map(cf) for cf in my_common_factors]
+ [expr])
else:
return expr
insns = new_insns
new_insns = []
subm = SubstitutionMapper(find_substitution)
for insn in insns:
if not isinstance(insn, Assignment) or is_assignee(insn):
new_insns.append(insn)
continue
new_insns.append(insn.with_transformed_expressions(subm))
# }}}
return kernel.copy(instructions=new_insns)
def __init__(self, data, space=None):
"""
:arg data: This may be one of the following:
* a :class:`numpy.ndarray`, which will be turned into a grade-1
multivector,
* a mapping from tuples of basis indices (together indicating a blade,
order matters and will be mapped to 'normalized' blades) to
coefficients,
* an array as described in :attr:`data`,
* a scalar--where everything that doesn't fall into the above cases
is viewed as a scalar.
:arg space: A :class:`Space` instance. If *None* or an integer,
:func:`get_euclidean_space` is called to obtain a default space with
the right number of dimensions for *data*. Note: dimension guessing only
works when a :class:`numpy.ndarray` is being passed for *data*.
"""
dimensions = None
if isinstance(data, np.ndarray):
if len(data.shape) != 1:
raise ValueError(
"only numpy vectors (not higher-rank objects) "
"are supported for 'data'")
dimensions, = data.shape
data = dict(((i, ), xi) for i, xi in enumerate(data))
elif isinstance(data, dict):
pass
else:
data = {0: data}
if space is None:
space = get_euclidean_space(dimensions)
else:
if dimensions is not None and space.dimensions != dimensions:
raise ValueError(
"dimension count of 'space' does not match that of 'data'")
# {{{ normalize data to bitmaps, if needed
from pytools import single_valued
from pymbolic.primitives import is_zero
if data and single_valued(
isinstance(k, tuple) for k in six.iterkeys(data)):
# data is in non-normalized non-bits tuple form
new_data = {}
for basis_indices, coeff in six.iteritems(data):
bits, sign = space.bits_and_sign(basis_indices)
new_coeff = new_data.setdefault(bits, 0) + sign * coeff
if is_zero(new_coeff):
del new_data[bits]
else:
new_data[bits] = new_coeff
data = new_data
# }}}
# assert that multivectors don't get nested
from pytools import any
assert not any(
isinstance(coeff, MultiVector) for coeff in six.itervalues(data))
self.space = space
self.data = data
def collect_common_factors_on_increment(kernel, var_name, vary_by_axes=()):
assert isinstance(kernel, LoopKernel)
# FIXME: Does not understand subst rules for now
if kernel.substitutions:
from loopy.transform.subst import expand_subst
kernel = expand_subst(kernel)
if var_name in kernel.temporary_variables:
var_descr = kernel.temporary_variables[var_name]
elif var_name in kernel.arg_dict:
var_descr = kernel.arg_dict[var_name]
else:
raise NameError("array '%s' was not found" % var_name)
# {{{ check/normalize vary_by_axes
if isinstance(vary_by_axes, str):
vary_by_axes = vary_by_axes.split(",")
from loopy.kernel.array import ArrayBase
if isinstance(var_descr, ArrayBase):
if var_descr.dim_names is not None:
name_to_index = {
name: idx
for idx, name in enumerate(var_descr.dim_names)
}
else:
name_to_index = {}
def map_ax_name_to_index(ax):
if isinstance(ax, str):
try:
return name_to_index[ax]
except KeyError:
raise LoopyError("axis name '%s' not understood " % ax)
else:
return ax
vary_by_axes = [map_ax_name_to_index(ax) for ax in vary_by_axes]
if (vary_by_axes
and (min(vary_by_axes) < 0
or max(vary_by_axes) > var_descr.num_user_axes())):
raise LoopyError("vary_by_axes refers to out-of-bounds axis index")
# }}}
from pymbolic.mapper.substitutor import make_subst_func
from pymbolic.primitives import (Sum, Product, is_zero, flattened_sum,
flattened_product, Subscript, Variable)
from loopy.symbolic import (get_dependencies, SubstitutionMapper,
UnidirectionalUnifier)
# {{{ common factor key list maintenance
# list of (index_key, common factors found)
common_factors = []
def find_unifiable_cf_index(index_key):
for i, (key, _val) in enumerate(common_factors):
unif = UnidirectionalUnifier(
lhs_mapping_candidates=get_dependencies(key))
unif_result = unif(key, index_key)
if unif_result:
assert len(unif_result) == 1
return i, unif_result[0]
return None, None
def extract_index_key(access_expr):
if isinstance(access_expr, Variable):
return ()
elif isinstance(access_expr, Subscript):
index = access_expr.index_tuple
return tuple(index[ax] for ax in vary_by_axes)
else:
raise ValueError("unexpected type of access_expr")
def is_assignee(insn):
return var_name in insn.assignee_var_names()
def iterate_as(cls, expr):
if isinstance(expr, cls):
yield from expr.children
else:
yield expr
# }}}
# {{{ find common factors
from loopy.kernel.data import Assignment
for insn in kernel.instructions:
if not is_assignee(insn):
continue
if not isinstance(insn, Assignment):
raise LoopyError("'%s' modified by non-single-assignment" %
var_name)
lhs = insn.assignee
rhs = insn.expression
if is_zero(rhs):
continue
index_key = extract_index_key(lhs)
cf_index, unif_result = find_unifiable_cf_index(index_key)
if cf_index is None:
# {{{ doesn't exist yet
assert unif_result is None
my_common_factors = None
for term in iterate_as(Sum, rhs):
if term == lhs:
continue
for part in iterate_as(Product, term):
if var_name in get_dependencies(part):
raise LoopyError("unexpected dependency on '%s' "
"in RHS of instruction '%s'" %
(var_name, insn.id))
product_parts = set(iterate_as(Product, term))
if my_common_factors is None:
my_common_factors = product_parts
else:
my_common_factors = my_common_factors & product_parts
if my_common_factors is not None:
common_factors.append((index_key, my_common_factors))
# }}}
else:
# {{{ match, filter existing common factors
_, my_common_factors = common_factors[cf_index]
unif_subst_map = SubstitutionMapper(
make_subst_func(unif_result.lmap))
for term in iterate_as(Sum, rhs):
if term == lhs:
continue
for part in iterate_as(Product, term):
if var_name in get_dependencies(part):
raise LoopyError("unexpected dependency on '%s' "
"in RHS of instruction '%s'" %
(var_name, insn.id))
product_parts = set(iterate_as(Product, term))
my_common_factors = {
cf
for cf in my_common_factors
if unif_subst_map(cf) in product_parts
}
common_factors[cf_index] = (index_key, my_common_factors)
# }}}
# }}}
common_factors = [(ik, cf) for ik, cf in common_factors if cf]
if not common_factors:
raise LoopyError("no common factors found")
# {{{ remove common factors
new_insns = []
for insn in kernel.instructions:
if not isinstance(insn, Assignment) or not is_assignee(insn):
new_insns.append(insn)
continue
index_key = extract_index_key(insn.assignee)
lhs = insn.assignee
rhs = insn.expression
if is_zero(rhs):
new_insns.append(insn)
continue
index_key = extract_index_key(lhs)
cf_index, unif_result = find_unifiable_cf_index(index_key)
if cf_index is None:
new_insns.append(insn)
continue
_, my_common_factors = common_factors[cf_index]
unif_subst_map = SubstitutionMapper(make_subst_func(unif_result.lmap))
mapped_my_common_factors = {
unif_subst_map(cf)
for cf in my_common_factors
}
new_sum_terms = []
for term in iterate_as(Sum, rhs):
if term == lhs:
new_sum_terms.append(term)
continue
new_sum_terms.append(
flattened_product([
part for part in iterate_as(Product, term)
if part not in mapped_my_common_factors
]))
new_insns.append(insn.copy(expression=flattened_sum(new_sum_terms)))
# }}}
# {{{ substitute common factors into usage sites
def find_substitution(expr):
if isinstance(expr, Subscript):
v = expr.aggregate.name
elif isinstance(expr, Variable):
v = expr.name
else:
return expr
if v != var_name:
return expr
index_key = extract_index_key(expr)
cf_index, unif_result = find_unifiable_cf_index(index_key)
unif_subst_map = SubstitutionMapper(make_subst_func(unif_result.lmap))
_, my_common_factors = common_factors[cf_index]
if my_common_factors is not None:
return flattened_product(
[unif_subst_map(cf) for cf in my_common_factors] + [expr])
else:
return expr
insns = new_insns
new_insns = []
subm = SubstitutionMapper(find_substitution)
for insn in insns:
if not isinstance(insn, Assignment) or is_assignee(insn):
new_insns.append(insn)
continue
new_insns.append(insn.with_transformed_expressions(subm))
# }}}
return kernel.copy(instructions=new_insns)
