Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 42 additions & 30 deletions src/underworld3/cython/petsc_generic_snes_solvers.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ from petsc4py import PETSc

import underworld3
import underworld3 as uw
from underworld3.utilities._jitextension import getext
from underworld3.utilities._jitextension import getext, JITCallbackSet
import underworld3.timing as timing

from underworld3.utilities._api_tools import uw_object
Expand Down Expand Up @@ -1542,15 +1542,19 @@ class SNES_Scalar(SolverBaseClass):
print(f"Scalar SNES: Jacobians complete, now compile", flush=True)

prim_field_list = [self.u]
_getext_result = getext(self.mesh,
tuple(fns_residual),
tuple(fns_jacobian),
[x.fn for x in self.essential_bcs],
tuple(fns_bd_residual),
tuple(fns_bd_jacobian),
primary_field_list=prim_field_list,
verbose=verbose,
debug=debug,)
_getext_result = getext(
self.mesh,
JITCallbackSet(
residual=tuple(fns_residual),
bcs=tuple(x.fn for x in self.essential_bcs),
jacobian=tuple(fns_jacobian),
bd_residual=tuple(fns_bd_residual),
bd_jacobian=tuple(fns_bd_jacobian),
),
prim_field_list,
verbose=verbose,
debug=debug,
)
self.compiled_extensions = _getext_result.ptrobj
self.ext_dict = _getext_result.fn_dicts
self.constants_manifest = _getext_result.constants_manifest
Expand Down Expand Up @@ -2291,15 +2295,19 @@ class SNES_Vector(SolverBaseClass):
# note also that the order here is important.

prim_field_list = [self.u,]
_getext_result = getext(self.mesh,
tuple(fns_residual),
tuple(fns_jacobian),
[x.fn for x in self.essential_bcs],
tuple(fns_bd_residual),
tuple(fns_bd_jacobian),
primary_field_list=prim_field_list,
verbose=verbose,
debug=debug,)
_getext_result = getext(
self.mesh,
JITCallbackSet(
residual=tuple(fns_residual),
bcs=tuple(x.fn for x in self.essential_bcs),
jacobian=tuple(fns_jacobian),
bd_residual=tuple(fns_bd_residual),
bd_jacobian=tuple(fns_bd_jacobian),
),
prim_field_list,
verbose=verbose,
debug=debug,
)
self.compiled_extensions = _getext_result.ptrobj
self.ext_dict = _getext_result.fn_dicts
self.constants_manifest = _getext_result.constants_manifest
Expand Down Expand Up @@ -3669,17 +3677,21 @@ class SNES_Stokes_SaddlePt(SolverBaseClass):
print(f"Stokes: Jacobians complete, now compile", flush=True)

prim_field_list = [self.u, self.p]
_getext_result = getext(self.mesh,
tuple(fns_residual),
tuple(fns_jacobian),
[x.fn for x in self.essential_bcs],
tuple(fns_bd_residual),
tuple(fns_bd_jacobian),
primary_field_list=prim_field_list,
verbose=verbose,
debug=debug,
debug_name=debug_name,
cache=False)
_getext_result = getext(
self.mesh,
JITCallbackSet(
residual=tuple(fns_residual),
bcs=tuple(x.fn for x in self.essential_bcs),
jacobian=tuple(fns_jacobian),
bd_residual=tuple(fns_bd_residual),
bd_jacobian=tuple(fns_bd_jacobian),
),
prim_field_list,
verbose=verbose,
debug=debug,
debug_name=debug_name,
cache=False,
)
self.compiled_extensions = _getext_result.ptrobj
self.ext_dict = _getext_result.fn_dicts
self.constants_manifest = _getext_result.constants_manifest
Expand Down
12 changes: 7 additions & 5 deletions src/underworld3/cython/petsc_maths.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ import sympy

import underworld3
import underworld3.timing as timing
from underworld3.utilities._jitextension import getext
from underworld3.utilities._jitextension import getext, JITCallbackSet

from petsc4py import PETSc

Expand Down Expand Up @@ -89,7 +89,8 @@ class Integral:
self.dm = self.mesh.dm # .clone()
mesh=self.mesh

_getext_result = getext(self.mesh, [self.fn,], [], [], [], [], self.mesh.vars.values(), verbose=verbose)
_getext_result = getext(self.mesh, JITCallbackSet(residual=(self.fn,)),
self.mesh.vars.values(), verbose=verbose)
cdef PtrContainer ext = _getext_result.ptrobj

# Pull out vec for variables, and go ahead with the integral
Expand Down Expand Up @@ -273,7 +274,8 @@ class CellWiseIntegral:
elif isinstance(self.fn, sympy.vector.Dyadic):
raise RuntimeError("Integral evaluation for Dyadic integrands not supported.")

cdef PtrContainer ext = getext(self.mesh, [self.fn,], [], [], [], [], self.mesh.vars.values()).ptrobj
cdef PtrContainer ext = getext(self.mesh, JITCallbackSet(residual=(self.fn,)),
self.mesh.vars.values()).ptrobj

# Pull out vec for variables, and go ahead with the integral
self.mesh.update_lvec()
Expand Down Expand Up @@ -388,8 +390,8 @@ class BdIntegral:

# Compile integrand using the boundary residual slot (includes petsc_n[] in signature)
_getext_result = getext(
self.mesh, [], [], [], [self.fn,], [], self.mesh.vars.values(), verbose=verbose
)
self.mesh, JITCallbackSet(bd_residual=(self.fn,)),
self.mesh.vars.values(), verbose=verbose)
cdef PtrContainer ext = _getext_result.ptrobj

# Prepare the solution vector
Expand Down
13 changes: 9 additions & 4 deletions src/underworld3/cython/petsc_types.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -22,21 +22,21 @@ ctypedef void(*PetscDSResidualFn)(PetscInt, PetscInt, PetscInt,
ctypedef void (*PetscDSJacobianFn)(PetscInt, PetscInt, PetscInt,
const PetscInt*, const PetscInt*, const PetscScalar*, const PetscScalar*, const PetscScalar*,
const PetscInt*, const PetscInt*, const PetscScalar*, const PetscScalar*, const PetscScalar*,
PetscReal, PetscReal, const PetscReal*, PetscInt, const PetscScalar*,
PetscReal, PetscReal, const PetscReal*, PetscInt, const PetscScalar*,
PetscScalar*)

ctypedef void(*PetscDSBdResidualFn)(
PetscInt, PetscInt, PetscInt,
const PetscInt*, const PetscInt*, const PetscScalar*, const PetscScalar*, const PetscScalar*,
const PetscInt*, const PetscInt*, const PetscScalar*, const PetscScalar*, const PetscScalar*,
PetscReal, const PetscReal*,const PetscReal*, PetscInt, const PetscScalar*,
PetscReal, const PetscReal*,const PetscReal*, PetscInt, const PetscScalar*,
PetscScalar* )

ctypedef void (*PetscDSBdJacobianFn)(
PetscInt, PetscInt, PetscInt,
const PetscInt*, const PetscInt*, const PetscScalar*, const PetscScalar*, const PetscScalar*,
const PetscInt*, const PetscInt*, const PetscScalar*, const PetscScalar*, const PetscScalar*,
PetscReal, PetscReal, const PetscReal*, const PetscReal*, PetscInt,
PetscReal, PetscReal, const PetscReal*, const PetscReal*, PetscInt,
const PetscScalar*, PetscScalar* )


Expand All @@ -47,4 +47,9 @@ cdef class PtrContainer:
cdef PetscDSBdResidualFn* fns_bd_residual
cdef PetscDSBdJacobianFn* fns_bd_jacobian


cpdef allocate(self, int n_res, int n_bcs, int n_jac, int n_bd_res, int n_bd_jac)
cpdef copy_residual_from(self, int dst, PtrContainer src, int src_idx)
cpdef copy_bcs_from(self, int dst, PtrContainer src, int src_idx)
cpdef copy_jacobian_from(self, int dst, PtrContainer src, int src_idx)
cpdef copy_bd_residual_from(self, int dst, PtrContainer src, int src_idx)
cpdef copy_bd_jacobian_from(self, int dst, PtrContainer src, int src_idx)
31 changes: 30 additions & 1 deletion src/underworld3/cython/petsc_types.pyx
Original file line number Diff line number Diff line change
@@ -1,2 +1,31 @@
from libc.stdlib cimport malloc

cdef class PtrContainer:
pass

cpdef allocate(self, int n_res, int n_bcs, int n_jac, int n_bd_res, int n_bd_jac):
"""Allocate function pointer arrays of the given sizes."""
self.fns_residual = <PetscDSResidualFn*> malloc(n_res * sizeof(PetscDSResidualFn))
self.fns_bcs = <PetscDSResidualFn*> malloc(n_bcs * sizeof(PetscDSResidualFn))
self.fns_jacobian = <PetscDSJacobianFn*> malloc(n_jac * sizeof(PetscDSJacobianFn))
self.fns_bd_residual = <PetscDSBdResidualFn*> malloc(n_bd_res * sizeof(PetscDSBdResidualFn))
self.fns_bd_jacobian = <PetscDSBdJacobianFn*> malloc(n_bd_jac * sizeof(PetscDSBdJacobianFn))
Comment on lines +5 to +11

Copilot AI Mar 25, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

PtrContainer.allocate() uses malloc() for several pointer arrays but PtrContainer has no dealloc (or equivalent) to free them. With the new per-function assembly path creating a fresh PtrContainer frequently, this becomes an unbounded native-memory leak. Add free() calls in a dealloc method (and consider freeing/replacing existing pointers inside allocate() when called more than once).

Copilot uses AI. Check for mistakes.

cpdef copy_residual_from(self, int dst, PtrContainer src, int src_idx):
"""Copy a residual function pointer from another container."""
self.fns_residual[dst] = src.fns_residual[src_idx]

cpdef copy_bcs_from(self, int dst, PtrContainer src, int src_idx):
"""Copy a BC function pointer from another container."""
self.fns_bcs[dst] = src.fns_bcs[src_idx]

cpdef copy_jacobian_from(self, int dst, PtrContainer src, int src_idx):
"""Copy a Jacobian function pointer from another container."""
self.fns_jacobian[dst] = src.fns_jacobian[src_idx]

cpdef copy_bd_residual_from(self, int dst, PtrContainer src, int src_idx):
"""Copy a boundary residual function pointer from another container."""
self.fns_bd_residual[dst] = src.fns_bd_residual[src_idx]

cpdef copy_bd_jacobian_from(self, int dst, PtrContainer src, int src_idx):
"""Copy a boundary Jacobian function pointer from another container."""
self.fns_bd_jacobian[dst] = src.fns_bd_jacobian[src_idx]
2 changes: 1 addition & 1 deletion src/underworld3/systems/solvers.py
Original file line number Diff line number Diff line change
Expand Up @@ -1275,7 +1275,7 @@ def solve(
timestep = self.delta_t.sym

if timestep != self.delta_t:
self._constitutive_model.Parameters.elastic_dt = timestep # this will force an initialisation because the functions need to be updated
self._constitutive_model.Parameters.dt_elastic = timestep # this will force an initialisation because the functions need to be updated

if _force_setup:
self.is_setup = False
Expand Down
Loading
Loading