compiler: Start extending the PetscSummary

Author: ZoeLeibowitz

devito/petsc/iet/logging.py

@@ -23,32 +23,54 @@ def __init__(self, level, **kwargs):
         if level <= PERF:
             self.function_list.extend([
                 'kspgetiterationnumber',
-                'snesgetiterationnumber'
+                'snesgetiterationnumber',
+                'kspgettolerances'
             ])
 
         # TODO: To be extended with if level <= DEBUG: ...
 
+        # from IPython import embed; embed()
+        # if str(self.inject_solve.expr.rhs.solver_parameters['ksp_rtol']) == '1e-15':
+
         name = self.sreg.make_name(prefix='petscinfo')
         pname = self.sreg.make_name(prefix='petscprofiler')
 
         self.statstruct = PetscInfo(
-            name, pname, self.logobjs, self.sobjs,
+            name, pname, self.petsc_option_mapper, self.sobjs,
             self.section_mapper, self.inject_solve,
             self.function_list
         )
+        # else:
+        #     name = self.sreg.make_name(prefix='petscinfooo')
+        #     pname = self.sreg.make_name(prefix='petscprofilerrrr')
+
+        #     self.statstruct = PetscInfo(
+        #         name, pname, self.petsc_option_mapper, self.sobjs,
+        #         self.section_mapper, self.inject_solve,
+        #         self.function_list
+        #     )
+
+        # from IPython import embed; embed()  # noqa: E402
+
+    # @property
+    # def statstruct(self):
+    #     return self._statstruct
 
     @cached_property
-    def logobjs(self):
+    def petsc_option_mapper(self):
         """
         Create PETSc objects specifically needed for logging solver statistics.
+
+        ADD EXTENDED DOCSTRING 
         """
-        return {
-            info.name: info.variable_type(
-                self.sreg.make_name(prefix=info.output_param)
-            )
-            for func_name in self.function_list
-            for info in [petsc_return_variable_dict[func_name]]
-        }
+        opts = {}
+        for func_name in self.function_list:
+            info = petsc_return_variable_dict[func_name]
+            opts[info.name] = {}
+            for vtype, out in zip(info.variable_type, info.output_param, strict=True):
+                opts[info.name][out] = vtype(self.sreg.make_name(prefix=out))
+        
+        return opts
 
     @cached_property
     def calls(self):
@@ -58,20 +80,21 @@ def calls(self):
         """
         struct = self.statstruct
         calls = []
-        for param in self.function_list:
-            param = petsc_return_variable_dict[param]
-
-            inputs = []
-            for i in param.input_params:
-                inputs.append(self.sobjs[i])
-
-            logobj = self.logobjs[param.name]
+        for func_name in self.function_list:
+            return_variable = petsc_return_variable_dict[func_name]
 
+            input = self.sobjs[return_variable.input_params]
+            output_params = self.petsc_option_mapper[return_variable.name].values()
+            outputs = [Byref(i) for i in output_params]
+            # from IPython import embed; embed()
             calls.append(
-                petsc_call(param.name, inputs + [Byref(logobj)])
+                petsc_call(return_variable.name, [input] + outputs)
             )
             # TODO: Perform a PetscCIntCast here?
-            expr = DummyExpr(FieldFromPointer(logobj._C_symbol, struct), logobj._C_symbol)
-            calls.append(expr)
+            exprs = [
+                DummyExpr(FieldFromPointer(i._C_symbol, struct), i._C_symbol)
+                for i in output_params
+            ]
+            calls.extend(exprs)
 
         return tuple(calls)

devito/petsc/iet/passes.py

@@ -73,7 +73,13 @@ def lower_petsc(iet, **kwargs):
 
     # Generate a shared callback used by all PETScSolve instances to set
     # individual PetscOptions
-    set_solver_option(efuncs)
+    # set_solver_option(efuncs)
+    # List of all callbacks that clear PetscOptions
+    # from IPython import embed; embed()
+    # TODO: throw a warning/error if the user passes a solver in with the same options_prefix
+    # it's going to lead to weird solver option behaviour. Note, if you use the options_prefix across
+    # different Operator runs, it will not be an issue
+    clear_options = []
 
     for iters, (inject_solve,) in inject_solve_mapper.items():
 
@@ -86,11 +92,17 @@ def lower_petsc(iet, **kwargs):
 
         efuncs.update(builder.cbbuilder.efuncs)
 
+        # clear_options.append(builder.cbbuilder._clear_options_efunc)
+        # from IPython import embed; embed()  # noqa: E402
+        clear_options.extend((petsc_call(
+            builder.cbbuilder._clear_options_efunc.name, []
+        ),))
+
     populate_matrix_context(efuncs)
 
     iet = Transformer(subs).visit(iet)
-
-    body = core + tuple(setup) + iet.body.body
+    # from IPython import embed; embed()
+    body = core + tuple(setup) + iet.body.body + tuple(clear_options)
     body = iet.body._rebuild(body=body)
     iet = iet._rebuild(body=body)
     metadata = {**core_metadata(), 'efuncs': tuple(efuncs.values())}

devito/petsc/iet/routines.py

@@ -1,3 +1,4 @@
+import sys
 from collections import OrderedDict
 from functools import cached_property
 import math
@@ -42,7 +43,8 @@ def __init__(self, **kwargs):
         self._efuncs = OrderedDict()
         self._struct_params = []
 
-        self._options_efunc = None
+        self._set_options_efunc = None
+        self._clear_options_efunc = None
         self._main_matvec_callback = None
         self._user_struct_callback = None
         self._F_efunc = None
@@ -120,26 +122,48 @@ def _make_options_callback(self):
         params = self.solver_parameters
         prefix = self.inject_solve.expr.rhs.formatted_prefix
 
-        body = []
+        set_body = []
+        clear_body = []
+
         for k, v in params.items():
-            option_key = String(f"-{prefix}{k}")
+            option = f'-{prefix}{k}'
+            if option in sys.argv:
+                # Ensures that the command line options take priority
+                continue
+            option_name = String(option)
             option_value = Null if v is None else String(str(v))
-            body.append(petsc_call('SetPetscOption', [option_key, option_value]))
+            set_body.append(petsc_call('PetscOptionsSetValue', [Null, option_name, option_value]))
+            clear_body.append(petsc_call('PetscOptionsClearValue', [Null, option_name]))
 
-        body = CallableBody(
-            List(body=body),
+        set_body = CallableBody(
+            List(body=set_body),
             init=(petsc_func_begin_user,),
             retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
         )
 
-        cb = PETScCallable(
+        clear_body = CallableBody(
+            List(body=clear_body),
+            init=(petsc_func_begin_user,),
+            retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
+        )
+
+        set_callback = PETScCallable(
             self.sregistry.make_name(prefix='SetPetscOptions'),
-            body,
+            set_body,
             retval=objs['err'],
             parameters=()
         )
-        self._options_efunc = cb
-        self._efuncs[cb.name] = cb
+
+        clear_callback = PETScCallable(
+            self.sregistry.make_name(prefix='ClearPetscOptions'),
+            clear_body,
+            retval=objs['err'],
+            parameters=()
+        )
+        self._set_options_efunc = set_callback
+        self._efuncs[set_callback.name] = set_callback
+        self._clear_options_efunc = clear_callback
+        self._efuncs[clear_callback.name] = clear_callback
 
     def _make_matvec(self, jacobian, prefix='MatMult'):
         # Compile `matvecs` into an IET via recursive compilation
@@ -1264,7 +1288,7 @@ def _setup(self):
         ) if self.formatted_prefix else None
 
         set_options = petsc_call(
-            self.cbbuilder._options_efunc.name, []
+            self.cbbuilder._set_options_efunc.name, []
         )
 
         snes_set_dm = petsc_call('SNESSetDM', [sobjs['snes'], dmda])
@@ -1421,7 +1445,7 @@ def _setup(self):
         ) if self.formatted_prefix else None
 
         set_options = petsc_call(
-            self.cbbuilder._options_efunc.name, []
+            self.cbbuilder._set_options_efunc.name, []
         )
 
         snes_set_dm = petsc_call('SNESSetDM', [sobjs['snes'], dmda])

devito/petsc/logging.py

@@ -3,7 +3,7 @@
 
 from devito.types import CompositeObject
 
-from devito.petsc.types import PetscInt
+from devito.petsc.types import PetscInt, PetscScalar
 from devito.petsc.utils import petsc_type_mappings
 
 
@@ -63,7 +63,9 @@ def petsc_entry(self, petscinfo):
         containing the values for each PETSc function call.
         """
         funcs = self._functions
+        # from IPython import embed; embed()
         values = tuple(getattr(petscinfo, c) for c in funcs)
+        # from IPython import embed; embed()
         return PetscEntry(**{k: v for k, v in zip(funcs, values)})
 
     def _add_properties(self):
