compiler: Drop petscgetargs callback and add more functions to logging/get_info

Author: ZoeLeibowitz

devito/ir/iet/visitors.py

@@ -292,6 +292,7 @@ def _gen_struct_decl(self, obj, masked=()):
         for i, (n, ct) in zip(fields, ctype._fields_):
             try:
                 entries.append(self._gen_value(i, 0, masked=('const',)))
+                
             except AttributeError:
                 cstr = self.ccode(ct)
                 if ct is c_restrict_void_p:

devito/petsc/iet/logging.py

@@ -27,6 +27,7 @@ def __init__(self, level, get_info=None, **kwargs):
                 'kspgetiterationnumber',
                 'kspgettolerances',
                 'kspgetconvergedreason',
+                'kspgettype',
                 # SNES specific
                 'snesgetiterationnumber',
             ]

devito/petsc/iet/passes.py

@@ -6,7 +6,7 @@
 from devito.ir.iet import (Transformer, MapNodes, Iteration, BlankLine,
                            DummyExpr, CallableBody, List, Call, Callable,
                            FindNodes, Section)
-from devito.symbolics import Byref, FieldFromPointer, Macro
+from devito.symbolics import Byref, FieldFromPointer, Macro, Null
 from devito.types import Symbol, Scalar
 from devito.types.basic import DataSymbol
 from devito.tools import frozendict
@@ -16,9 +16,8 @@
                                 PointerIS, Mat, CallbackVec, Vec, CallbackMat, SNES,
                                 DummyArg, PetscInt, PointerDM, PointerMat, MatReuse,
                                 CallbackPointerIS, CallbackPointerDM, JacobianStruct,
-                                SubMatrixStruct, Initialize, Finalize, ArgvSymbol,
-                                GetArgs, ArgvSymbolPtr, ArgcPtr)
-from devito.petsc.types.macros import petsc_func_begin_user, Null
+                                SubMatrixStruct, Initialize, Finalize, ArgvSymbol)
+from devito.petsc.types.macros import petsc_func_begin_user
 from devito.petsc.iet.nodes import PetscMetaData
 from devito.petsc.utils import core_metadata, petsc_languages
 from devito.petsc.iet.routines import (CBBuilder, CCBBuilder, BaseObjectBuilder,
@@ -52,9 +51,6 @@ def lower_petsc(iet, **kwargs):
     if any(filter(lambda i: isinstance(i.expr.rhs, Finalize), data)):
         return finalize(iet), core_metadata()
 
-    if any(filter(lambda i: isinstance(i.expr.rhs, GetArgs), data)):
-        return get_args(iet), core_metadata()
-
     unique_grids = {i.expr.rhs.grid for (i,) in inject_solve_mapper.values()}
     # Assumption is that all solves are on the same grid
     if len(unique_grids) > 1:
@@ -138,18 +134,6 @@ def finalize(iet):
     return iet._rebuild(body=finalize_body)
 
 
-def get_args(iet):
-    argc = ArgcPtr(name='argc', dtype=np.int32)
-    argv = ArgvSymbolPtr(name='argv')
-
-    body = petsc_call('PetscGetArgs', [argc, argv])
-    body = CallableBody(
-        body=(petsc_func_begin_user, body),
-        retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
-    )
-    return iet._rebuild(body=body)
-
-
 def make_core_petsc_calls(objs, comm):
     call_mpi = petsc_call_mpi('MPI_Comm_size', [comm, Byref(objs['size'])])
     return call_mpi, BlankLine

devito/petsc/iet/routines.py

@@ -7,7 +7,7 @@
                            retrieve_iteration_tree, filter_iterations, Iteration,
                            PointerCast)
 from devito.symbolics import (Byref, FieldFromPointer, cast, VOID,
-                              FieldFromComposite, IntDiv, Deref, Mod, String)
+                              FieldFromComposite, IntDiv, Deref, Mod, String, Null)
 from devito.symbolics.unevaluation import Mul
 from devito.types.basic import AbstractFunction
 from devito.types import Temp, Dimension
