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 _layer_index(self):
if self.constant_layers:
return FixedIndex(0)
if self.subset:
return self._loop_index
else:
return self.loop_index
def bottom_layer(self):
if self.iteration_region == ON_TOP:
return Materialise(PackInst(),
Indexed(self._layers_array, (self._layer_index, FixedIndex(0))),
MultiIndex())
else:
start, _ = self.layer_extents
return start
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 _rvalue(self, multiindex, loop_indices=None):
"""Returns indexed Dat and masking condition to apply to reads/writes.
If the masking condition is None, no mask is applied,
otherwise the pack/unpack will be wrapped in When(mask, expr).
This is used for the case where maps might have negative entries.
"""
f, i, *j = multiindex
n, layer = self.pick_loop_indices(*loop_indices)
if self.view_index is not None:
j = tuple(j) + tuple(FixedIndex(i) for i in self.view_index)
map_, (f, i) = self.map_.indexed((n, i, f), layer=layer)
return Indexed(self.outer, MultiIndex(map_, *j)), self._mask(map_)
def kernel_arg(self, loop_indices=None):
if self.map_ is None:
if loop_indices is None:
raise ValueError("Need iteration index")
n, layer = self.pick_loop_indices(*loop_indices)
shape = self.outer.shape
if self.view_index is None:
multiindex = (n, ) + tuple(Index(e) for e in shape[1:])
else:
multiindex = (n, ) + tuple(FixedIndex(i) for i in self.view_index)
return Indexed(self.outer, multiindex)
else:
pack = self.pack(loop_indices)
shape = pack.shape
return Indexed(pack, (Index(e) for e in shape))
def replace_indices_index(node, self):
if node.extent == 1:
return FixedIndex(0)
return self.subst.get(node, node)
def _layer_index(self):
if self.constant_layers:
return FixedIndex(0)
else:
return self.loop_index