def _build_casts(self, iet): iet = super(Operator, self)._build_casts(iet) # Add YASK solution pointer for use in C-land soln_obj = Object(namespace['code-soln-name'], namespace['type-solution']) # Add YASK user and local grids pointers for use in C-land grid_objs = [YaskGridObject(i.name) for i in self.input if i.from_YASK] grid_objs.extend([YaskGridObject(i) for i in self.yk_soln.local_grids]) # Build pointer casts casts = [PointerCast(soln_obj)] + [PointerCast(i) for i in grid_objs] return List(body=casts + [iet])
def _build_casts(self, iet): iet = super(Operator, self)._build_casts(iet) # Add YASK solution pointer for use in C-land soln_objs = [YaskSolnObject(cname) for _, cname in self.yk_solns] # Add YASK user and local grids pointers for use in C-land grid_objs = [YaskGridObject(i.name) for i in self.input if i.from_YASK] grid_objs.extend([YaskGridObject(i) for i in self._local_grids]) # Build pointer casts casts = [PointerCast(i) for i in soln_objs] + [PointerCast(i) for i in grid_objs] return List(body=casts + [iet])
def _specialize_iet(self, iet, **kwargs): """ Transform the Iteration/Expression tree to offload the computation of one or more loop nests onto YASK. This involves calling the YASK compiler to generate YASK code. Such YASK code is then called from within the transformed Iteration/Expression tree. """ mapper = {} self.yk_solns = OrderedDict() for n, (section, trees) in enumerate(find_affine_trees(iet).items()): dimensions = tuple( filter_ordered(i.dim.root for i in flatten(trees))) context = contexts.fetch(dimensions, self._dtype) # A unique name for the 'real' compiler and kernel solutions name = namespace['jit-soln'](Signer._digest( configuration, *[i.root for i in trees])) # Create a YASK compiler solution for this Operator yc_soln = context.make_yc_solution(name) try: # Generate YASK grids and populate `yc_soln` with equations local_grids = yaskit(trees, yc_soln) # Build the new IET nodes yk_soln_obj = YaskSolnObject(namespace['code-soln-name'](n)) funcall = make_sharedptr_funcall(namespace['code-soln-run'], ['time'], yk_soln_obj) funcall = Offloaded(funcall, self._dtype) mapper[trees[0].root] = funcall mapper.update({i.root: mapper.get(i.root) for i in trees}) # Drop trees # Mark `funcall` as an external function call self._func_table[namespace['code-soln-run']] = MetaCall( None, False) # JIT-compile the newly-created YASK kernel yk_soln = context.make_yk_solution(name, yc_soln, local_grids) self.yk_solns[(dimensions, yk_soln_obj)] = yk_soln # Print some useful information about the newly constructed solution log("Solution '%s' contains %d grid(s) and %d equation(s)." % (yc_soln.get_name(), yc_soln.get_num_grids(), yc_soln.get_num_equations())) except NotImplementedError as e: log("Unable to offload a candidate tree. Reason: [%s]" % str(e)) iet = Transformer(mapper).visit(iet) if not self.yk_solns: log("No offloadable trees found") # Some Iteration/Expression trees are not offloaded to YASK and may # require further processing to be executed in YASK, due to the differences # in storage layout employed by Devito and YASK yk_grid_objs = { i.name: YaskGridObject(i.name) for i in self._input if i.from_YASK } yk_grid_objs.update({i: YaskGridObject(i) for i in self._local_grids}) iet = make_grid_accesses(iet, yk_grid_objs) # Finally optimize all non-yaskized loops iet = super(OperatorYASK, self)._specialize_iet(iet, **kwargs) return iet