@@ -21,8 +21,7 @@
                                 KSP, PC, SNES, PetscInt, StartPtr, PointerIS, PointerDM,
                                 VecScatter, DMCast, JacobianStruct, SubMatrixStruct,
                                 CallbackDM)
-from devito.petsc.types.macros import petsc_func_begin_user, Null
-# from devito.petsc.initialize import PetscGetArgs
+from devito.petsc.types.macros import petsc_func_begin_user
 
 
 class CBBuilder:
@@ -150,10 +149,10 @@ def _make_options_callback(self):
         for k, v in params.items():
             option = f'-{prefix}{k}'
 
-            # TODO: drop the global variable _petsc_clargs..
-            # from devito.petsc.initialize import PetscGetArgs
-            # PetscGetArgs()
-
+            # TODO: Revisit use of a global variable here.
+            # Consider replacing this with a call to `PetscGetArgs`, though
+            # initial attempts failed, possibly because the argv pointer is
+            # created in Python?..
             import devito.petsc.initialize
             if option in devito.petsc.initialize._petsc_clargs:
                 # Ensures that the command line args take priority

devito/petsc/initialize.py

@@ -1,12 +1,12 @@
 import os
 import sys
-from ctypes import POINTER, cast, c_char, c_int, c_char_p, byref
+from ctypes import POINTER, cast, c_char
 import atexit
 
 from devito import Operator, switchconfig
 from devito.types import Symbol
 from devito.types.equation import PetscEq
-from devito.petsc.types import Initialize, Finalize, GetArgs
+from devito.petsc.types import Initialize, Finalize
 
 global _petsc_initialized
 _petsc_initialized = False
@@ -26,7 +26,6 @@ def PetscInitialize(clargs=sys.argv):
         if clargs is not sys.argv:
             clargs = [sys.argv[0], *clargs]
 
-        # TODO: Drop this global variable
         _petsc_clargs = clargs
 
         # TODO: Potentially just use cgen + the compiler machinery in Devito
@@ -43,24 +42,19 @@ def PetscInitialize(clargs=sys.argv):
                 name='kernel_finalize', opt='noop'
             )
 
+        # Convert each string to a bytes object (e.g: '-ksp_type' -> b'-ksp_type')
         # `argv_bytes` must be a list so the memory address persists
         # `os.fsencode` should be preferred over `string().encode('utf-8')`
         # in case there is some system specific encoding in use
         argv_bytes = list(map(os.fsencode, clargs))
+
+        # POINTER(c_char) is equivalent to char * in C
+        # (POINTER(c_char) * len(clargs)) creates a C array type: char *[len(clargs)]
+        # Instantiating it with (*map(...)) casts each bytes object to a char * and
+        # fills the array. The result is a char *argv[]
         argv_pointer = (POINTER(c_char)*len(clargs))(
             *map(lambda s: cast(s, POINTER(c_char)), argv_bytes)
         )
         op_init.apply(argc=len(clargs), argv=argv_pointer)
         atexit.register(op_finalize.apply)
         _petsc_initialized = True
-
-
-def PetscGetArgs():
-    with switchconfig(language='petsc'):
-        op_get_args = Operator(
-            [PetscEq(dummy, GetArgs(dummy))],
-            name='kernel_get_args', opt='noop'
-        )
-    argc_ptr = c_int()
-    argv_ptr = (POINTER(c_char_p))()
-    op_get_args.apply(argc=byref(argc_ptr), argv=byref(argv_ptr))

devito/petsc/logging.py

@@ -1,10 +1,12 @@
 from collections import namedtuple, OrderedDict
 from dataclasses import dataclass
+from functools import cached_property
+from cgen import Struct, Value
 
 from devito.types import CompositeObject
 
-from devito.petsc.types import PetscInt, PetscScalar
-from devito.petsc.utils import petsc_type_mappings
+from devito.petsc.types import PetscInt, PetscScalar, KSPType
+from devito.petsc.utils import petsc_type_mappings, fixed_petsc_type_mappings
 
 
 class PetscEntry:
@@ -140,13 +142,26 @@ def __init__(self, name, pname, petsc_option_mapper, sobjs, section_mapper,
 
         mapper = {v: k for k, v in petsc_type_mappings.items()}
 
-        fields = []
+        pfields = []
+        self._tmp_fields = []
+
+        # obj_mapper is e.g {'kspits': kspits0}
+        # TODO: think this can just be for (x,y) in ....
         for obj_mapper in petsc_option_mapper.values():
             for petsc_option in obj_mapper.values():
-                ctype = mapper[str(petsc_option._C_ctype)]
-                fields.append((petsc_option.name, ctype))
+                self._tmp_fields.append(petsc_option)
+                # petsc_type is e.g. 'PetscInt', 'PetscScalar', 'KSPType'
+                petsc_type = str(petsc_option.dtype)
+                if petsc_type in mapper:
+                    ctype = mapper[petsc_type]
+                else:
+                    ctype = fixed_petsc_type_mappings[petsc_type] 
+                pfields.append((petsc_option.name, ctype))
+        super().__init__(name, pname, pfields)
 
-        super().__init__(name, pname, fields)
+    @property
+    def fields(self):
+        return self._tmp_fields
 
     @property
     def section(self):
@@ -177,9 +192,9 @@ def __getattr__(self, attr):
         # - If the function returns multiple values (e.g., KSPGetTolerances),
         #   return a dictionary mapping each output name to its value,
         #   e.g., {'rtol': val0, 'atol': val1, ...}.
-        if len(obj_mapper) == 1:
-            return get_val(next(iter(obj_mapper.values())))
-        return {k: get_val(v) for k, v in obj_mapper.items()}
+        if len(obj_mapper) > 1:
+            return {k: get_val(v) for k, v in obj_mapper.items()}
+        return get_val(next(iter(obj_mapper.values())))
 
 
 @dataclass
@@ -217,6 +232,12 @@ class PetscReturnVariable:
         input_params='ksp',
         output_param=('reason',),
     ),
