compiler: Create dummy op for petscgetargs

Author: ZoeLeibowitz

devito/petsc/iet/passes.py

@@ -16,7 +16,8 @@
                                 PointerIS, Mat, CallbackVec, Vec, CallbackMat, SNES,
                                 DummyArg, PetscInt, PointerDM, PointerMat, MatReuse,
                                 CallbackPointerIS, CallbackPointerDM, JacobianStruct,
-                                SubMatrixStruct, Initialize, Finalize, ArgvSymbol)
+                                SubMatrixStruct, Initialize, Finalize, ArgvSymbol,
+                                GetArgs, ArgvSymbolPtr, ArgcPtr)
 from devito.petsc.types.macros import petsc_func_begin_user, Null
 from devito.petsc.iet.nodes import PetscMetaData
 from devito.petsc.utils import core_metadata, petsc_languages
@@ -51,6 +52,9 @@ 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:
@@ -73,8 +77,6 @@ def lower_petsc(iet, **kwargs):
     prefixes = [d.expr.rhs.user_prefix for d in data if d.expr.rhs.user_prefix]
     duplicates = {p for p in prefixes if prefixes.count(p) > 1}
 
-    # TODO: Avoid the other exception raised - think due to exception being
-    # raised inside the @iet_pass?
     if duplicates:
         dup_list = ", ".join(repr(p) for p in sorted(duplicates))
         raise ValueError(
@@ -136,6 +138,18 @@ 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

@@ -1,4 +1,3 @@
-import sys
 from collections import OrderedDict
 from functools import cached_property
 import math
@@ -23,6 +22,7 @@
                                 VecScatter, DMCast, JacobianStruct, SubMatrixStruct,
                                 CallbackDM)
 from devito.petsc.types.macros import petsc_func_begin_user, Null
+# from devito.petsc.initialize import PetscGetArgs
 
 
 class CBBuilder:
@@ -149,14 +149,16 @@ def _make_options_callback(self):
 
         for k, v in params.items():
             option = f'-{prefix}{k}'
-            # if option in sys.argv:
-            # TODO: Pre-build the KSPGetArgs operator and run it here
-            # to drop the global _petsc_clargs
-            # tmp = petsc_get_args_op.apply()
+
+            # TODO: drop the global variable _petsc_clargs..
+            # from devito.petsc.initialize import PetscGetArgs
+            # PetscGetArgs()
+
             import devito.petsc.initialize
             if option in devito.petsc.initialize._petsc_clargs:
                 # Ensures that the command line args take priority
                 continue
+
             option_name = String(option)
             # For options without a value e.g `ksp_view`, pass Null
             option_value = Null if v is None else String(str(v))

devito/petsc/initialize.py

@@ -1,29 +1,34 @@
 import os
 import sys
-from ctypes import POINTER, cast, c_char
+from ctypes import POINTER, cast, c_char, c_int, c_char_p, byref
 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
+from devito.petsc.types import Initialize, Finalize, GetArgs
 
 global _petsc_initialized
 _petsc_initialized = False
 
 
 global _petsc_clargs
 
+
+dummy = Symbol(name='d')
+
+
 def PetscInitialize(clargs=sys.argv):
     global _petsc_initialized
     global _petsc_clargs
 
     if not _petsc_initialized:
         if clargs is not sys.argv:
             clargs = [sys.argv[0], *clargs]
-        # TODO: drop this
+
+        # TODO: Drop this global variable
         _petsc_clargs = clargs
-        dummy = Symbol(name='d')
+
         # TODO: Potentially just use cgen + the compiler machinery in Devito
         # to generate these "dummy_ops" instead of using the Operator class.
         # This would prevent circular imports when initializing during import
@@ -46,6 +51,18 @@ def PetscInitialize(clargs=sys.argv):
             *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
+
+
+with switchconfig(language='petsc'):
+    op_get_args = Operator(
+        [PetscEq(dummy, GetArgs(dummy))],
+        name='kernel_get_args', opt='noop'
+    )
+
+
+def PetscGetArgs():
+    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

@@ -23,6 +23,7 @@ def __repr__(self):
 
 class PetscSummary(dict):
     """
+    # TODO: Actually print to screen when DEBUG of PERF is enabled
     A summary of PETSc statistics collected for all solver runs
     associated with a single operator during execution.
     """

devito/petsc/types/object.py

@@ -1,4 +1,4 @@
-from ctypes import POINTER, c_char, c_char_p
+from ctypes import POINTER, c_char, c_char_p, c_int
 
 from devito.tools import CustomDtype, dtype_to_ctype, as_tuple, CustomIntType
 from devito.types import (LocalObject, LocalCompositeObject, ModuloDimension,
@@ -305,6 +305,18 @@ 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):

devito/petsc/types/types.py

@@ -28,6 +28,10 @@ class Finalize(MetaData):
     pass
 
 
