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