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 _loop_index(self):
start, end = self.loop_extents
return RuntimeIndex(start, end,
LogicalAnd(
Comparison("<=", Zero((), numpy.int32), start),
Comparison("<=", start, end)),
name="n")
def pack(self, loop_indices=None):
if self.map_ is None:
return None
if hasattr(self, "_pack"):
return self._pack
if self.interior_horizontal:
shape = (2, )
else:
shape = (1, )
shape = shape + self.map_.shape[1:]
if self.view_index is None:
shape = shape + self.outer.shape[1:]
if self.access in {INC, WRITE}:
val = Zero((), self.outer.dtype)
multiindex = MultiIndex(*(Index(e) for e in shape))
self._pack = Materialise(PackInst(), val, multiindex)
elif self.access in {READ, RW, MIN, MAX}:
multiindex = MultiIndex(*(Index(e) for e in shape))
expr, mask = self._rvalue(multiindex, loop_indices=loop_indices)
if mask is not None:
expr = When(mask, expr)
self._pack = Materialise(PackInst(), expr, multiindex)
else:
raise ValueError(
"Don't know how to initialise pack for '%s' access" %
self.access)
return self._pack
def pack(self, loop_indices=None):
if hasattr(self, "_pack"):
return self._pack
shape = self.outer.shape
if self.access is READ:
# No packing required
return self.outer
# We don't need to pack for memory layout, however packing
# globals that are written is required such that subsequent
# vectorisation loop transformations privatise these reduction
# variables. The extra memory movement cost is minimal.
loop_indices = self.pick_loop_indices(*loop_indices)
if self.init_with_zero:
also_zero = {MIN, MAX}
else:
also_zero = set()
if self.access in {INC, WRITE} | also_zero:
val = Zero((), self.outer.dtype)
multiindex = MultiIndex(*(Index(e) for e in shape))
self._pack = Materialise(PackInst(loop_indices), val, multiindex)
elif self.access in {READ, RW, MIN, MAX} - also_zero:
multiindex = MultiIndex(*(Index(e) for e in shape))
expr = Indexed(self.outer, multiindex)
self._pack = Materialise(PackInst(loop_indices), expr, multiindex)
else:
raise ValueError("Don't know how to initialise pack for '%s' access" % self.access)
return self._pack
def emit_pack_instruction(self, *, loop_indices=None):
shape = self.outer.shape
if self.access is WRITE:
zero = Zero((), self.outer.dtype)
multiindex = MultiIndex(*(Index(e) for e in shape))
yield Accumulate(PackInst(), Indexed(self.outer, multiindex), zero)
else:
return ()
def layer_index(self):
if self.extruded:
start, end = self.layer_extents
return RuntimeIndex(start, end,
LogicalAnd(
Comparison("<=", Zero((), numpy.int32), start),
Comparison("<=", start, end)),
name="layer")
else:
return None
def pack(self, loop_indices=None, only_declare=False):
if hasattr(self, "_pack"):
return self._pack
shape = tuple(itertools.chain(*self.shapes))
if only_declare:
pack = Variable(f"matpack{next(self.count)}", shape, self.dtype)
self._pack = pack
if self.access in {WRITE, INC}:
val = Zero((), self.dtype)
multiindex = MultiIndex(*(Index(e) for e in shape))
pack = Materialise(PackInst(), val, multiindex)
self._pack = pack
else:
raise ValueError("Unexpected access type")
return self._pack
def pack(self, loop_indices=None):
if hasattr(self, "_pack"):
return self._pack
((rdim, cdim), ), = self.dims
rmap, cmap = self.maps
if self.interior_horizontal:
shape = (2, )
else:
shape = (1, )
rshape = shape + rmap.shape[1:] + (rdim, )
cshape = shape + cmap.shape[1:] + (cdim, )
if self.access in {WRITE, INC}:
val = Zero((), self.dtype)
multiindex = MultiIndex(*(Index(e) for e in (rshape + cshape)))
pack = Materialise(PackInst(), val, multiindex)
self._pack = pack
return pack
else:
raise ValueError("Unexpected access type")
def pack(self, loop_indices=None):
if hasattr(self, "_pack"):
return self._pack
rshape = 0
cshape = 0
# Need to compute row and col shape based on individual pack shapes
for p in self.packs[:, 0]:
shape, _ = p.shapes
rshape += numpy.prod(shape, dtype=int)
for p in self.packs[0, :]:
_, shape = p.shapes
cshape += numpy.prod(shape, dtype=int)
shape = (rshape, cshape)
if self.access in {WRITE, INC}:
val = Zero((), self.dtype)
multiindex = MultiIndex(*(Index(e) for e in shape))
pack = Materialise(PackInst(), val, multiindex)
self._pack = pack
return pack
else:
raise ValueError("Unexpected access type")