+class GetArgs(MetaData):
+    pass
+
+
 class SolveExpr(MetaData):
     """
     A symbolic expression passed through the Operator, containing the metadata

devito/petsc/utils.py

@@ -1,7 +1,6 @@
 import os
 import ctypes
 from pathlib import Path
-from petsctools import get_petscvariables
 
 from devito.tools import memoized_func, filter_ordered
 from devito.types import Symbol, SteppingDimension
@@ -48,7 +47,32 @@ def core_metadata():
     }
 
 
-petsc_variables = get_petscvariables()
+@memoized_func
+def get_petsc_variables():
+    """
+    Taken from https://www.firedrakeproject.org/_modules/firedrake/petsc.html
+    Get a dict of PETSc environment variables from the file:
+    $PETSC_DIR/$PETSC_ARCH/lib/petsc/conf/petscvariables
+    """
+    try:
+        petsc_dir = get_petsc_dir()
+    except PetscOSError:
+        petsc_variables = {}
+    else:
+        path = [petsc_dir[-1], 'lib', 'petsc', 'conf', 'petscvariables']
+        variables_path = Path(*path)
+
+        with open(variables_path) as fh:
+            # Split lines on first '=' (assignment)
+            splitlines = (line.split("=", maxsplit=1) for line in fh.readlines())
+        petsc_variables = {k.strip(): v.strip() for k, v in splitlines}
+
+    return petsc_variables
+
+
+petsc_variables = get_petsc_variables()
+# TODO: use petsctools get_petscvariables() instead?
+# petsc_variables = get_petscvariables()
 
 
 def get_petsc_type_mappings():

requirements.txt

@@ -9,4 +9,4 @@ multidict<6.3
 anytree>=2.4.3,<=2.13.0
 cloudpickle<3.1.2
 packaging<25.1
-petsctools
+petsctools==2025.1.dev0

tests/test_petsc.py

@@ -3,7 +3,6 @@
 import numpy as np
 import os
 import re
-import sys
 from collections import OrderedDict
 
 from conftest import skipif
@@ -24,26 +23,26 @@
 from devito.petsc.logging import PetscSummary
 from devito.petsc.solver_parameters import linear_solve_defaults
 
+
 @pytest.fixture(scope='session')
 def command_line():
-    # one prefix per test
+    # One random prefix to use per test that "tests" the command line args
     prefix = ('d17weqroegn', 'riabfodkj')
 
     petsc_option = (
         ('ksp_rtol',),
-        ('ksp_rtol','ksp_atol')
+        ('ksp_rtol', 'ksp_atol')
     )
     value = (
-        ('1e-8',),
-        ('1e-11','1e-15'),
+        (1e-8,),
+        (1e-11, 1e-15),
     )
     argv = []
-
     expected = {}
     for p, opt, val in zip(prefix, petsc_option, value, strict=True):
         for o, v in zip(opt, val, strict=True):
-            argv.extend([f'-{p}_{o}', v])
-        expected[p] = zip(opt,val)
+            argv.extend([f'-{p}_{o}', str(v)])
+        expected[p] = zip(opt, val)
     return argv, expected
 
 
@@ -1448,7 +1447,7 @@ def test_logging(self, log_level):
         # the tolerances should match the default linear values.
         tols = entry0.KSPGetTolerances
         assert tols['rtol'] == linear_solve_defaults['ksp_rtol']
-        assert tols['abstol'] == linear_solve_defaults['ksp_atol']
+        assert tols['atol'] == linear_solve_defaults['ksp_atol']
         assert tols['dtol'] == linear_solve_defaults['ksp_divtol']
         assert tols['maxits'] == linear_solve_defaults['ksp_max_it']
 
@@ -1655,7 +1654,7 @@ def test_tolerances(self, log_level):
         # Test that the tolerances have been set correctly and therefore
         # appear as expected in the `PetscSummary`.
         assert tolerances['rtol'] == params['ksp_rtol']
-        assert tolerances['abstol'] == params['ksp_atol']
+        assert tolerances['atol'] == params['ksp_atol']
         assert tolerances['dtol'] == params['ksp_divtol']
         assert tolerances['maxits'] == params['ksp_max_it']
 
@@ -1754,7 +1753,7 @@ def test_command_line_priority_1(self, command_line):
         petsc_summary = summary.petsc
         entry = petsc_summary.get_entry('section0', prefix)
         for opt, val in expected[prefix]:
-            assert str(entry.KSPGetTolerances[opt.removeprefix('ksp_')]) == val
+            assert entry.KSPGetTolerances[opt.removeprefix('ksp_')] == val
 
     @skipif('petsc')
     def test_command_line_priority_2(self, command_line):
@@ -1774,7 +1773,7 @@ def test_command_line_priority_2(self, command_line):
         petsc_summary = summary.petsc
         entry = petsc_summary.get_entry('section0', prefix)
         for opt, val in expected[prefix]:
-            assert str(entry.KSPGetTolerances[opt.removeprefix('ksp_')]) == val
+            assert entry.KSPGetTolerances[opt.removeprefix('ksp_')] == val
 
 
 class TestHashing: