def nodes(self): r""" :returns: object array of shape ``(ambient_dim,)`` containing :class:`~meshmode.dof_array.DOFArray`\ s of node coordinates. """ actx = self._setup_actx @memoize_in(actx, (Discretization, "nodes_prg")) def prg(): return make_loopy_program("""{[iel,idof,j]: 0<=iel<nelements and 0<=idof<ndiscr_nodes and 0<=j<nmesh_nodes}""", """ result[iel, idof] = \ sum(j, resampling_mat[idof, j] * nodes[iel, j]) """, name="nodes") return make_obj_array([ _DOFArray.from_list(None, [ actx.freeze( actx.call_loopy( prg(), resampling_mat=actx.from_numpy( grp.from_mesh_interp_matrix()), nodes=actx.from_numpy( grp.mesh_el_group.nodes[iaxis]))["result"]) for grp in self.groups ]) for iaxis in range(self.ambient_dim) ])
def num_reference_derivative(self, ref_axes, vec): actx = vec.array_context @memoize_in(actx, (Discretization, "reference_derivative_prg")) def prg(): return make_loopy_program( """{[iel,idof,j]: 0<=iel<nelements and 0<=idof,j<nunit_dofs}""", "result[iel,idof] = sum(j, diff_mat[idof, j] * vec[iel, j])", name="diff") def get_mat(grp): mat = None for ref_axis in ref_axes: next_mat = grp.diff_matrices()[ref_axis] if mat is None: mat = next_mat else: mat = np.dot(next_mat, mat) return mat return _DOFArray.from_list(actx, [ actx.call_loopy(prg(), diff_mat=actx.from_numpy(get_mat(grp)), vec=vec[grp.index])["result"] for grp in self.groups ])
def _interleave_dof_arrays(self, ary1, ary2): if not isinstance(ary1, DOFArray) or not isinstance(ary2, DOFArray): raise TypeError("non-array passed to connection") @memoize_in(self.array_context, (CenterGranularityConnection, "interleave")) def prg(): from meshmode.array_context import make_loopy_program return make_loopy_program( """{[iel, idof]: 0<=iel<nelements and 0<=idof<nunit_dofs}""", """ dst[iel, 2*idof] = src1[iel, idof] dst[iel, 2*idof + 1] = src2[iel, idof] """, [ lp.GlobalArg("src1", shape="(nelements, nunit_dofs)"), lp.GlobalArg("src2", shape="(nelements, nunit_dofs)"), lp.GlobalArg("dst", shape="(nelements, 2*nunit_dofs)"), "...", ], name="interleave") results = [] for grp, src1, src2 in zip(self.discr.groups, ary1, ary2): if src1.dtype != src2.dtype: raise ValueError("dtype mismatch in inputs") result = self.array_context.empty( (grp.nelements, 2 * grp.nunit_dofs), dtype=src1.dtype) self.array_context.call_loopy( prg(), src1=src1, src2=src2, dst=result, nelements=grp.nelements, nunit_dofs=grp.nunit_dofs) results.append(result) return DOFArray.from_list(self.array_context, results)
def _new_array(self, actx, creation_func, dtype=None): if dtype is None: dtype = self.real_dtype elif dtype == "c": dtype = self.complex_dtype else: dtype = np.dtype(dtype) return _DOFArray.from_list(actx, [ creation_func(shape=(grp.nelements, grp.nunit_dofs), dtype=dtype) for grp in self.groups ])
def _map_elementwise_reduction(self, reduction_name, expr): @memoize_in(self.places, "elementwise_node_" + reduction_name) def node_knl(): from meshmode.array_context import make_loopy_program return make_loopy_program("""{[iel, idof, jdof]: 0<=iel<nelements and 0<=idof, jdof<ndofs}""", """ result[iel, idof] = %s(jdof, operand[iel, jdof]) """ % reduction_name, name="nodewise_reduce") @memoize_in(self.places, "elementwise_" + reduction_name) def element_knl(): from meshmode.array_context import make_loopy_program return make_loopy_program("""{[iel, jdof]: 0<=iel<nelements and 0<=jdof<ndofs} """, """ result[iel, 0] = %s(jdof, operand[iel, jdof]) """ % reduction_name, name="elementwise_reduce") discr = self.places.get_discretization(expr.dofdesc.geometry, expr.dofdesc.discr_stage) operand = self.rec(expr.operand) assert operand.shape == (len(discr.groups), ) def _reduce(knl, result): for grp in discr.groups: self.array_context.call_loopy(knl, operand=operand[grp.index], result=result[grp.index]) return result dtype = operand.entry_dtype granularity = expr.dofdesc.granularity if granularity is sym.GRANULARITY_NODE: return _reduce(node_knl(), discr.empty(self.array_context, dtype=dtype)) elif granularity is sym.GRANULARITY_ELEMENT: result = DOFArray.from_list(self.array_context, [ self.array_context.empty((grp.nelements, 1), dtype=dtype) for grp in discr.groups ]) return _reduce(element_knl(), result) else: raise ValueError('unsupported granularity: %s' % granularity)
def map_insn_loopy_kernel(self, insn, profile_data=None): kdescr = insn.kernel_descriptor discr = self.discrwb.discr_from_dd(kdescr.governing_dd) dof_array_kwargs = {} other_kwargs = {} for name, expr in kdescr.input_mappings.items(): v = self.rec(expr) if isinstance(v, DOFArray): dof_array_kwargs[name] = v else: other_kwargs[name] = v for name in kdescr.scalar_args(): v = other_kwargs[name] if isinstance(v, (int, float)): other_kwargs[name] = discr.real_dtype.type(v) elif isinstance(v, complex): other_kwargs[name] = discr.complex_dtype.type(v) elif isinstance(v, np.number): pass else: raise ValueError("unrecognized scalar type for variable '%s': %s" % (name, type(v))) result = {} for grp in discr.groups: kwargs = other_kwargs.copy() kwargs["nelements"] = grp.nelements kwargs["nunit_dofs"] = grp.nunit_dofs for name, ary in dof_array_kwargs.items(): kwargs[name] = ary[grp.index] knl_result = self.array_context.call_loopy( kdescr.loopy_kernel, **kwargs) for name, val in knl_result.items(): result.setdefault(name, []).append(val) result = { name: DOFArray.from_list(self.array_context, val) for name, val in result.items()} return list(result.items()), []
def quad_weights(self): """:returns: A :class:`~meshmode.dof_array.DOFArray` with quadrature weights. """ actx = self._setup_actx @memoize_in(actx, (Discretization, "quad_weights_prg")) def prg(): return make_loopy_program( "{[iel,idof]: 0<=iel<nelements and 0<=idof<nunit_dofs}", "result[iel,idof] = weights[idof]", name="quad_weights") return _DOFArray.from_list(None, [ actx.freeze( actx.call_loopy( prg(), weights=actx.from_numpy(grp.weights), nelements=grp.nelements, )["result"]) for grp in self.groups ])
def test_node_reduction(ctx_factory): ctx = ctx_factory() queue = cl.CommandQueue(ctx) actx = PyOpenCLArrayContext(queue) # {{{ build discretization target_order = 4 nelements = 32 mesh = make_curve_mesh(starfish, np.linspace(0.0, 1.0, nelements + 1), target_order) discr = Discretization( actx, mesh, InterpolatoryQuadratureSimplexGroupFactory(target_order)) # }}} # {{{ test # create a shuffled [1, nelements + 1] array ary = [] el_nr_base = 0 for grp in discr.groups: x = 1 + np.arange(el_nr_base, grp.nelements) np.random.shuffle(x) ary.append(actx.freeze(actx.from_numpy(x.reshape(-1, 1)))) el_nr_base += grp.nelements from meshmode.dof_array import DOFArray ary = DOFArray.from_list(actx, ary) for func, expected in [ (sym.NodeSum, nelements * (nelements + 1) // 2), (sym.NodeMax, nelements), (sym.NodeMin, 1), ]: r = bind(discr, func(sym.var("x")))(actx, x=ary) assert abs(r - expected) < 1.0e-15, r
def test_to_fd_idempotency(ctx_factory, mm_mesh, fspace_degree): """ Make sure mm->fd->mm and (mm->)->fd->mm->fd are identity """ cl_ctx = ctx_factory() queue = cl.CommandQueue(cl_ctx) actx = PyOpenCLArrayContext(queue) # make sure degree is higher order than mesh fspace_degree += mm_mesh.groups[0].order # Make a function space and a function with unique values at each node factory = InterpolatoryQuadratureSimplexGroupFactory(fspace_degree) discr = Discretization(actx, mm_mesh, factory) fdrake_connection = build_connection_to_firedrake(discr) fdrake_mesh = fdrake_connection.firedrake_fspace().mesh() dtype = fdrake_mesh.coordinates.dat.data.dtype mm_unique = discr.zeros(actx, dtype=dtype) unique_vals = np.arange(np.size(mm_unique[0]), dtype=dtype) mm_unique[0].set(unique_vals.reshape(mm_unique[0].shape)) mm_unique_copy = DOFArray.from_list(actx, [mm_unique[0].copy()]) # Test for idempotency mm->fd->mm fdrake_unique = fdrake_connection.from_meshmode(mm_unique) fdrake_connection.from_firedrake(fdrake_unique, out=mm_unique_copy) np.testing.assert_allclose(actx.to_numpy(mm_unique_copy[0]), actx.to_numpy(mm_unique[0]), atol=CLOSE_ATOL) # Test for idempotency (mm->)fd->mm->fd fdrake_unique_copy = fdrake_connection.from_meshmode(mm_unique_copy) np.testing.assert_allclose(fdrake_unique_copy.dat.data, fdrake_unique.dat.data, atol=CLOSE_ATOL)
def __call__(self, ary): from meshmode.dof_array import DOFArray if not isinstance(ary, DOFArray): raise TypeError("non-array passed to discretization connection") actx = ary.array_context @memoize_in(actx, ( DirectDiscretizationConnection, "resample_by_mat_knl", self.is_surjective, )) def mat_knl(): if self.is_surjective: domains = [ """ {[iel, idof, j]: 0<=iel<nelements and 0<=idof<n_to_nodes and 0<=j<n_from_nodes} """, ] instructions = """ result[to_element_indices[iel], idof] \ = sum(j, resample_mat[idof, j] \ * ary[from_element_indices[iel], j]) """ else: domains = [ """ {[iel_init, idof_init]: 0<=iel_init<nelements_result and 0<=idof_init<n_to_nodes} """, """ {[iel, idof, j]: 0<=iel<nelements and 0<=idof<n_to_nodes and 0<=j<n_from_nodes} """, ] instructions = """ result[iel_init, idof_init] = 0 {id=init} ... gbarrier {id=barrier, dep=init} result[to_element_indices[iel], idof] \ = sum(j, resample_mat[idof, j] \ * ary[from_element_indices[iel], j]) {dep=barrier} """ knl = make_loopy_program( domains, instructions, [ lp.GlobalArg("result", None, shape="nelements_result, n_to_nodes", offset=lp.auto), lp.GlobalArg("ary", None, shape="nelements_vec, n_from_nodes", offset=lp.auto), lp.ValueArg("nelements_result", np.int32), lp.ValueArg("nelements_vec", np.int32), lp.ValueArg("n_from_nodes", np.int32), "...", ], name="resample_by_mat") return knl @memoize_in(actx, (DirectDiscretizationConnection, "resample_by_picking_knl", self.is_surjective)) def pick_knl(): if self.is_surjective: domains = [ """{[iel, idof]: 0<=iel<nelements and 0<=idof<n_to_nodes}""" ] instructions = """ result[to_element_indices[iel], idof] \ = ary[from_element_indices[iel], pick_list[idof]] """ else: domains = [ """ {[iel_init, idof_init]: 0<=iel_init<nelements_result and 0<=idof_init<n_to_nodes} """, """ {[iel, idof]: 0<=iel<nelements and 0<=idof<n_to_nodes} """ ] instructions = """ result[iel_init, idof_init] = 0 {id=init} ... gbarrier {id=barrier, dep=init} result[to_element_indices[iel], idof] \ = ary[from_element_indices[iel], pick_list[idof]] {dep=barrier} """ knl = make_loopy_program( domains, instructions, [ lp.GlobalArg("result", None, shape="nelements_result, n_to_nodes", offset=lp.auto), lp.GlobalArg("ary", None, shape="nelements_vec, n_from_nodes", offset=lp.auto), lp.ValueArg("nelements_result", np.int32), lp.ValueArg("nelements_vec", np.int32), lp.ValueArg("n_from_nodes", np.int32), "...", ], name="resample_by_picking") return knl if ary.shape != (len(self.from_discr.groups), ): raise ValueError("invalid shape of incoming resampling data") group_idx_to_result = [] for i_tgrp, (tgrp, cgrp) in enumerate(zip(self.to_discr.groups, self.groups)): kernels = [] # get kernels for each batch; to be fused eventually kwargs = {} # kwargs to the fused kernel for i_batch, batch in enumerate(cgrp.batches): if batch.from_element_indices.size == 0: continue point_pick_indices = self._resample_point_pick_indices( actx, i_tgrp, i_batch) if point_pick_indices is None: knl = mat_knl() knl = lp.rename_argument(knl, "resample_mat", f"resample_mat_{i_batch}") kwargs[f"resample_mat_{i_batch}"] = (self._resample_matrix( actx, i_tgrp, i_batch)) else: knl = pick_knl() knl = lp.rename_argument(knl, "pick_list", f"pick_list_{i_batch}") kwargs[f"pick_list_{i_batch}"] = point_pick_indices # {{{ enforce different namespaces for the kernels for iname in knl.all_inames(): knl = lp.rename_iname(knl, iname, f"{iname}_{i_batch}") knl = lp.rename_argument(knl, "ary", f"ary_{i_batch}") knl = lp.rename_argument(knl, "from_element_indices", f"from_element_indices_{i_batch}") knl = lp.rename_argument(knl, "to_element_indices", f"to_element_indices_{i_batch}") knl = lp.rename_argument(knl, "nelements", f"nelements_{i_batch}") # }}} kwargs[f"ary_{i_batch}"] = ary[batch.from_group_index] kwargs[f"from_element_indices_{i_batch}"] = ( batch.from_element_indices) kwargs[f"to_element_indices_{i_batch}"] = ( batch.to_element_indices) kernels.append(knl) fused_knl = lp.fuse_kernels(kernels) # order of operations doesn't matter fused_knl = lp.add_nosync(fused_knl, "global", "writes:result", "writes:result", bidirectional=True, force=True) result_dict = actx.call_loopy(fused_knl, nelements_result=tgrp.nelements, n_to_nodes=tgrp.nunit_dofs, **kwargs) group_idx_to_result.append(result_dict["result"]) from meshmode.dof_array import DOFArray return DOFArray.from_list(actx, group_idx_to_result)
def map_insn_loopy_kernel(self, insn, profile_data): kdescr = insn.kernel_descriptor discr = self.inner_mapper.discrwb.discr_from_dd(kdescr.governing_dd) dof_array_kwargs = {} other_kwargs = {} for name, expr in kdescr.input_mappings.items(): v = self.inner_mapper.rec(expr) if isinstance(v, DOFArray): dof_array_kwargs[name] = v if profile_data is not None: size = dof_array_nbytes(v) profile_data["bytes_read"] = ( profile_data.get("bytes_read", 0) + size) profile_data["bytes_read_by_scalar_assignments"] = ( profile_data.get("bytes_read_by_scalar_assignments", 0) + size) else: assert not isinstance(v, np.ndarray) other_kwargs[name] = v for name in kdescr.scalar_args(): v = other_kwargs[name] if isinstance(v, (int, float)): other_kwargs[name] = discr.real_dtype.type(v) elif isinstance(v, complex): other_kwargs[name] = discr.complex_dtype.type(v) elif isinstance(v, np.number): pass else: raise ValueError( "unrecognized scalar type for variable '%s': %s" % (name, type(v))) result = {} for grp in discr.groups: kwargs = other_kwargs.copy() kwargs["nelements"] = grp.nelements kwargs["nunit_dofs"] = grp.nunit_dofs for name, ary in dof_array_kwargs.items(): kwargs[name] = ary[grp.index] knl_result = self.inner_mapper.array_context.call_loopy( kdescr.loopy_kernel, **kwargs) for name, val in knl_result.items(): result.setdefault(name, []).append(val) result = { name: DOFArray.from_list(self.inner_mapper.array_context, val) for name, val in result.items() } for val in result.values(): assert isinstance(val, DOFArray) if profile_data is not None: size = dof_array_nbytes(val) profile_data["bytes_written"] = ( profile_data.get("bytes_written", 0) + size) profile_data["bytes_written_by_scalar_assignments"] = ( profile_data.get("bytes_written_by_scalar_assignments", 0) + size) return list(result.items()), []
def map_if(self, expr): bool_crit = self.rec(expr.condition) if isinstance(bool_crit, DOFArray): # continues below pass elif isinstance(bool_crit, np.number): if bool_crit: return self.rec(expr.then) else: return self.rec(expr.else_) else: raise TypeError( "Expected criterion to be of type np.number or DOFArray") assert isinstance(bool_crit, DOFArray) ngroups = len(bool_crit) from pymbolic import var iel = var("iel") idof = var("idof") subscript = (iel, idof) then = self.rec(expr.then) else_ = self.rec(expr.else_) import pymbolic.primitives as p var = p.Variable if isinstance(then, DOFArray): sym_then = var("a")[subscript] def get_then(igrp): return then[igrp] elif isinstance(then, np.number): sym_then = var("a") def get_then(igrp): return then else: raise TypeError( "Expected 'then' to be of type np.number or DOFArray") if isinstance(else_, DOFArray): sym_else = var("b")[subscript] def get_else(igrp): return else_[igrp] elif isinstance(else_, np.number): sym_else = var("b") def get_else(igrp): return else_ else: raise TypeError( "Expected 'else' to be of type np.number or DOFArray") @memoize_in(self.array_context, (ExecutionMapper, "map_if_knl")) def knl(sym_then, sym_else): return make_loopy_program( "{[iel, idof]: 0<=iel<nelements and 0<=idof<nunit_dofs}", [ lp.Assignment(var("out")[iel, idof], p.If(var("crit")[iel, idof], sym_then, sym_else)) ]) return DOFArray.from_list(self.array_context, [ self.array_context.call_loopy( knl(sym_then, sym_else), crit=bool_crit[igrp], a=get_then(igrp), b=get_else(igrp)) for igrp in range(ngroups)])