def emit_unpack_instruction(self, *, loop_indices=None):
pack = self.pack(loop_indices)
if self.access is READ:
return ()
else:
if self.interior_horizontal:
_shape = (2,)
else:
_shape = (1,)
offset = 0
for p in self.packs:
shape = _shape + p.map_.shape[1:] + p.outer.shape[1:]
mi = MultiIndex(*(Index(e) for e in shape))
rvalue, mask = p._rvalue(mi, loop_indices)
extents = [numpy.prod(shape[i+1:], dtype=numpy.int32) for i in range(len(shape))]
index = reduce(Sum, [Product(i, Literal(IntType.type(e), casting=False)) for i, e in zip(mi, extents)], Literal(IntType.type(0), casting=False))
indices = MultiIndex(Sum(index, Literal(IntType.type(offset), casting=False)),)
rhs = Indexed(pack, indices)
offset += numpy.prod(shape, dtype=numpy.int32)
if self.access in {INC, MIN, MAX}:
op = {INC: Sum,
MIN: Min,
MAX: Max}[self.access]
rhs = op(rvalue, rhs)
acc = Accumulate(UnpackInst(), rvalue, rhs)
if mask is None:
yield acc
else:
yield When(mask, acc)
def layer_extents(self):
if self.iteration_region == ON_BOTTOM:
start = Indexed(self._layers_array,
(self._layer_index, FixedIndex(0)))
end = Sum(
Indexed(self._layers_array,
(self._layer_index, FixedIndex(0))),
Literal(IntType.type(1)))
elif self.iteration_region == ON_TOP:
start = Sum(
Indexed(self._layers_array,
(self._layer_index, FixedIndex(1))),
Literal(IntType.type(-2)))
end = Sum(
Indexed(self._layers_array,
(self._layer_index, FixedIndex(1))),
Literal(IntType.type(-1)))
elif self.iteration_region == ON_INTERIOR_FACETS:
start = Indexed(self._layers_array,
(self._layer_index, FixedIndex(0)))
end = Sum(
Indexed(self._layers_array,
(self._layer_index, FixedIndex(1))),
Literal(IntType.type(-2)))
elif self.iteration_region == ALL:
start = Indexed(self._layers_array,
(self._layer_index, FixedIndex(0)))
end = Sum(
Indexed(self._layers_array,
(self._layer_index, FixedIndex(1))),
Literal(IntType.type(-1)))
else:
raise ValueError("Unknown iteration region")
return (Materialise(PackInst(), start, MultiIndex()),
Materialise(PackInst(), end, MultiIndex()))
def pack(self, loop_indices=None):
if hasattr(self, "_pack"):
return self._pack
flat_shape = numpy.sum(
tuple(
numpy.prod(p.map_.shape[1:] + p.outer.shape[1:])
for p in self.packs))
if self.interior_horizontal:
_shape = (2, )
flat_shape *= 2
else:
_shape = (1, )
if self.access in {INC, WRITE}:
val = Zero((), self.dtype)
multiindex = MultiIndex(Index(flat_shape))
self._pack = Materialise(PackInst(), val, multiindex)
elif self.access in {READ, RW, MIN, MAX}:
multiindex = MultiIndex(Index(flat_shape))
val = Zero((), self.dtype)
expressions = []
offset = 0
for p in self.packs:
shape = _shape + p.map_.shape[1:] + p.outer.shape[1:]
mi = MultiIndex(*(Index(e) for e in shape))
expr, mask = p._rvalue(mi, loop_indices)
extents = [
numpy.prod(shape[i + 1:], dtype=numpy.int32)
for i in range(len(shape))
]
index = reduce(Sum, [
Product(i, Literal(IntType.type(e), casting=False))
for i, e in zip(mi, extents)
], Literal(IntType.type(0), casting=False))
indices = MultiIndex(
Sum(index, Literal(IntType.type(offset), casting=False)), )
offset += numpy.prod(shape, dtype=numpy.int32)
if mask is not None:
expr = When(mask, expr)
expressions.append(expr)
expressions.append(indices)
self._pack = Materialise(PackInst(), val, multiindex, *expressions)
else:
raise ValueError(
"Don't know how to initialise pack for '%s' access" %
self.access)
return self._pack
def indexed(self, multiindex, layer=None):
n, i, f = multiindex
if layer is not None and self.offset is not None:
# For extruded mesh, prefetch the indirections for each map, so that they don't
# need to be recomputed. Different f values need to be treated separately.
key = f.extent
if key is None:
key = 1
if key not in self.prefetch:
bottom_layer, _ = self.layer_bounds
offset_extent, = self.offset.shape
j = Index(offset_extent)
base = Indexed(self.values, (n, j))
if f.extent:
k = Index(f.extent)
else:
k = Index(1)
offset = Sum(
Sum(layer, Product(Literal(numpy.int32(-1)),
bottom_layer)), k)
offset = Product(offset, Indexed(self.offset, (j, )))
self.prefetch[key] = Materialise(PackInst(), Sum(base, offset),
MultiIndex(k, j))
return Indexed(self.prefetch[key], (f, i)), (f, i)
else:
assert f.extent == 1 or f.extent is None
base = Indexed(self.values, (n, i))
return base, (f, i)
def __init__(self, map_, interior_horizontal, layer_bounds,
values=None, offset=None, unroll=False):
self.variable = map_.iterset._extruded and not map_.iterset.constant_layers
self.unroll = unroll
self.layer_bounds = layer_bounds
self.interior_horizontal = interior_horizontal
self.prefetch = {}
if values is not None:
raise RuntimeError
self.values = values
if map_.offset is not None:
assert offset is not None
self.offset = offset
return
offset = map_.offset
shape = (None, ) + map_.shape[1:]
values = Argument(shape, dtype=map_.dtype, pfx="map")
if offset is not None:
if len(set(map_.offset)) == 1:
offset = Literal(offset[0], casting=True)
else:
offset = NamedLiteral(offset, name=values.name + "_offset")
self.values = values
self.offset = offset
def top_layer(self):
if self.iteration_region == ON_BOTTOM:
return Materialise(PackInst(),
Sum(Indexed(self._layers_array, (self._layer_index, FixedIndex(1))),
Literal(IntType.type(-1))),
MultiIndex())
else:
_, end = self.layer_extents
return end
def indexed_vector(self, n, shape, layer=None):
shape = self.shape[1:] + shape
if self.interior_horizontal:
shape = (2, ) + shape
else:
shape = (1, ) + shape
f, i, j = (Index(e) for e in shape)
base, (f, i) = self.indexed((n, i, f), layer=layer)
init = Sum(Product(base, Literal(numpy.int32(j.extent))), j)
pack = Materialise(PackInst(), init, MultiIndex(f, i, j))
multiindex = tuple(Index(e) for e in pack.shape)
return Indexed(pack, multiindex), multiindex
def emit_unpack_instruction(self, *,
loop_indices=None):
pack = self.pack(loop_indices=loop_indices)
mixed_to_local = []
local_to_global = []
roffset = 0
for row in self.packs:
coffset = 0
for p in row:
rshape, cshape = p.shapes
pack_ = p.pack(loop_indices=loop_indices, only_declare=True)
rindices = tuple(Index(e) for e in rshape)
cindices = tuple(Index(e) for e in cshape)
indices = MultiIndex(*rindices, *cindices)
lvalue = Indexed(pack_, indices)
rextents = [numpy.prod(rshape[i+1:], dtype=numpy.int32) for i in range(len(rshape))]
cextents = [numpy.prod(cshape[i+1:], dtype=numpy.int32) for i in range(len(cshape))]
flat_row_index = reduce(Sum, [Product(i, Literal(IntType.type(e), casting=False))
for i, e in zip(rindices, rextents)],
Literal(IntType.type(0), casting=False))
flat_col_index = reduce(Sum, [Product(i, Literal(IntType.type(e), casting=False))
for i, e in zip(cindices, cextents)],
Literal(IntType.type(0), casting=False))
flat_index = MultiIndex(Sum(flat_row_index, Literal(IntType.type(roffset), casting=False)),
Sum(flat_col_index, Literal(IntType.type(coffset), casting=False)))
rvalue = Indexed(pack, flat_index)
# Copy from local mixed element tensor into non-mixed
mixed_to_local.append(Accumulate(PreUnpackInst(), lvalue, rvalue))
# And into global matrix.
local_to_global.extend(p.emit_unpack_instruction(loop_indices=loop_indices))
coffset += numpy.prod(cshape, dtype=numpy.int32)
roffset += numpy.prod(rshape, dtype=numpy.int32)
yield from iter(mixed_to_local)
yield from iter(local_to_global)
def indexed(self, multiindex, layer=None):
n, i, f = multiindex
if layer is not None and self.offset is not None:
# For extruded mesh, prefetch the indirections for each map, so that they don't
# need to be recomputed.
# First prefetch the base map (not dependent on layers)
base_key = None
if base_key not in self.prefetch:
j = Index()
base = Indexed(self.values, (n, j))
self.prefetch[base_key] = Materialise(PackInst(), base, MultiIndex(j))
base = self.prefetch[base_key]
# Now prefetch the extruded part of the map (inside the layer loop).
# This is necessary so loopy DTRT for MatSetValues
# Different f values need to be treated separately.
key = f.extent
if key is None:
key = 1
if key not in self.prefetch:
bottom_layer, _ = self.layer_bounds
k = Index(f.extent if f.extent is not None else 1)
offset = Sum(Sum(layer, Product(Literal(numpy.int32(-1)), bottom_layer)), k)
j = Index()
# Inline map offsets where all entries are identical.
if self.offset.shape == ():
offset = Product(offset, self.offset)
else:
offset = Product(offset, Indexed(self.offset, (j,)))
base = Indexed(base, (j, ))
self.prefetch[key] = Materialise(PackInst(), Sum(base, offset), MultiIndex(k, j))
return Indexed(self.prefetch[key], (f, i)), (f, i)
else:
assert f.extent == 1 or f.extent is None
base = Indexed(self.values, (n, i))
return base, (f, i)
def _mask(self, map_):
if self.needs_mask:
return Comparison(">=", map_, Literal(numpy.int32(0)))
else:
return None