def emit_unpack_instruction(self, *, loop_indices=None):
pack = self.pack(loop_indices)
if pack is None:
return ()
elif self.access is READ:
return ()
elif self.access in {INC, MIN, MAX}:
op = {INC: Sum,
MIN: Min,
MAX: Max}[self.access]
multiindex = tuple(Index(e) for e in pack.shape)
rvalue, mask = self._rvalue(multiindex, loop_indices=loop_indices)
acc = Accumulate(UnpackInst(), rvalue, op(rvalue, Indexed(pack, multiindex)))
if mask is None:
yield acc
else:
yield When(mask, acc)
else:
multiindex = tuple(Index(e) for e in pack.shape)
rvalue, mask = self._rvalue(multiindex, loop_indices=loop_indices)
acc = Accumulate(UnpackInst(), rvalue, Indexed(pack, multiindex))
if mask is None:
yield acc
else:
yield When(mask, acc)
def replace_materialise_materialise(node, self):
v = Variable(node.name, node.shape, node.dtype)
inits = list(map(self, node.children))
label = node.label
accs = []
for rvalue, indices in zip(*(inits[0::2], inits[1::2])):
lvalue = Indexed(v, indices)
if isinstance(rvalue, When):
when, rvalue = rvalue.children
acc = When(when, Accumulate(label, lvalue, rvalue))
else:
acc = Accumulate(label, lvalue, rvalue)
accs.append(acc)
self.initialisers.append(tuple(accs))
return v
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 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 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 emit_unpack_instruction(self, *, loop_indices=None):
pack = self.pack(loop_indices)
loop_indices = self.pick_loop_indices(*loop_indices)
if pack is None:
return ()
elif self.access is READ:
return ()
elif self.access in {INC, MIN, MAX}:
op = {INC: Sum,
MIN: Min,
MAX: Max}[self.access]
multiindex = tuple(Index(e) for e in pack.shape)
rvalue = Indexed(self.outer, multiindex)
yield Accumulate(UnpackInst(loop_indices), rvalue, op(rvalue, Indexed(pack, multiindex)))
else:
multiindex = tuple(Index(e) for e in pack.shape)
rvalue = Indexed(self.outer, multiindex)
yield Accumulate(UnpackInst(loop_indices), rvalue, Indexed(pack, multiindex))