def __setstate__(self, state): yk_solns = state.pop('yk_solns') super(OperatorYASK, self).__setstate__(state) # Restore the YASK solutions (see __getstate__ for more info) self.yk_solns = OrderedDict() for (dimensions, yk_soln_obj, name, soname, libfile, pyfile, pysofile, local_grids) in yk_solns: path = Path(namespace['yask-lib'], 'lib%s.so' % soname) if not path.is_file(): with open(path, 'wb') as f: f.write(libfile) path = Path(namespace['yask-pylib'], '%s.py' % name) if not path.is_file(): with open(path, 'w') as f: f.write(pyfile.decode()) path = Path(namespace['yask-pylib'], '_%s.so' % name) if not path.is_file(): with open(path, 'wb') as f: f.write(pysofile) # Finally reinstantiate the YASK solution -- no code generation or JIT # will happen at this point, as all necessary files have been restored context = contexts.fetch(dimensions, self._dtype) local_grids = [i for i in self.parameters if i.name in local_grids] yk_soln = context.make_yk_solution(name, None, local_grids) self.yk_solns[(dimensions, yk_soln_obj)] = yk_soln
def wrapper(self): if self._data is None: log("Allocating memory for %s%s" % (self.name, self.shape_allocated)) # Fetch the appropriate context context = contexts.fetch(self.dimensions, self.dtype) # Create a YASK grid; this allocates memory grid = context.make_grid(self) # /self._padding/ must be updated as (from the YASK docs): # "The value may be slightly larger [...] due to rounding" padding = [] for i in self.dimensions: if i.is_Space: padding.append((grid.get_left_extra_pad_size(i.name), grid.get_right_extra_pad_size(i.name))) else: # time and misc dimensions padding.append((0, 0)) self._padding = tuple(padding) self._data = Data(grid, self.shape_allocated, self.indices, self.dtype) self._data.reset() return func(self)
def wrapper(self): if self._data is None: log("Allocating memory for %s%s" % (self.name, self.shape_allocated)) # Fetch the appropriate context context = contexts.fetch(self.dimensions, self.dtype) # Create a YASK grid; this allocates memory grid = context.make_grid(self) # `self._padding` must be updated as (from the YASK docs): # "The value may be slightly larger [...] due to rounding" padding = [] for i in self.dimensions: if i.is_Space: padding.append((grid.get_left_extra_pad_size(i.name), grid.get_right_extra_pad_size(i.name))) else: # time and misc dimensions padding.append((0, 0)) self._padding = tuple(padding) del self.shape_allocated # Invalidate cached_property self._data = Data(grid, self.shape_allocated, self.indices, self.dtype) self._data.reset() return func(self)
def wrapper(self): if self._data is None: log("Allocating memory for %s%s" % (self.name, self.shape_allocated)) # Free memory carried by stale symbolic objects # TODO: see issue #944 # CacheManager.clear(dump_contexts=False, force=False) # Fetch the appropriate context context = contexts.fetch(self.dimensions, self.dtype) # Create a YASK var; this allocates memory var = context.make_var(self) # `self._padding` must be updated as (from the YASK docs): # "The value may be slightly larger [...] due to rounding" padding = [] for i in self.dimensions: if i.is_Space: padding.append((var.get_left_extra_pad_size(i.name), var.get_right_extra_pad_size(i.name))) else: # time and misc dimensions padding.append((0, 0)) self._padding = tuple(padding) del self.shape_allocated # Invalidate cached_property self._data = Data(var, self.shape_allocated, self.indices, self.dtype) self._data.reset() return func(self)
def _specialize_iet(self, nodes): """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.""" log("Specializing a Devito Operator for YASK...") self.context = YaskNullContext() self.yk_soln = YaskNullKernel() offloadable = find_offloadable_trees(nodes) if len(offloadable) == 0: log("No offloadable trees found") elif len(offloadable) == 1: tree, grid, dtype = offloadable[0] self.context = contexts.fetch(grid, dtype) # Create a YASK compiler solution for this Operator yc_soln = self.context.make_yc_solution(namespace['jit-yc-soln']) transform = sympy2yask(self.context, yc_soln) try: for i in tree[-1].nodes: transform(i.expr) funcall = make_sharedptr_funcall(namespace['code-soln-run'], ['time'], namespace['code-soln-name']) funcall = Element(c.Statement(ccode(funcall))) nodes = Transformer({tree[1]: funcall}).visit(nodes) # Track /funcall/ as an external function call self.func_table[namespace['code-soln-run']] = MetaCall( None, False) # JIT-compile the newly-created YASK kernel local_grids = [i for i in transform.mapper if i.is_Array] self.yk_soln = self.context.make_yk_solution( namespace['jit-yk-soln'], yc_soln, local_grids) # 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: log("Unable to offload a candidate tree.") else: exit("Found more than one offloadable trees in a single Operator") # 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 nodes = make_grid_accesses(nodes) log("Specialization successfully performed!") return nodes
def _allocate_memory(self): """Allocate memory in terms of YASK grids.""" # TODO: YASK assumes that self.shape == "domain_size" (in YASK jargon) # This is exactly how Devito will work too once the Grid abstraction lands # Fetch the appropriate context context = contexts.fetch(self.grid, self.dtype) # TODO : the following will fail if not using a SteppingDimension, # eg with save=True one gets /time/ instead /t/ self._data_object = context.make_grid(self)
def wrapper(self): if self._data is None: log("Allocating memory for %s%s" % (self.name, self.shape_allocated)) # Fetch the appropriate context context = contexts.fetch(self.grid, self.dtype) # TODO : the following will fail if not using a SteppingDimension, # eg with save=True one gets /time/ instead /t/ grid = context.make_grid(self) # /self._padding/ must be updated as (from the YASK docs): # "The value may be slightly larger [...] due to rounding" pad = [(0, 0) if i.is_Time else (grid.get_left_extra_pad_size(i.name), grid.get_right_extra_pad_size(i.name)) for i in self.indices] self._padding = pad self._data = Data(grid, self.shape_allocated, self.indices, self.dtype) self._data.reset() return func(self)
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
def _specialize(self, nodes, parameters): """ Create a YASK representation of this Iteration/Expression tree. ``parameters`` is modified in-place adding YASK-related arguments. """ log("Specializing a Devito Operator for YASK...") self.context = YaskNullContext() self.yk_soln = YaskNullKernel() local_grids = [] offloadable = find_offloadable_trees(nodes) if len(offloadable) == 0: log("No offloadable trees found") elif len(offloadable) == 1: tree, grid, dtype = offloadable[0] self.context = contexts.fetch(grid, dtype) # Create a YASK compiler solution for this Operator yc_soln = self.context.make_yc_solution(namespace['jit-yc-soln']) transform = sympy2yask(self.context, yc_soln) try: for i in tree[-1].nodes: transform(i.expr) funcall = make_sharedptr_funcall(namespace['code-soln-run'], ['time'], namespace['code-soln-name']) funcall = Element(c.Statement(ccode(funcall))) nodes = Transformer({tree[1]: funcall}).visit(nodes) # Track /funcall/ as an external function call self.func_table[namespace['code-soln-run']] = MetaCall( None, False) # JIT-compile the newly-created YASK kernel local_grids += [i for i in transform.mapper if i.is_Array] self.yk_soln = self.context.make_yk_solution( namespace['jit-yk-soln'], yc_soln, local_grids) # Now we must drop a pointer to the YASK solution down to C-land parameters.append( Object(namespace['code-soln-name'], namespace['type-solution'], self.yk_soln.rawpointer)) # 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: log("Unable to offload a candidate tree.") else: exit("Found more than one offloadable trees in a single Operator") # 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 nodes = make_grid_accesses(nodes) # Update the parameters list adding all necessary YASK grids for i in list(parameters) + local_grids: try: if i.from_YASK: parameters.append( Object(namespace['code-grid-name'](i.name), namespace['type-grid'])) except AttributeError: # Ignore e.g. Dimensions pass log("Specialization successfully performed!") return nodes
def _specialize(self, nodes, parameters): """ Create a YASK representation of this Iteration/Expression tree. ``parameters`` is modified in-place adding YASK-related arguments. """ log("Specializing a Devito Operator for YASK...") # Find offloadable Iteration/Expression trees offloadable = [] for tree in retrieve_iteration_tree(nodes): parallel = filter_iterations(tree, lambda i: i.is_Parallel) if not parallel: # Cannot offload non-parallel loops continue if not (IsPerfectIteration().visit(tree) and all(i.is_Expression for i in tree[-1].nodes)): # Don't know how to offload this Iteration/Expression to YASK continue functions = flatten(i.functions for i in tree[-1].nodes) keys = set((i.indices, i.shape, i.dtype) for i in functions if i.is_TimeData) if len(keys) == 0: continue elif len(keys) > 1: exit("Cannot handle Operators w/ heterogeneous grids") dimensions, shape, dtype = keys.pop() if len(dimensions) == len(tree) and\ all(i.dim == j for i, j in zip(tree, dimensions)): # Detected a "full" Iteration/Expression tree (over both # time and space dimensions) offloadable.append((tree, dimensions, shape, dtype)) # Construct YASK ASTs given Devito expressions. New grids may be allocated. if len(offloadable) == 0: # No offloadable trees found self.context = YaskNullContext() self.yk_soln = YaskNullSolution() processed = nodes log("No offloadable trees found") elif len(offloadable) == 1: # Found *the* offloadable tree for this Operator tree, dimensions, shape, dtype = offloadable[0] self.context = contexts.fetch(dimensions, shape, dtype) # Create a YASK compiler solution for this Operator # Note: this can be dropped as soon as the kernel has been built yc_soln = self.context.make_yc_solution(namespace['jit-yc-soln']) transform = sympy2yask(self.context, yc_soln) try: for i in tree[-1].nodes: transform(i.expr) funcall = make_sharedptr_funcall(namespace['code-soln-run'], ['time'], namespace['code-soln-name']) funcall = Element(c.Statement(ccode(funcall))) processed = Transformer({tree[1]: funcall}).visit(nodes) # Track this is an external function call self.func_table[namespace['code-soln-run']] = FunMeta( None, False) # JIT-compile the newly-created YASK kernel self.yk_soln = self.context.make_yk_solution( namespace['jit-yk-soln'], yc_soln) # Now we must drop a pointer to the YASK solution down to C-land parameters.append( Object(namespace['code-soln-name'], namespace['type-solution'], self.yk_soln.rawpointer)) # 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: self.yk_soln = YaskNullSolution() processed = nodes log("Unable to offload a candidate tree.") else: exit("Found more than one offloadable trees in a single Operator") # 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 processed = make_grid_accesses(processed) # Update the parameters list adding all necessary YASK grids for i in list(parameters): try: if i.from_YASK: parameters.append( Object(namespace['code-grid-name'](i.name), namespace['type-grid'], i.data.rawpointer)) except AttributeError: # Ignore e.g. Dimensions pass log("Specialization successfully performed!") return processed
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. """ offloadable = find_offloadable_trees(iet) if len(offloadable.trees) == 0: self.yk_soln = YaskNullKernel() log("No offloadable trees found") else: context = contexts.fetch(offloadable.grid, offloadable.dtype) # A unique name for the 'real' compiler and kernel solutions name = namespace['jit-soln'](Signer._digest(iet, configuration)) # Create a YASK compiler solution for this Operator yc_soln = context.make_yc_solution(name) try: trees = offloadable.trees # Generate YASK grids and populate `yc_soln` with equations mapper = yaskizer(trees, yc_soln) local_grids = [i for i in mapper if i.is_Array] # Transform the IET funcall = make_sharedptr_funcall(namespace['code-soln-run'], ['time'], namespace['code-soln-name']) funcall = Element(c.Statement(ccode(funcall))) mapper = {trees[0].root: funcall} mapper.update({i.root: mapper.get(i.root) for i in trees}) # Drop trees iet = Transformer(mapper).visit(iet) # Mark `funcall` as an external function call self.func_table[namespace['code-soln-run']] = MetaCall( None, False) # JIT-compile the newly-created YASK kernel self.yk_soln = context.make_yk_solution( name, yc_soln, local_grids) # 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: self.yk_soln = YaskNullKernel() log("Unable to offload a candidate tree. Reason: [%s]" % str(e)) # 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 iet = make_grid_accesses(iet) # Finally optimize all non-yaskized loops iet = super(Operator, self)._specialize_iet(iet, **kwargs) return iet
def make_yask_kernels(iet, **kwargs): yk_solns = kwargs.pop('yk_solns') mapper = {} for n, (section, trees) in enumerate(find_affine_trees(iet).items()): dimensions = tuple(filter_ordered(i.dim.root for i in flatten(trees))) # Retrieve the section dtype exprs = FindNodes(Expression).visit(section) dtypes = {e.dtype for e in exprs} if len(dtypes) != 1: log("Unable to offload in presence of mixed-precision arithmetic") continue dtype = dtypes.pop() context = contexts.fetch(dimensions, 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 vars and populate `yc_soln` with equations local_vars = 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, dtype) mapper[trees[0].root] = funcall mapper.update({i.root: mapper.get(i.root) for i in trees}) # Drop trees # JIT-compile the newly-created YASK kernel yk_soln = context.make_yk_solution(name, yc_soln, local_vars) yk_solns[(dimensions, yk_soln_obj)] = yk_soln # Print some useful information about the newly constructed solution log("Solution '%s' contains %d var(s) and %d equation(s)." % (yc_soln.get_name(), yc_soln.get_num_vars(), 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 yk_solns: log("No offloadable trees found") # Some Iteration/Expression trees are not offloaded to YASK and may # require further processing to be executed through YASK, due to the # different storage layout yk_var_objs = { i.name: YASKVarObject(i.name) for i in FindSymbols().visit(iet) if i.from_YASK } yk_var_objs.update({i: YASKVarObject(i) for i in get_local_vars(yk_solns)}) iet = make_var_accesses(iet, yk_var_objs) # The signature needs to be updated # TODO: this could be done automagically through the iet pass engine, but # currently it only supports *appending* to the parameters list. While here # we actually need to change it as some parameters may disappear (x_m, x_M, ...) parameters = derive_parameters(iet, True) iet = iet._rebuild(parameters=parameters) return 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
def _specialize_iet(cls, 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 = {} yk_solns = kwargs.pop('yk_solns') for n, (section, trees) in enumerate(find_affine_trees(iet).items()): dimensions = tuple( filter_ordered(i.dim.root for i in flatten(trees))) # Retrieve the section dtype exprs = FindNodes(Expression).visit(section) dtypes = {e.dtype for e in exprs} if len(dtypes) != 1: log("Unable to offload in presence of mixed-precision arithmetic" ) continue dtype = dtypes.pop() context = contexts.fetch(dimensions, 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 vars and populate `yc_soln` with equations local_vars = 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, dtype) mapper[trees[0].root] = funcall mapper.update({i.root: mapper.get(i.root) for i in trees}) # Drop trees # JIT-compile the newly-created YASK kernel yk_soln = context.make_yk_solution(name, yc_soln, local_vars) yk_solns[(dimensions, yk_soln_obj)] = yk_soln # Print some useful information about the newly constructed solution log("Solution '%s' contains %d var(s) and %d equation(s)." % (yc_soln.get_name(), yc_soln.get_num_vars(), 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 yk_solns: log("No offloadable trees found") # Some Iteration/Expression trees are not offloaded to YASK and may # require further processing to be executed through YASK, due to the # different storage layout yk_var_objs = { i.name: YaskVarObject(i.name) for i in FindSymbols().visit(iet) if i.from_YASK } yk_var_objs.update( {i: YaskVarObject(i) for i in cls._get_local_vars(yk_solns)}) iet = make_var_accesses(iet, yk_var_objs) # The signature needs to be updated parameters = derive_parameters(iet, True) iet = iet._rebuild(parameters=parameters) return super(OperatorYASK, cls)._specialize_iet(iet, **kwargs)