compiler: Rework /petsc/iet and PETScArrays by ZoeLeibowitz · Pull Request #38 · ZoeLeibowitz/devito

ZoeLeibowitz · 2025-01-09T20:54:59Z

Refactors petsc/iet by introducing more classes to reduce the number of free functions, and enable flexibility to accommodate different types of solvers with varying setups, callbacks, and other configurations etc. (many free functions inside petsc/iet/passes.py have moved into new classes inside petsc/iet/routines.py)
Simplifies the PETScArray class, as it is consistently constructed from Function-like objects.
Ensure that a unique DMDA and matrix context struct are created for each SNES solve.

devito/petsc/iet/passes.py

devito/petsc/iet/routines.py

tests/test_petsc.py

mloubout

Mostly minor comments.

The iet layer could use more docstring and comments for maintainability but it can wait for later updates.

mloubout · 2025-01-13T12:52:23Z

devito/petsc/iet/passes.py

+    for iters, (injectsolve,) in injectsolve_mapper.items():

-    builder = PETScCallbackBuilder(**kwargs)
+        builder_classes = get_builder_classes(injectsolve)


Could make it an object or namedtuple so you can can have buillder.objbuilder and such, might make it easier to extend/maintain as well.

I've now created the Builder class inside petsc/iet/passes to encapsulate this. What do you think?

mloubout · 2025-01-13T12:53:33Z

devito/petsc/iet/passes.py

-        solver_objs = build_solver_objs(injectsolve, iters, **kwargs)
+        time_dep = dep(injectsolve, **kwargs)
+
+        solver_objs = ObjBuilder(dep=time_dep, **kwargs).build(injectsolve, iters)


They all seem to take the same input so you'd be able to just init the builder with all its components at once

mloubout · 2025-01-13T12:54:20Z

devito/petsc/iet/passes.py

-        subs.update({space_iter: List(body=runsolve)})
+        # Transform the spatial iteration loop with the calls to execute the solver
+        space_iter, = spatial_loop_nest(iters, injectsolve)
+        runsolve = SolverRun(dep=time_dep).runsolve(


personal nitpicking: i'd avoid using "run"

mloubout · 2025-01-13T12:57:36Z

devito/petsc/iet/passes.py

+        mapper = {solver_objs['dummyctx']: solver_objs['localctx']}
+        return Uxreplace(mapper).visit(efunc)
+
+    return {k: replace(v) for k, v in efuncs.items()}


why define an inner function instead of a plain loop

mapper = {} vistor = Uxreplace({solver_objs['dummyctx']: solver_objs['localctx']}) for k, v in efuncs.items(): mapper.update(vistor.visit(v)) return mapper

Either way Uxreplace(mapper) is the same for all efuncs so should only need to make it once

mloubout · 2025-01-13T12:58:06Z

devito/petsc/iet/routines.py

        obj = object.__new__(cls)
        obj.rcompile = rcompile
        obj.sregistry = sregistry
+        obj.concretize_mapper = kwargs.get('concretize_mapper', {})


So this is tied to the OSS PR right?

I've actually dropped it now since I use the same SymbolRegistry throughout, even with rcompile

mloubout · 2025-01-13T13:03:11Z

devito/petsc/iet/routines.py

+        }
+
+
+class SetupSolver:


Does this really need to be a class if it's stateless instead of three plain functions?

I've now added a hook for subclasses to add additional setup calls

mloubout · 2025-01-13T13:04:14Z

devito/petsc/iet/routines.py

+        if not time_iter:
+            return mapper
+
+        for d in time_iter[0].dimensions:


why only time_iter[0] ?

I've corrected this now

devito/petsc/iet/routines.py

mloubout · 2025-01-13T13:06:11Z

devito/petsc/types/array.py

+
+    @classmethod
+    def __indices_setup__(cls, *args, **kwargs):
+        dimensions = kwargs['target'].space_dimensions


why not just dimensions ?

PETScArray objects correspond to PETSc Vec objects, which represent the spatial domain and are used at each time step.

mloubout · 2025-01-13T13:07:32Z

devito/petsc/types/types.py

        'ksp_type': 'gmres',
        'pc_type': 'jacobi',
-        'ksp_rtol': 1e-7,  # Relative tolerance
+        'ksp_rtol': 1e-5,  # Relative tolerance


These need to be documented in docstring

FabioLuporini · 2025-01-29T08:46:26Z

devito/petsc/solve.py

+
    matvecs, formfuncs, formrhs = (
-        [eq.subs(time_mapper) for eq in lst] for lst in (matvecs, formfuncs, formrhs)
+        [eq.xreplace(time_mapper) for eq in lst] for lst in (matvecs, formfuncs, formrhs)


out of curiosity why this change?

I think both work fine. I can't remember exactly why I changed it. Should it stay as subs?

not worth another spin of CI imho -- up to you :)

xreplace is what introduced a bug with custom coeff legacy API, it sometimes only replaces the first occurrence it finds instead of all so if there is more than one occurrence it's likely to break

Oh ok, will change it back in my next PR.

georgebisbas · 2025-01-29T15:38:03Z

devito/petsc/iet/routines.py

-        calls = [call_struct_callback] + calls_set_app_ctx
+    def _spatial_loop_nest(self, iters, injectsolve):
+        spatial_body = []
+        for tree in retrieve_iteration_tree(iters[0]):


What is special about iters[0] here?

georgebisbas · 2025-01-29T15:42:05Z

devito/petsc/types/object.py

