def restrict_kernel(Vf, Vc): hierarchy, level = utils.get_level(Vc.ufl_domain()) levelf = level + Fraction(1 / hierarchy.refinements_per_level) cache = hierarchy._shared_data_cache["transfer_kernels"] coordinates = Vc.ufl_domain().coordinates key = (("restrict", ) + Vf.ufl_element().value_shape() + entity_dofs_key(Vf.finat_element.entity_dofs()) + entity_dofs_key(Vc.finat_element.entity_dofs()) + entity_dofs_key( coordinates.function_space().finat_element.entity_dofs())) try: return cache[key] except KeyError: mesh = coordinates.ufl_domain() evaluate_kernel = compile_element(firedrake.TestFunction(Vc), Vf) to_reference_kernel = to_reference_coordinates( coordinates.ufl_element()) coords_element = create_element(coordinates.ufl_element()) element = create_element(Vc.ufl_element()) eval_args = evaluate_kernel.args[:-1] args = eval_args[-1].gencode(not_scope=True) R, fine = (a.sym.symbol for a in eval_args) my_kernel = """#include <petsc.h> %(to_reference)s %(evaluate)s __attribute__((noinline)) /* Clang bug */ static void pyop2_kernel_restrict(PetscScalar *R, %(args)s, const PetscScalar *X, const PetscScalar *Xc) { PetscScalar Xref[%(tdim)d]; int cell = -1; int bestcell = -1; double bestdist = 1e10; for (int i = 0; i < %(ncandidate)d; i++) { const PetscScalar *Xci = Xc + i*%(Xc_cell_inc)d; double celldist = 2*bestdist; to_reference_coords_kernel(Xref, X, Xci); if (%(inside_cell)s) { cell = i; break; } %(compute_celldist)s /* fprintf(stderr, "cell %%d celldist: %%.14e\\n", i, celldist); fprintf(stderr, "Xref: %%.14e %%.14e %%.14e\\n", Xref[0], Xref[1], Xref[2]); */ if (celldist < bestdist) { bestdist = celldist; bestcell = i; } } if (cell == -1) { /* We didn't find a cell that contained this point exactly. Did we find one that was close enough? */ if (bestdist < 10) { cell = bestcell; } else { fprintf(stderr, "Could not identify cell in transfer operator. Point: "); for (int coord = 0; coord < %(spacedim)s; coord++) { fprintf(stderr, "%%.14e ", X[coord]); } fprintf(stderr, "\\n"); fprintf(stderr, "Number of candidates: %%d. Best distance located: %%14e", %(ncandidate)d, bestdist); abort(); } } { const PetscScalar *Ri = %(R)s + cell*%(coarse_cell_inc)d; pyop2_kernel_evaluate(Ri, %(fine)s, Xref); } } """ % { "to_reference": str(to_reference_kernel), "evaluate": str(evaluate_kernel), "ncandidate": hierarchy.fine_to_coarse_cells[levelf].shape[1], "inside_cell": inside_check(element.cell, eps=1e-8, X="Xref"), "compute_celldist": compute_celldist(element.cell, X="Xref", celldist="celldist"), "Xc_cell_inc": coords_element.space_dimension(), "coarse_cell_inc": element.space_dimension(), "args": args, "spacedim": element.cell.get_spatial_dimension(), "R": R, "fine": fine, "tdim": mesh.topological_dimension() } return cache.setdefault( key, op2.Kernel(my_kernel, name="pyop2_kernel_restrict"))
def inject_kernel(Vf, Vc): hierarchy, level = utils.get_level(Vc.ufl_domain()) cache = hierarchy._shared_data_cache["transfer_kernels"] coordinates = Vf.ufl_domain().coordinates key = ( ("inject", ) + Vf.ufl_element().value_shape() + entity_dofs_key(Vc.finat_element.entity_dofs()) + entity_dofs_key(Vf.finat_element.entity_dofs()) + entity_dofs_key( Vc.mesh().coordinates.function_space().finat_element.entity_dofs()) + entity_dofs_key( coordinates.function_space().finat_element.entity_dofs())) try: return cache[key] except KeyError: ncandidate = hierarchy.coarse_to_fine_cells[level].shape[1] if Vc.finat_element.entity_dofs( ) == Vc.finat_element.entity_closure_dofs(): return cache.setdefault( key, (dg_injection_kernel(Vf, Vc, ncandidate), True)) coordinates = Vf.ufl_domain().coordinates evaluate_kernel = compile_element(ufl.Coefficient(Vf)) to_reference_kernel = to_reference_coordinates( coordinates.ufl_element()) coords_element = create_element(coordinates.ufl_element()) Vf_element = create_element(Vf.ufl_element()) kernel = """ %(to_reference)s %(evaluate)s __attribute__((noinline)) /* Clang bug */ static void pyop2_kernel_inject(PetscScalar *R, const PetscScalar *X, const PetscScalar *f, const PetscScalar *Xf) { PetscScalar Xref[%(tdim)d]; int cell = -1; int bestcell = -1; double bestdist = 1e10; for (int i = 0; i < %(ncandidate)d; i++) { const PetscScalar *Xfi = Xf + i*%(Xf_cell_inc)d; double celldist = 2*bestdist; to_reference_coords_kernel(Xref, X, Xfi); if (%(inside_cell)s) { cell = i; break; } %(compute_celldist)s if (celldist < bestdist) { bestdist = celldist; bestcell = i; } } if (cell == -1) { /* We didn't find a cell that contained this point exactly. Did we find one that was close enough? */ if (bestdist < 10) { cell = bestcell; } else { fprintf(stderr, "Could not identify cell in transfer operator. Point: "); for (int coord = 0; coord < %(spacedim)s; coord++) { fprintf(stderr, "%%.14e ", X[coord]); } fprintf(stderr, "\\n"); fprintf(stderr, "Number of candidates: %%d. Best distance located: %%14e", %(ncandidate)d, bestdist); abort(); } } const PetscScalar *fi = f + cell*%(f_cell_inc)d; for ( int i = 0; i < %(Rdim)d; i++ ) { R[i] = 0; } pyop2_kernel_evaluate(R, fi, Xref); } """ % { "to_reference": str(to_reference_kernel), "evaluate": str(evaluate_kernel), "inside_cell": inside_check(Vc.finat_element.cell, eps=1e-8, X="Xref"), "spacedim": Vc.finat_element.cell.get_spatial_dimension(), "compute_celldist": compute_celldist( Vc.finat_element.cell, X="Xref", celldist="celldist"), "tdim": Vc.ufl_domain().topological_dimension(), "ncandidate": ncandidate, "Rdim": numpy.prod(Vf_element.value_shape), "Xf_cell_inc": coords_element.space_dimension(), "f_cell_inc": Vf_element.space_dimension() } return cache.setdefault( key, (op2.Kernel(kernel, name="pyop2_kernel_inject"), False))
def prolong_kernel(expression): meshc = expression.ufl_domain() hierarchy, level = utils.get_level(expression.ufl_domain()) levelf = level + Fraction(1 / hierarchy.refinements_per_level) cache = hierarchy._shared_data_cache["transfer_kernels"] coordinates = expression.ufl_domain().coordinates if meshc.cell_set._extruded: idx = levelf * hierarchy.refinements_per_level assert idx == int(idx) level_ratio = (hierarchy._meshes[int(idx)].layers - 1) // (meshc.layers - 1) else: level_ratio = 1 key = (("prolong", level_ratio) + expression.ufl_element().value_shape() + entity_dofs_key( expression.function_space().finat_element.entity_dofs()) + entity_dofs_key( coordinates.function_space().finat_element.entity_dofs())) try: return cache[key] except KeyError: mesh = coordinates.ufl_domain() evaluate_kernel = compile_element(expression) to_reference_kernel = to_reference_coordinates( coordinates.ufl_element()) element = create_element(expression.ufl_element()) eval_args = evaluate_kernel.args[:-1] coords_element = create_element(coordinates.ufl_element()) args = eval_args[-1].gencode(not_scope=True) R, coarse = (a.sym.symbol for a in eval_args) my_kernel = """#include <petsc.h> %(to_reference)s %(evaluate)s __attribute__((noinline)) /* Clang bug */ static void pyop2_kernel_prolong(PetscScalar *R, %(args)s, const PetscScalar *X, const PetscScalar *Xc) { PetscScalar Xref[%(tdim)d]; int cell = -1; int bestcell = -1; double bestdist = 1e10; for (int i = 0; i < %(ncandidate)d; i++) { const PetscScalar *Xci = Xc + i*%(Xc_cell_inc)d; double celldist = 2*bestdist; to_reference_coords_kernel(Xref, X, Xci); if (%(inside_cell)s) { cell = i; break; } %(compute_celldist)s if (celldist < bestdist) { bestdist = celldist; bestcell = i; } } if (cell == -1) { /* We didn't find a cell that contained this point exactly. Did we find one that was close enough? */ if (bestdist < 10) { cell = bestcell; } else { fprintf(stderr, "Could not identify cell in transfer operator. Point: "); for (int coord = 0; coord < %(spacedim)s; coord++) { fprintf(stderr, "%%.14e ", X[coord]); } fprintf(stderr, "\\n"); fprintf(stderr, "Number of candidates: %%d. Best distance located: %%14e", %(ncandidate)d, bestdist); abort(); } } const PetscScalar *coarsei = %(coarse)s + cell*%(coarse_cell_inc)d; for ( int i = 0; i < %(Rdim)d; i++ ) { %(R)s[i] = 0; } pyop2_kernel_evaluate(%(R)s, coarsei, Xref); } """ % { "to_reference": str(to_reference_kernel), "evaluate": str(evaluate_kernel), "args": args, "R": R, "spacedim": element.cell.get_spatial_dimension(), "coarse": coarse, "ncandidate": hierarchy.fine_to_coarse_cells[levelf].shape[1] * level_ratio, "Rdim": numpy.prod(element.value_shape), "inside_cell": inside_check(element.cell, eps=1e-8, X="Xref"), "compute_celldist": compute_celldist(element.cell, X="Xref", celldist="celldist"), "Xc_cell_inc": coords_element.space_dimension(), "coarse_cell_inc": element.space_dimension(), "tdim": mesh.topological_dimension() } return cache.setdefault( key, op2.Kernel(my_kernel, name="pyop2_kernel_prolong"))