def add_argument(self, arg):
interior_horizontal = self.iteration_region == ON_INTERIOR_FACETS
if arg._is_dat:
if arg._is_mixed:
packs = []
for a in arg:
shape = (None, *a.data.shape[1:])
argument = Argument(shape, a.data.dtype, pfx="mdat")
packs.append(
a.data.pack(argument,
arg.access,
self.map_(a.map, unroll=a.unroll_map),
interior_horizontal=interior_horizontal))
self.arguments.append(argument)
pack = MixedDatPack(packs,
arg.access,
arg.dtype,
interior_horizontal=interior_horizontal)
self.packed_args.append(pack)
self.argument_accesses.append(arg.access)
return
if arg._is_dat_view:
view_index = arg.data.index
data = arg.data._parent
else:
view_index = None
data = arg.data
shape = (None, *data.shape[1:])
argument = Argument(shape, arg.data.dtype, pfx="dat")
pack = arg.data.pack(argument,
arg.access,
self.map_(arg.map, unroll=arg.unroll_map),
interior_horizontal=interior_horizontal,
view_index=view_index)
elif arg._is_global:
argument = Argument(arg.data.dim, arg.data.dtype, pfx="glob")
pack = GlobalPack(argument, arg.access)
elif arg._is_mat:
argument = Argument((), PetscMat(), pfx="mat")
map_ = tuple(self.map_(m, unroll=arg.unroll_map) for m in arg.map)
pack = arg.data.pack(argument,
arg.access,
map_,
arg.data.dims,
arg.data.dtype,
interior_horizontal=interior_horizontal)
else:
raise ValueError("Unhandled argument type")
self.arguments.append(argument)
self.packed_args.append(pack)
self.argument_accesses.append(arg.access)
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:
offset = NamedLiteral(offset, name=values.name + "_offset")
self.values = values
self.offset = offset
def __init__(self, *, kernel, iterset, iteration_region=None, single_cell=False,
pass_layer_to_kernel=False, forward_arg_types=()):
self.kernel = kernel
self.arguments = []
self.argument_accesses = []
self.packed_args = []
self.indices = []
self.maps = OrderedDict()
self.iterset = iterset
if iteration_region is None:
self.iteration_region = ALL
else:
self.iteration_region = iteration_region
self.pass_layer_to_kernel = pass_layer_to_kernel
self.single_cell = single_cell
self.forward_arguments = tuple(Argument((), fa, pfx="farg") for fa in forward_arg_types)
def _rename_node_argument(node, self):
name = self.replace.get(node, node.name)
return Argument(node.shape, node.dtype, name=name)
def _layers_array(self):
if self.constant_layers:
return Argument((1, 2), IntType, name="layers")
else:
return Argument((None, 2), IntType, name="layers")
def _subset_indices(self):
return Argument(("end", ), IntType, name="subset_indices")
def loop_extents(self):
return (Argument((), IntType,
name="start"), Argument((), IntType, name="end"))