def freesurface(model, eq): """ Generate the stencil that mirrors the field as a free surface modeling for the acoustic wave equation. Parameters ---------- model : Model Physical model. eq : Eq Time-stepping stencil (time update) to mirror at the freesurface. """ lhs, rhs = eq.evaluate.args # Get vertical dimension and corresponding subdimension zfs = model.grid.subdomains['fsdomain'].dimensions[-1] z = zfs.parent # Functions present in the stencil funcs = retrieve_functions(rhs) mapper = {} # Antisymmetric mirror at negative indices # TODO: Make a proper "mirror_indices" tool function for f in funcs: zind = f.indices[-1] if (zind - z).as_coeff_Mul()[0] < 0: s = sign(zind.subs({z: zfs, z.spacing: 1})) mapper.update({f: s * f.subs({zind: INT(abs(zind))})}) return Eq(lhs, rhs.subs(mapper), subdomain=model.grid.subdomains['fsdomain'])
def inject(self, field, expr, offset=0): """ Generate equations injecting an arbitrary expression into a field. Parameters ---------- field : Function Input field into which the injection is performed. expr : expr-like Injected expression. offset : int, optional Additional offset from the boundary. """ expr = indexify(expr) field = indexify(field) p, _ = self.gridpoints.indices dim_subs = [] coeffs = [] for i, d in enumerate(self.grid.dimensions): rd = DefaultDimension(name="r%s" % d.name, default_value=self.r) dim_subs.append((d, INT(rd + self.gridpoints[p, i]))) coeffs.append(self.interpolation_coeffs[p, i, rd]) rhs = prod(coeffs) * expr field = field.subs(dim_subs) return [Eq(field, field + rhs.subs(dim_subs))]
def _coordinate_indices(self): """Symbol for each grid index according to the coordinates.""" return tuple([ INT(FLOOR((c - o) / i.spacing)) for c, o, i in zip(self._coordinate_symbols, self.grid.origin, self.grid.dimensions[:self.grid.dim]) ])
def interpolate(self, expr, offset=0, increment=False, self_subs={}): """ Generate equations interpolating an arbitrary expression into ``self``. Parameters ---------- expr : expr-like Input expression to interpolate. offset : int, optional Additional offset from the boundary. increment: bool, optional If True, generate increments (Inc) rather than assignments (Eq). """ expr = indexify(expr) p, _, _ = self.interpolation_coeffs.indices dim_subs = [] coeffs = [] for i, d in enumerate(self.grid.dimensions): rd = DefaultDimension(name="r%s" % d.name, default_value=self.r) dim_subs.append((d, INT(rd + self.gridpoints[p, i]))) coeffs.append(self.interpolation_coeffs[p, i, rd]) # Apply optional time symbol substitutions to lhs of assignment lhs = self.subs(self_subs) rhs = prod(coeffs) * expr.subs(dim_subs) return [Eq(lhs, lhs + rhs)]
def _coordinate_indices(self): """Symbol for each grid index according to the coordinates.""" indices = self.grid.dimensions return tuple([ INT(sympy.Function('floor')((c - o) / i.spacing)) for c, o, i in zip(self._coordinate_symbols, self.grid.origin, indices[:self.grid.dim]) ])
def make_var_gets(expr): mapper = {} indexeds = retrieve_indexed(expr) data_carriers = [i for i in indexeds if i.base.function.from_YASK] for i in data_carriers: args = [ ListInitializer([INT(make_var_gets(j)) for j in i.indices]) ] mapper[i] = make_sharedptr_funcall( namespace['code-var-get'], args, yk_var_objs[i.base.function.name]) return expr.xreplace(mapper)
def _make_partree(self, candidates, nthreads=None): """Parallelize the `candidates` Iterations attaching suitable OpenMP pragmas.""" assert candidates root = candidates[0] # Get the collapsable Iterations collapsable = self._find_collapsable(root, candidates) ncollapse = 1 + len(collapsable) # Prepare to build a ParallelTree if all(i.is_Affine for i in candidates): bundles = FindNodes(ExpressionBundle).visit(root) sops = sum(i.ops for i in bundles) if sops >= self.dynamic_work: schedule = 'dynamic' else: schedule = 'static' if nthreads is None: # pragma omp for ... schedule(..., 1) nthreads = self.nthreads body = OpenMPIteration(schedule=schedule, ncollapse=ncollapse, **root.args) else: # pragma omp parallel for ... schedule(..., 1) body = OpenMPIteration(schedule=schedule, parallel=True, ncollapse=ncollapse, nthreads=nthreads, **root.args) prefix = [] else: # pragma omp for ... schedule(..., expr) assert nthreads is None nthreads = self.nthreads_nonaffine chunk_size = Symbol(name='chunk_size') body = OpenMPIteration(ncollapse=ncollapse, chunk_size=chunk_size, **root.args) niters = prod([root.symbolic_size] + [j.symbolic_size for j in collapsable]) value = INT(Max(niters / (nthreads * self.chunk_nonaffine), 1)) prefix = [Expression(DummyEq(chunk_size, value, dtype=np.int32))] # Create a ParallelTree partree = ParallelTree(prefix, body, nthreads=nthreads) collapsed = [partree] + collapsable return root, partree, collapsed
def callback(): _expr = indexify(expr) _field = indexify(field) p, _ = self.obj.gridpoints.indices dim_subs = [] coeffs = [] for i, d in enumerate(self.obj.grid.dimensions): rd = DefaultDimension(name="r%s" % d.name, default_value=self.r) dim_subs.append((d, INT(rd + self.obj.gridpoints[p, i]))) coeffs.append(self.obj.interpolation_coeffs[p, i, rd]) rhs = prod(coeffs) * _expr _field = _field.subs(dim_subs) return [Eq(_field, _field + rhs.subs(dim_subs))]
def make_var_accesses(node, yk_var_objs): """ Construct a new Iteration/Expression based on ``node``, in which all :class:`types.Indexed` accesses have been converted into YASK var accesses. """ def make_var_gets(expr): mapper = {} indexeds = retrieve_indexed(expr) data_carriers = [i for i in indexeds if i.base.function.from_YASK] for i in data_carriers: args = [ ListInitializer([INT(make_var_gets(j)) for j in i.indices]) ] mapper[i] = make_sharedptr_funcall( namespace['code-var-get'], args, yk_var_objs[i.base.function.name]) return expr.xreplace(mapper) mapper = {} for i, e in enumerate(FindNodes(Expression).visit(node)): if e.is_ForeignExpression: continue lhs, rhs = e.expr.args # RHS translation rhs = make_var_gets(rhs) # LHS translation if e.write.from_YASK: args = [rhs] args += [ ListInitializer([INT(make_var_gets(i)) for i in lhs.indices]) ] call = namespace['code-var-add' if e. is_Increment else 'code-var-put'] handle = make_sharedptr_funcall(call, args, yk_var_objs[e.write.name]) processed = ForeignExpression(handle, e.dtype, is_Increment=e.is_Increment) else: # Writing to a scalar temporary processed = e._rebuild(expr=e.expr.func(lhs, rhs)) mapper.update({e: processed}) return Transformer(mapper).visit(node)
def callback(): _expr = indexify(expr) p, _, _ = self.obj.interpolation_coeffs.indices dim_subs = [] coeffs = [] for i, d in enumerate(self.obj.grid.dimensions): rd = DefaultDimension(name="r%s" % d.name, default_value=self.r) dim_subs.append((d, INT(rd + self.obj.gridpoints[p, i]))) coeffs.append(self.obj.interpolation_coeffs[p, i, rd]) # Apply optional time symbol substitutions to lhs of assignment lhs = self.obj.subs(self_subs) rhs = prod(coeffs) * _expr.subs(dim_subs) return [Eq(lhs, lhs + rhs)]
def _make_partree(self, candidates, nthreads=None): """Parallelize the `candidates` Iterations attaching suitable OpenMP pragmas.""" assert candidates root = candidates[0] # Get the collapsable Iterations collapsable = self._find_collapsable(root, candidates) ncollapse = 1 + len(collapsable) # Prepare to build a ParallelTree prefix = [] if all(i.is_Affine for i in candidates): if nthreads is None: # pragma omp for ... schedule(..., 1) nthreads = self.nthreads omp_pragma = self.lang['for'](ncollapse, 1) else: # pragma omp parallel for ... schedule(..., 1) omp_pragma = self.lang['par-for'](ncollapse, 1, nthreads) else: # pragma omp for ... schedule(..., expr) assert nthreads is None nthreads = self.nthreads_nonaffine chunk_size = Symbol(name='chunk_size') omp_pragma = self.lang['for'](ncollapse, chunk_size) niters = prod([root.symbolic_size] + [j.symbolic_size for j in collapsable]) value = INT(Max(niters / (nthreads * self.CHUNKSIZE_NONAFFINE), 1)) prefix.append( Expression(DummyEq(chunk_size, value, dtype=np.int32))) # Create a ParallelTree body = root._rebuild(pragmas=root.pragmas + (omp_pragma, ), properties=root.properties + (COLLAPSED(ncollapse), )) partree = ParallelTree(prefix, body, nthreads=nthreads) collapsed = [partree] + collapsable return root, partree, collapsed
def freesurface(model, eq): """ Generate the stencil that mirrors the field as a free surface modeling for the acoustic wave equation Parameters ---------- model: Model Physical model eq: Eq or List of Eq Equation to apply mirror to """ fs_eq = [] for eq_i in eq: for p in eq_i._flatten: lhs, rhs = p.evaluate.args # Add modulo replacements to to rhs zfs = model.grid.subdomains['fsdomain'].dimensions[-1] z = zfs.parent funcs = retrieve_functions(rhs.evaluate) mapper = {} for f in funcs: zind = f.indices[-1] if (zind - z).as_coeff_Mul()[0] < 0: s = sign((zind - z.symbolic_min).subs({ z: zfs, z.spacing: 1 })) mapper.update({f: s * f.subs({zind: INT(abs(zind))})}) fs_eq.append( Eq(lhs, sign(lhs.indices[-1] - z.symbolic_min) * rhs.subs(mapper), subdomain=model.grid.subdomains['fsdomain'])) return fs_eq
def _make_partree(self, candidates, nthreads=None): """Parallelize `root` attaching a suitable OpenMP pragma.""" assert candidates root = candidates[0] # Get the collapsable Iterations collapsable = [] if ncores() >= Ompizer.COLLAPSE_NCORES and IsPerfectIteration().visit( root): for n, i in enumerate(candidates[1:], 1): # The OpenMP specification forbids collapsed loops to use iteration # variables in initializer expressions. E.g., the following is forbidden: # # #pragma omp ... collapse(2) # for (i = ... ) # for (j = i ...) # ... # # Here, we make sure this won't happen if any(j.dim in i.symbolic_min.free_symbols for j in candidates[:n]): break # Also, we do not want to collapse vectorizable Iterations if i.is_Vectorizable: break # Would there be enough work per parallel iteration? try: work = prod( [int(j.dim.symbolic_size) for j in candidates[n + 1:]]) if work < Ompizer.COLLAPSE_WORK: break except TypeError: pass collapsable.append(i) ncollapse = 1 + len(collapsable) # Prepare to build a ParallelTree prefix = [] if all(i.is_Affine for i in candidates): if nthreads is None: # pragma omp for ... schedule(..., 1) nthreads = self.nthreads omp_pragma = self.lang['for'](ncollapse, 1) else: # pragma omp parallel for ... schedule(..., 1) omp_pragma = self.lang['par-for'](ncollapse, 1, nthreads) else: # pragma omp for ... schedule(..., expr) assert nthreads is None nthreads = self.nthreads_nonaffine chunk_size = Symbol(name='chunk_size') omp_pragma = self.lang['for'](ncollapse, chunk_size) niters = prod([root.symbolic_size] + [j.symbolic_size for j in collapsable]) value = INT(Max(niters / (nthreads * self.CHUNKSIZE_NONAFFINE), 1)) prefix.append( Expression(DummyEq(chunk_size, value, dtype=np.int32))) # Create a ParallelTree body = root._rebuild(pragmas=root.pragmas + (omp_pragma, ), properties=root.properties + (COLLAPSED(ncollapse), )) partree = ParallelTree(prefix, body, nthreads=nthreads) collapsed = [partree] + collapsable return root, partree, collapsed