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"))