+        These are not set in the `PopulateMatContext` callback.
+        """
        return [f for f in self.fields
                if isinstance(f, (ModuloDimension, TimeDimension))]


if isinstance(f, (ModuloDimension, TimeDimension))]
could be replaced with
if f.is_Time
?

EdCaunt · 2025-02-03T13:54:19Z

devito/petsc/iet/routines.py

-        matvec_callback, formfunc_callback, formrhs_callback = self.make_all(
-            injectsolve, objs, solver_objs
-        )
+    @property


Nitpick: should be cached_property

EdCaunt · 2025-02-03T14:02:41Z

devito/petsc/iet/routines.py

+
+        # Global dimensions
+        args.extend(list(grid.shape)[::-1])
+        # No.of processors in each dimension


Nitpick: missing space in comment

EdCaunt · 2025-02-03T14:04:10Z

devito/petsc/iet/routines.py

+    - Modulo dimensions are updated in the matrix context struct at each time
+      step and can be accessed in the callback functions where needed.
+    """
+    @property


Nitpick: should probably be cached_property, but tbh is so cheap it doesn't really matter

EdCaunt · 2025-02-03T14:04:32Z

devito/petsc/iet/routines.py

-    }
-    subs = {symb: true_dims[time_mapper[symb]] for symb in time_mapper}
-    return Uxreplace(subs).visit(body)
+    @property


Nitpick: should use cached_property

EdCaunt · 2025-02-03T14:04:41Z

devito/petsc/iet/routines.py

+        target_time = target_time.pop()
+        return target_time
+
+    @property


EdCaunt · 2025-02-03T14:10:57Z

devito/petsc/types/array.py

@@ -107,11 +115,6 @@ def _C_ctype(self):
    @property


EdCaunt · 2025-02-03T14:11:14Z

devito/petsc/types/array.py

    def symbolic_shape(self):
        field_from_composites = [
-            FieldFromComposite('g%sm' % d.name, self.dmda.info) for d in self.dimensions]
+            FieldFromComposite('g%sm' % d.name, self.localinfo) for d in self.dimensions]


Nitpick: fstring this?

EdCaunt · 2025-02-03T14:14:50Z

devito/petsc/types/object.py

@@ -160,9 +182,21 @@ def fields(self):

    @property


Nitpick: probably should be cached

EdCaunt · 2025-02-03T14:14:58Z

devito/petsc/types/object.py

        return [f for f in self.fields
                if isinstance(f, (ModuloDimension, TimeDimension))]

+    @property


JDBetteridge · 2025-02-05T14:24:49Z

devito/petsc/iet/routines.py

-
-    def create_formrhs_body(self, injectsolve, body, solver_objs, objs):
-        linsolveexpr = injectsolve.expr.rhs
+    def _create_formrhs_body(self, injectsolve, body, solver_objs, objs):


Suggested change

def _create_formrhs_body(self, injectsolve, body, solver_objs, objs):

def _create_form_rhs_body(self, injectsolve, body, solver_objs, objs):

?

ZoeLeibowitz marked this pull request as ready for review January 10, 2025 08:13

ZoeLeibowitz requested review from EdCaunt, FabioLuporini, georgebisbas and mloubout January 10, 2025 08:18

EdCaunt reviewed Jan 10, 2025

View reviewed changes

ZoeLeibowitz force-pushed the cleanup2 branch from 841a8da to 727e77c Compare January 12, 2025 14:02

mloubout reviewed Jan 13, 2025

View reviewed changes

FabioLuporini reviewed Jan 29, 2025

View reviewed changes

FabioLuporini approved these changes Jan 29, 2025

View reviewed changes

georgebisbas reviewed Jan 29, 2025

View reviewed changes

ZoeLeibowitz requested a review from JDBetteridge January 31, 2025 11:08

EdCaunt reviewed Feb 3, 2025

View reviewed changes

JDBetteridge reviewed Feb 11, 2025

View reviewed changes

ZoeLeibowitz force-pushed the FieldFromPointer branch from 64870de to 579045f Compare February 18, 2025 13:33

FabioLuporini and others added 15 commits February 18, 2025 15:55

misc: Split Dockerfile.cpu into .cpu and .intel

cfa7218

misc: Update file due to incorrect rebase

fd8c53f

misc: Split Dockerfile.cpu into .cpu and .intel

d09e654

misc: Update file due to incorrect rebase

62a4e1c

dsl: Dispatch to sympy.Add not both Add and EvalDerivative

0f862ba

compiler: Create mixin for ThreadCallable and PETScCallable

b8de614

types: Edit freesymbols in FieldFromPointer

82b434b

compiler: Use classes

8cb18d3

compiler: Switch to using a separate struct and dm for each SNES solve

da2bbf3

compiler: Rework petscarrays

c445140

misc: Clean up

670e1d2

compiler: Timedep abstraction

ad5fb49

More cleanup, add callbackdm etc

89cc06e

Address some comments

2c4a568

compiler: Move uxreplace_efuncs into Builder

b83ec92

ZoeLeibowitz added 3 commits February 18, 2025 16:17

misc: Clean up

4d75af8

misc: Add timeloop abstraction docs

5f25b0f

misc: Clean up and docs

37b62b3

ZoeLeibowitz force-pushed the cleanup2 branch from eec492e to 0106f9a Compare February 18, 2025 16:35

ZoeLeibowitz added 2 commits February 18, 2025 17:51

misc: Spellings

10af16f

flake8

ef7ee34

ZoeLeibowitz force-pushed the cleanup2 branch from e26e832 to ef7ee34 Compare February 18, 2025 17:51

ZoeLeibowitz merged commit 3b80b8b into FieldFromPointer Feb 18, 2025
29 checks passed

	def _create_formrhs_body(self, injectsolve, body, solver_objs, objs):
	def _create_form_rhs_body(self, injectsolve, body, solver_objs, objs):

Conversation

ZoeLeibowitz commented Jan 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mloubout left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ZoeLeibowitz commented Jan 9, 2025 •

edited

Loading