@@ -119,20 +121,28 @@ def __getitem__(self, key):
 
 class PetscInfo(CompositeObject):
 
-    __rargs__ = ('name', 'pname', 'logobjs', 'sobjs', 'section_mapper',
+    __rargs__ = ('name', 'pname', 'petsc_option_mapper', 'sobjs', 'section_mapper',
                  'inject_solve', 'function_list')
 
-    def __init__(self, name, pname, logobjs, sobjs, section_mapper,
+    def __init__(self, name, pname, petsc_option_mapper, sobjs, section_mapper,
                  inject_solve, function_list):
 
-        self.logobjs = logobjs
+        # TODO: change name to match new name elsewehere
+        self.petsc_option_mapper = petsc_option_mapper
         self.sobjs = sobjs
         self.section_mapper = section_mapper
         self.inject_solve = inject_solve
         self.function_list = function_list
 
         mapper = {v: k for k, v in petsc_type_mappings.items()}
-        fields = [(str(i), mapper[str(i._C_ctype)]) for i in logobjs.values()]
+
+        self.formatted_prefix = inject_solve.expr.rhs.formatted_prefix
+
+        fields = [
+            (str(ptype), mapper[str(ptype._C_ctype)])
+            for option in petsc_option_mapper.values() for ptype in option.values()
+        ]
+
         super().__init__(name, pname, fields)
 
     @property
@@ -143,20 +153,53 @@ def section(self):
     @property
     def summary_key(self):
         user_prefix = self.inject_solve.expr.rhs.user_prefix
+        # TODO: this will be the case when using the default options prefix provided by Devito
+        # if user_prefix is None:
+        #     user_prefix = self.formatted_prefix
         return (self.section, user_prefix)
 
     def __getattr__(self, attr):
-        if attr in self.logobjs.keys():
-            return getattr(self.value._obj, self.logobjs[attr].name)
+        if attr in self.petsc_option_mapper.keys():
+            if len(self.petsc_option_mapper[attr].values()) > 1:
+                tmp = {}
+                for i, j in self.petsc_option_mapper[attr].items():
+                    tmp2 = getattr(self.value._obj, j.name)
+                    tmp[i] = tmp2
+                return tmp
+            else:
+                # TODO: CLEANNN
+                first_val = list(self.petsc_option_mapper[attr].values())[0]
+                return getattr(self.value._obj, first_val.name)
+        # from IPython import embed; embed()  # noqa: E402
         raise AttributeError(f"{attr} not found in PETSc return variables")
-
-
+    
+    # TODO: maybe just overrider _hashable_content??
+
+    # def _hashable_content(self):
+    #     """
+    #     Return a tuple of the formatted prefix and section for hashing.
+    #     This is used to ensure that two PetscInfo objects with the same
+    #     formatted prefix and section are considered equal.
+    #     """
+    #     return (self.name, self.dtype, self.inject_solve.expr.rhs)
+    
+    # def __eq__(self, other):
+    #     if not isinstance(other, PetscInfo):
+    #         return NotImplemented
+    #     # return self.formatted_prefix == other.formatted_prefix
+    #     return self.inject_solve.expr.rhs == other.inject_solve.expr.rhs
+
+    # def __hash__(self):
+    #     return hash(self.inject_solve.expr.rhs)
+
+
+# TODO: change the lists to tuples
 @dataclass
 class PetscReturnVariable:
     name: str
-    variable_type: None
-    input_params: list
-    output_param: str
+    variable_type: list
+    input_params: str
+    output_param: list[str]
 
 
 # NOTE:
@@ -168,14 +211,20 @@ class PetscReturnVariable:
 petsc_return_variable_dict = {
     'kspgetiterationnumber': PetscReturnVariable(
         name='KSPGetIterationNumber',
-        variable_type=PetscInt,
-        input_params=['ksp'],
-        output_param='kspiter'
+        variable_type=[PetscInt],
+        input_params='ksp',
+        output_param=['kspiter']
     ),
     'snesgetiterationnumber': PetscReturnVariable(
         name='SNESGetIterationNumber',
-        variable_type=PetscInt,
-        input_params=['snes'],
-        output_param='snesiter',
+        variable_type=[PetscInt],
+        input_params='snes',
+        output_param=['snesiter'],
+    ),
+    'kspgettolerances': PetscReturnVariable(
+        name='KSPGetTolerances',
+        variable_type=[PetscScalar, PetscScalar, PetscScalar, PetscInt],
+        input_params='ksp',
+        output_param=['rtol', 'abstol', 'dtol', 'maxits'],
     )
 }