def make_grid(self, obj): """ Create and return a new :class:`Data`, a YASK grid wrapper. Memory is allocated. :param obj: The :class:`Function` for which a YASK grid is allocated. """ if set(obj.indices) < set(self.space_dimensions): exit("Need a Function[x,y,z] to create a YASK grid.") name = 'devito_%s_%d' % (obj.name, contexts.ngrids) # Set up the YASK grid and allocate memory grid = self.yk_hook.new_grid(name, obj) for i, s, h in zip(obj.indices, obj.shape, obj._halo): if i.is_Time: assert grid.is_dim_used(i.name) assert grid.get_alloc_size(i.name) == s else: # Note: # 1) The halo is set to a value which is the max between the number # of points on the left and the number of points on the right of # the approximation (the same with a centered approximation) # 2) from the YASK docs: "If the halo is set to a value larger than # the padding size, the padding size will be automatically increased # to accomodate it grid.set_halo_size(i.name, max(h)) grid.alloc_storage() self.grids[name] = grid return grid
def find_offloadable_trees(nodes): """ Return the trees within ``nodes`` that can be computed by YASK. A tree is "offloadable to YASK" if it is embedded in a time stepping loop *and* all of the grids accessed by the enclosed equations are homogeneous (i.e., same dimensions and data type). """ 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.grid, i.dtype) for i in functions if i.is_TimeFunction) if len(keys) == 0: continue elif len(keys) > 1: exit("Cannot handle Operators w/ heterogeneous grids") grid, dtype = keys.pop() # Is this a "complete" tree iterating over the entire grid? dims = [i.dim for i in tree] if all(i in dims for i in grid.dimensions) and\ any(j in dims for j in [grid.time_dim, grid.stepping_dim]): offloadable.append((tree, grid, dtype)) return offloadable
def find_offloadable_trees(nodes): """ Return the trees within ``nodes`` that can be computed by YASK. A tree is "offloadable to YASK" if it is embedded in a time stepping loop *and* all of the grids accessed by the enclosed equations are homogeneous (i.e., same dimensions, shape, data type). """ 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_TimeFunction) 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)) return offloadable
def fetch(self, dimensions, shape, dtype): """ Fetch the :class:`YaskContext` in ``self`` uniquely identified by ``dimensions``, ``shape``, and ``dtype``. Create a new (empty) :class:`YaskContext` on miss. """ # Sanity checks assert len(dimensions) == len(shape) dimensions = [str(i) for i in dimensions] if set(dimensions) < {'x', 'y', 'z'}: exit("Need a Function[x,y,z] for initialization") # The time dimension is dropped as implicit to the context domain = OrderedDict([(i, j) for i, j in zip(dimensions, shape) if i != namespace['time-dim']]) # A unique key for this context. key = tuple([configuration['isa'], dtype] + list(domain.items())) # Fetch or create a YaskContext if key in self: log("Fetched existing context from cache") else: self[key] = YaskContext('ctx%d' % self.ncontexts, domain, dtype) self.ncontexts += 1 log("Context successfully created!") return self[key]
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 fetch(self, dimensions, dtype): """ Fetch the :class:`YaskContext` in ``self`` uniquely identified by ``dimensions`` and ``dtype``. """ key = self._getkey(None, dtype, dimensions) context = self.get(key, self._partial_map.get(key)) if context is not None: log("Fetched existing YaskContext from cache") return context else: exit("Couldn't find YaskContext for key=`%s`" % str(key))
def fetch(self, dimensions, dtype): """ Fetch the YaskContext in ``self`` uniquely identified by ``dimensions`` and ``dtype``. """ key = self._getkey(None, dtype, dimensions) context = self.get(key, self._partial_map.get(key)) if context is not None: log("Fetched existing YaskContext from cache") return context else: exit("Couldn't find YaskContext for key=`%s`" % str(key))
def fetch(self, grid, dtype, dimensions=None): """ Fetch the :class:`YaskContext` in ``self`` uniquely identified by ``grid`` and ``dtype``. """ key = self._getkey(grid, dtype, dimensions) context = self.get(key, self._partial_map.get(key)) if context is not None: log("Fetched existing context from cache") return context else: exit("Couldn't find context for grid %s" % grid)
def make_grid(self, obj): """ Create and return a new :class:`YaskGrid`, a YASK grid wrapper. Memory is allocated. :param obj: The symbolic data object for which a YASK grid is allocated. """ if set(obj.indices) < set(self.space_dimensions): exit("Need a Function[x,y,z] to create a YASK grid.") name = 'devito_%s_%d' % (obj.name, contexts.ngrids) log("Allocating YaskGrid for %s (%s)" % (obj.name, str(obj.shape))) grid = self.yk_hook.new_grid(name, obj) wrapper = YaskGrid(grid, obj.shape, obj.space_order, obj.dtype) self.grids[name] = wrapper return wrapper
def make_grid(self, obj): """ Create and return a new :class:`Data`, a YASK grid wrapper. Memory is allocated. :param obj: The :class:`Function` for which a YASK grid is allocated. """ if set(obj.indices) < set(self.space_dimensions): exit("Need a Function[x,y,z] to create a YASK grid.") name = 'devito_%s_%d' % (obj.name, contexts.ngrids) # Create the YASK grid grid = self.yk_hook.new_grid(name, obj) # Where should memory be allocated ? alloc = obj._allocator if alloc.is_Numa: if alloc.put_onnode: grid.set_numa_preferred(alloc.node) elif alloc.put_local: grid.set_numa_preferred(namespace['numa-put-local']) for i, s, h in zip(obj.indices, obj.shape_allocated, obj._extent_halo): if i.is_Time: assert grid.is_dim_used(i.name) assert grid.get_alloc_size(i.name) == s else: # Note: # From the YASK docs: "If the halo is set to a value larger than # the padding size, the padding size will be automatically increased # to accomodate it." grid.set_left_halo_size(i.name, h.left) grid.set_right_halo_size(i.name, h.right) grid.alloc_storage() self.grids[name] = grid return grid
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 __init__(self, name, yc_soln, local_grids=None): """ Write out a YASK kernel, build it using YASK's Makefiles, import the corresponding SWIG-generated Python module, and finally create a YASK kernel solution object. :param name: Unique name of this YaskKernel. :param yc_soln: YaskCompiler solution. :param local_grids: A local grid is necessary to run the YaskKernel, but its final content can be ditched. Indeed, local grids are hidden to users -- for example, they could represent temporary arrays introduced by the DSE. This parameter tells which of the ``yc_soln``'s grids are local. """ self.name = name # Shared object name self.soname = "%s.%s.%s" % (name, yc_soln.get_name(), configuration['platform']) # It's necessary to `clean` the YASK kernel directory *before* # writing out the first `yask_stencil_code.hpp` make(namespace['path'], ['-C', namespace['kernel-path'], 'clean']) # Write out the stencil file if not os.path.exists(namespace['kernel-path-gen']): os.makedirs(namespace['kernel-path-gen']) yc_soln.format(configuration['isa'], ofac.new_file_output(namespace['kernel-output'])) # JIT-compile it try: compiler = configuration.yask['compiler'] opt_level = 1 if configuration.yask['develop-mode'] else 3 make( namespace['path'], [ '-j3', 'YK_CXX=%s' % compiler.cc, 'YK_CXXOPT=-O%d' % opt_level, 'mpi=0', # Disable MPI for now # "EXTRA_MACROS=TRACE", 'YK_BASE=%s' % str(name), 'stencil=%s' % yc_soln.get_name(), 'arch=%s' % configuration['platform'], '-C', namespace['kernel-path'], 'api' ]) except CompilationError: exit("Kernel solution compilation") # Import the corresponding Python (SWIG-generated) module try: yk = getattr(__import__('yask', fromlist=[name]), name) except ImportError: exit("Python YASK kernel bindings") try: yk = reload(yk) except NameError: # Python 3.5 compatibility yk = importlib.reload(yk) # Create the YASK solution object kfac = yk.yk_factory() self.env = kfac.new_env() self.soln = kfac.new_solution(self.env) # MPI setup: simple rank configuration in 1st dim only. # TODO: in production runs, the ranks would be distributed along all # domain dimensions. self.soln.set_num_ranks(self.space_dimensions[0], self.env.get_num_ranks()) # Redirect stdout/strerr to a string or file if configuration.yask['dump']: filename = 'yk_dump.%s.%s.%s.txt' % ( self.name, configuration['platform'], configuration['isa']) filename = os.path.join(configuration.yask['dump'], filename) self.output = yk.yask_output_factory().new_file_output(filename) else: self.output = yk.yask_output_factory().new_string_output() self.soln.set_debug_output(self.output) # Users may want to run the same Operator (same domain etc.) with # different grids. self.grids = {i.get_name(): i for i in self.soln.get_grids()} self.local_grids = { i.name: self.grids[i.name] for i in (local_grids or []) }
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 run_py(self, ntimesteps): exit("Cannot run a NullSolution through YASK's Python bindings")
def __init__(self, name, yc_soln, domain): """ Write out a YASK kernel, build it using YASK's Makefiles, import the corresponding SWIG-generated Python module, and finally create a YASK kernel solution object. :param name: Unique name of this YaskKernel. :param yc_soln: YaskCompiler solution. :param domain: A mapper from space dimensions to their domain size. """ self.name = name # Shared object name self.soname = "%s.%s.%s" % (name, yc_soln.get_name(), yask_configuration['arch']) # It's necessary to `clean` the YASK kernel directory *before* # writing out the first `yask_stencil_code.hpp` make(namespace['path'], ['-C', namespace['kernel-path'], 'clean']) # Write out the stencil file if not os.path.exists(namespace['kernel-path-gen']): os.makedirs(namespace['kernel-path-gen']) yc_soln.format(yask_configuration['isa'], ofac.new_file_output(namespace['kernel-output'])) # JIT-compile it try: opt_level = 1 if yask_configuration['develop-mode'] else 3 make( os.environ['YASK_HOME'], [ '-j', 'YK_CXXOPT=-O%d' % opt_level, # "EXTRA_MACROS=TRACE", 'YK_BASE=%s' % str(name), 'stencil=%s' % yc_soln.get_name(), 'arch=%s' % yask_configuration['arch'], '-C', namespace['kernel-path'], 'api' ]) except CompilationError: exit("Kernel solution compilation") # Import the corresponding Python (SWIG-generated) module try: yk = importlib.import_module(name) except ImportError: exit("Python YASK kernel bindings") try: yk = reload(yk) except NameError: # Python 3.5 compatibility yk = importlib.reload(yk) # Create the YASK solution object kfac = yk.yk_factory() self.env = kfac.new_env() self.soln = kfac.new_solution(self.env) # MPI setup: simple rank configuration in 1st dim only. # TODO: in production runs, the ranks would be distributed along all # domain dimensions. self.soln.set_num_ranks(self.soln.get_domain_dim_names()[0], self.env.get_num_ranks()) # Redirect stdout/strerr to a string self.output = yk.yask_output_factory().new_string_output() self.soln.set_debug_output(self.output) # Set up the solution domain size for k, v in domain.items(): self.soln.set_rank_domain_size(k, v)
def __init__(self, name, yc_soln, local_grids=None): """ Write out a YASK kernel, build it using YASK's Makefiles, import the corresponding SWIG-generated Python module, and finally create a YASK kernel solution object. :param name: Unique name of this YaskKernel. :param yc_soln: YaskCompiler solution. :param local_grids: A local grid is necessary to run the YaskKernel, but its final content can be ditched. Indeed, local grids are hidden to users -- for example, they could represent temporary arrays introduced by the DSE. This parameter tells which of the ``yc_soln``'s grids are local. """ self.name = name # Shared object name self.soname = "%s.devito.%s" % (name, configuration['platform']) # The directory in which the YASK-generated code (.hpp) will be placed yk_codegen = namespace['yask-codegen'](name, 'devito', configuration['platform']) if not os.path.exists(yk_codegen): os.makedirs(yk_codegen) # Write out the stencil file yk_codegen_file = os.path.join(yk_codegen, namespace['yask-codegen-file']) yc_soln.format(configuration['isa'], ofac.new_file_output(yk_codegen_file)) # JIT-compile it try: compiler = configuration.yask['compiler'] if configuration['develop-mode']: if yc_soln.get_num_equations() == 0: # YASK will compile more quickly, and no price has to be paid # in terms of performance, as this is a void kernel opt_level = 0 else: opt_level = 1 else: opt_level = 3 args = [ '-j', 'YK_CXX=%s' % compiler.cc, 'YK_CXXOPT=-O%d' % opt_level, # No MPI support at the moment 'mpi=0', # To locate the YASK compiler 'YC_EXEC=%s' % os.path.join(namespace['path'], 'bin'), # Error out if a grid not explicitly defined in the compiler is created 'allow_new_grid_types=0', # To give a unique name to the generated Python modules, rather # than creating `yask_kernel.py` 'YK_BASE=%s' % name, # `stencil` and `arch` should always be provided 'stencil=%s' % 'devito', 'arch=%s' % configuration['platform'], # The root directory of generated code files, shared libs, Python modules 'YASK_OUTPUT_DIR=%s' % namespace['yask-output-dir'], # Pick the YASK kernel Makefile, i.e. the one under `yask/src/kernel` '-C', namespace['kernel-path'], 'api' ] # Other potentially useful args: # - "EXTRA_MACROS=TRACE", -- debugging option make(namespace['path'], args) except CompilationError: exit("Kernel solution compilation") # Import the corresponding Python (SWIG-generated) module try: sys.path.append(os.path.join(namespace['yask-output-dir'], 'yask')) importlib.invalidate_caches() yk = importlib.import_module(name) except ImportError: exit("Python YASK kernel bindings") # Create the YASK solution object kfac = yk.yk_factory() self.env = kfac.new_env() self.soln = kfac.new_solution(self.env) # Apply any user-provided options, if any. # These are applied here instead of just before prepare_solution() # so that applicable options will apply to all API calls. self.soln.apply_command_line_options(configuration.yask['options'] or '') # MPI setup: simple rank configuration in 1st dim only. # TODO: in production runs, the ranks would be distributed along all # domain dimensions. self.soln.set_num_ranks(self.space_dimensions[0], self.env.get_num_ranks()) # Redirect stdout to a string or file if configuration.yask['dump']: filename = 'yk_dump.%s.%s.%s.txt' % ( name, configuration['platform'], configuration['isa']) filename = os.path.join(configuration.yask['dump'], filename) self.output = yk.yask_output_factory().new_file_output(filename) else: self.output = yk.yask_output_factory().new_string_output() self.soln.set_debug_output(self.output) # Users may want to run the same Operator (same domain etc.) with # different grids. self.grids = {i.get_name(): i for i in self.soln.get_grids()} self.local_grids = { i.name: self.grids[i.name] for i in (local_grids or []) }