+    'kspgettype': PetscReturnVariable(
+        name='KSPGetType',
+        variable_type=(KSPType,),
+        input_params='ksp',
+        output_param=('ksptype',),
+    ),
     # SNES specific
     'snesgetiterationnumber': PetscReturnVariable(
         name='SNESGetIterationNumber',

devito/petsc/solve.py

@@ -46,13 +46,35 @@ def PETScSolve(target_exprs, target=None, solver_parameters=None,
     solver_parameters : dict, optional
         PETSc solver options.
 
+        Descriptions (not exhaustive):
+        - 'ksp_type':    Specifies the Krylov method (e.g., 'gmres', 'cg').
+        - 'pc_type':     Specifies the preconditioner type (e.g., ...).
+        - 'ksp_rtol':    Relative convergence tolerance for KSP solvers.
+        - 'ksp_atol':    Absolute convergence tolerance for KSP solvers.
+        - 'ksp_divtol':  Divergence tolerance, amount residual norm can increase before
+                        `KSPConvergedDefault()` concludes that the method is diverging.
+        - 'ksp_max_it':  Maximum number of KSP iterations to use.
+        - 'snes_type':   Type of SNES solver; 'ksponly' is used for linear solves.
+
+        References:
+        - KSP types: https://petsc.org/main/manualpages/KSP/KSPType/
+        - PC types: https://petsc.org/main/manualpages/PC/PCType/
+        - KSP tolerances: https://petsc.org/main/manualpages/KSP/KSPSetTolerances/
+        - SNES type: https://petsc.org/main/manualpages/SNES/SNESType/
+
     options_prefix : str, optional
         Prefix for the solver, used to configure options via the command line. If not
         provided, a default prefix is generated by Devito.
 
     get_info : list[str], optional
         A list of PETSc API functions to collect statistics from the solver.
-        For example, `['kspgetiterationnumber', 'kspgettolerances']`.
+        For example, `['kspgetiterationnumber', 'kspgettolerances']`. 
+        Capitalisation does not matter; e.g. `'KSPGetIterationNumber'` and
+        `'kspgetiterationnumber'` are treated the same.
+
+        List of available functions:
+        - ['kspgetiterationnumber', 'kspgettolerances', 'kspgetconvergedreason',
+           'snesgetiterationnumber']
 
     Returns
     -------
@@ -76,7 +98,7 @@ def __init__(self, solver_parameters=None, target_exprs=None, options_prefix=Non
         self.target_exprs = target_exprs
         self.user_prefix = options_prefix
         self.formatted_prefix = format_options_prefix(options_prefix)
-        self.get_info = get_info
+        self.get_info = [f.lower() for f in get_info]
 
     def build_expr(self):
         target, funcs, field_data = self.linear_solve_args()

devito/petsc/solver_parameters.py

@@ -1,24 +1,6 @@
 from petsctools import flatten_parameters
 import itertools
 
-"""
-Parameter descriptions:
-- 'ksp_type':    Specifies the Krylov method (e.g., 'gmres', 'cg').
-- 'pc_type':     Specifies the preconditioner type (e.g., ...).
-- 'ksp_rtol':    Relative convergence tolerance for KSP solvers.
-- 'ksp_atol':    Absolute convergence tolerance for KSP solvers.
-- 'ksp_divtol':  Divergence tolerance, amount residual norm can increase before
-                `KSPConvergedDefault()` concludes that the method is diverging.
-- 'ksp_max_it':  Maximum number of KSP iterations to use.
-- 'snes_type':   Type of SNES solver; 'ksponly' is used for linear solves.
-
-References:
-- KSP types: https://petsc.org/main/manualpages/KSP/KSPType/
-- PC types: https://petsc.org/main/manualpages/PC/PCType/
-- KSP tolerances: https://petsc.org/main/manualpages/KSP/KSPSetTolerances/
-- SNES type: https://petsc.org/main/manualpages/SNES/SNESType/
-"""
-
 
 # NOTE: Will be extended, the default preconditioner is not going to be 'none'
 base_solve_defaults = {
@@ -51,9 +33,9 @@ def linear_solver_parameters(solver_parameters):
 
 def format_options_prefix(options_prefix):
     # NOTE: Modified from the `OptionsManager` inside petsctools
-    if options_prefix is None:
+    if not options_prefix:
         options_prefix = f"devito_{next(_options_prefix_counter)}_"
     else:
-        if len(options_prefix) and not options_prefix.endswith("_"):
+        if not options_prefix.endswith("_"):
             options_prefix += "_"
     return options_prefix

devito/petsc/types/macros.py

@@ -1,8 +1,5 @@
 import cgen as c
-from devito.symbolics import Macro
 
 
-Null = Macro('NULL')
-
 # TODO: Don't use c.Line here?
 petsc_func_begin_user = c.Line('PetscFunctionBeginUser;')

devito/petsc/types/object.py

@@ -1,4 +1,4 @@
-from ctypes import POINTER, c_char, c_char_p, c_int
+from ctypes import POINTER, c_char, c_char_p
 
 from devito.tools import CustomDtype, dtype_to_ctype, as_tuple, CustomIntType
 from devito.types import (LocalObject, LocalCompositeObject, ModuloDimension,
@@ -113,6 +113,10 @@ class KSP(PetscObject):
     dtype = CustomDtype('KSP')
 
 
+class KSPType(PetscObject):
+    dtype = CustomDtype('KSPType')
+
+
 class CallbackSNES(PetscObject):
     """
     PETSc SNES : Non-Linear Systems Solvers.
@@ -305,22 +309,10 @@ def _C_ctype(self):
         return POINTER(POINTER(c_char))
 
 
-class ArgvSymbolPtr(DataSymbol):
-    @property
-    def _C_ctype(self):
-        return POINTER(POINTER(c_char_p))
-
-
-class ArgcPtr(DataSymbol):
-    @property
-    def _C_ctype(self):
-        return POINTER(c_int)
-
-
-class CharPtr(DataSymbol):
-    @property
-    def _C_ctype(self):
-        return c_char_p
+# class CharPtr(DataSymbol):
+#     @property
+#     def _C_ctype(self):
+#         return c_char_p
 
 
 FREE_PRIORITY = {