compiler: Add PETSc language and printer (#67)

ZoeLeibowitz · web-flow · commit f61d4dd27e04 · 2025-03-26T22:33:17.000Z
diff --git a/conftest.py b/conftest.py
@@ -14,7 +14,7 @@
 from devito.ir.iet import (FindNodes, FindSymbols, Iteration, ParallelBlock,
                            retrieve_iteration_tree)
 from devito.tools import as_tuple
-from devito.petsc.utils import get_petsc_dir, get_petsc_arch
+from devito.petsc.utils import PetscOSError, get_petsc_dir
 
 try:
     from mpi4py import MPI  # noqa
@@ -95,18 +95,11 @@ def skipif(items, whole_module=False):
             skipit = "pyrevolve not installed"
             break
         if i == 'petsc':
-            petsc_dir = get_petsc_dir()
-            petsc_arch = get_petsc_arch()
-            if petsc_dir is None or petsc_arch is None:
-                skipit = "PETSC_DIR or PETSC_ARCH are not set"
+            try:
+                _ = get_petsc_dir()
+            except PetscOSError:
+                skipit = "PETSc is not installed"
                 break
-            else:
-                petsc_installed = os.path.join(
-                    petsc_dir, petsc_arch, 'include', 'petscconf.h'
-                )
-                if not os.path.isfile(petsc_installed):
-                    skipit = "PETSc is not installed"
-                    break
 
     if skipit is False:
         return pytest.mark.skipif(False, reason='')
diff --git a/devito/core/__init__.py b/devito/core/__init__.py
@@ -6,7 +6,8 @@
     Cpu64CustomOperator, Cpu64CustomCXXOperator,
     Cpu64CXXNoopCOperator, Cpu64CXXNoopOmpOperator,
     Cpu64AdvCXXOperator, Cpu64AdvCXXOmpOperator,
-    Cpu64FsgCXXOperator, Cpu64FsgCXXOmpOperator
+    Cpu64FsgCXXOperator, Cpu64FsgCXXOmpOperator,
+    Cpu64NoopPetscOperator, Cpu64AdvPetscOperator
 )
 from devito.core.intel import (
     Intel64AdvCOperator, Intel64AdvOmpOperator,
@@ -40,11 +41,13 @@
 operator_registry.add(Cpu64NoopOmpOperator, Cpu64, 'noop', 'openmp')
 operator_registry.add(Cpu64CXXNoopCOperator, Cpu64, 'noop', 'CXX')
 operator_registry.add(Cpu64CXXNoopOmpOperator, Cpu64, 'noop', 'CXXopenmp')
+operator_registry.add(Cpu64NoopPetscOperator, Cpu64, 'noop', 'petsc')
 
 operator_registry.add(Cpu64AdvCOperator, Cpu64, 'advanced', 'C')
 operator_registry.add(Cpu64AdvOmpOperator, Cpu64, 'advanced', 'openmp')
 operator_registry.add(Cpu64AdvCXXOperator, Cpu64, 'advanced', 'CXX')
 operator_registry.add(Cpu64AdvCXXOmpOperator, Cpu64, 'advanced', 'CXXopenmp')
+operator_registry.add(Cpu64AdvPetscOperator, Cpu64, 'advanced', 'petsc')
 
 operator_registry.add(Cpu64FsgCOperator, Cpu64, 'advanced-fsg', 'C')
 operator_registry.add(Cpu64FsgOmpOperator, Cpu64, 'advanced-fsg', 'openmp')
diff --git a/devito/core/cpu.py b/devito/core/cpu.py
@@ -11,13 +11,14 @@
 from devito.passes.iet import (CTarget, CXXTarget, COmpTarget, CXXOmpTarget,
                                avoid_denormals, linearize,
                                mpiize, hoist_prodders, relax_incr_dimensions,
-                               check_stability)
+                               check_stability, PetscTarget)
 from devito.tools import timed_pass
 
 __all__ = ['Cpu64NoopCOperator', 'Cpu64NoopOmpOperator', 'Cpu64AdvCOperator',
            'Cpu64AdvOmpOperator', 'Cpu64FsgCOperator', 'Cpu64FsgOmpOperator',
            'Cpu64CustomOperator', 'Cpu64CustomCXXOperator', 'Cpu64AdvCXXOperator',
-           'Cpu64AdvCXXOmpOperator', 'Cpu64FsgCXXOperator', 'Cpu64FsgCXXOmpOperator']
+           'Cpu64AdvCXXOmpOperator', 'Cpu64FsgCXXOperator', 'Cpu64FsgCXXOmpOperator',
+           'Cpu64NoopPetscOperator']
 
 
 class Cpu64OperatorMixin:
@@ -346,6 +347,15 @@ class Cpu64CXXNoopOmpOperator(Cpu64NoopOperator):
     LINEARIZE = True
 
 
+class Cpu64NoopPetscOperator(Cpu64NoopOperator):
+    _Target = PetscTarget
+
+    @classmethod
+    def _rcompile_wrapper(cls, **kwargs0):
+        kwargs0['language'] = 'petsc'
+        return super()._rcompile_wrapper(**kwargs0)
+
+
 class Cpu64AdvCOperator(Cpu64AdvOperator):
     _Target = CTarget
 
@@ -355,6 +365,15 @@ class Cpu64AdvCXXOperator(Cpu64AdvOperator):
     LINEARIZE = True
 
 
+class Cpu64AdvPetscOperator(Cpu64AdvOperator):
+    _Target = PetscTarget
+
+    @classmethod
+    def _rcompile_wrapper(cls, **kwargs0):
+        kwargs0['language'] = 'petsc'
+        return super()._rcompile_wrapper(**kwargs0)
+
+
 class Cpu64AdvOmpOperator(Cpu64AdvOperator):
     _Target = COmpTarget
 
diff --git a/devito/operator/registry.py b/devito/operator/registry.py
@@ -27,7 +27,7 @@ class OperatorRegistry(OrderedDict, metaclass=Singleton):
 
     _modes = ('noop', 'advanced', 'advanced-fsg')
     _languages = ('C', 'CXX', 'openmp', 'Copenmp', 'CXXopenmp',
-                  'openacc', 'cuda', 'hip', 'sycl')
+                  'openacc', 'cuda', 'hip', 'sycl', 'petsc')
     _accepted = _modes + tuple(product(_modes, _languages))
 
     def add(self, operator, platform, mode, language='C'):
diff --git a/devito/passes/iet/languages/C.py b/devito/passes/iet/languages/C.py
@@ -6,6 +6,7 @@
 from devito.passes.iet.orchestration import Orchestrator
 from devito.passes.iet.langbase import LangBB
 from devito.symbolics.extended_dtypes import c_complex, c_double_complex
+from devito.petsc.utils import petsc_type_mappings
 
 __all__ = ['CBB', 'CDataManager', 'COrchestrator']
 
@@ -55,3 +56,9 @@ class CPrinter(BasePrinter, C99CodePrinter):
 
     def _print_ImaginaryUnit(self, expr):
         return '_Complex_I'
+
+
+class PetscCPrinter(CPrinter):
+    _restrict_keyword = ''
+
+    type_mappings = {**CPrinter.type_mappings, **petsc_type_mappings}
diff --git a/devito/passes/iet/languages/targets.py b/devito/passes/iet/languages/targets.py
@@ -1,4 +1,5 @@
-from devito.passes.iet.languages.C import CDataManager, COrchestrator, CPrinter
+from devito.passes.iet.languages.C import (CDataManager, COrchestrator, CPrinter,
+                                           PetscCPrinter)
 from devito.passes.iet.languages.CXX import CXXPrinter
 from devito.passes.iet.languages.openmp import (SimdOmpizer, Ompizer, DeviceOmpizer,
                                                 OmpDataManager, DeviceOmpDataManager,
@@ -8,7 +9,7 @@
 from devito.passes.iet.instrument import instrument
 
 __all__ = ['CTarget', 'OmpTarget', 'COmpTarget', 'DeviceOmpTarget', 'DeviceAccTarget',
-           'CXXTarget', 'CXXOmpTarget', 'DeviceCXXOmpTarget']
+           'CXXTarget', 'CXXOmpTarget', 'DeviceCXXOmpTarget', 'PetscTarget']
 
 
 class Target:
@@ -47,6 +48,10 @@ class COmpTarget(Target):
     Printer = CPrinter
 
 
+class PetscTarget(CTarget):
+    Printer = PetscCPrinter
+
+
 OmpTarget = COmpTarget
 
 
diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py
@@ -17,7 +17,7 @@
                                 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
+from devito.petsc.utils import core_metadata, petsc_languages
 from devito.petsc.iet.routines import (CBBuilder, CCBBuilder, BaseObjectBuilder,
                                        CoupledObjectBuilder, BaseSetup, CoupledSetup,
                                        Solver, CoupledSolver, TimeDependent,
@@ -34,6 +34,12 @@ def lower_petsc(iet, **kwargs):
     if not injectsolve_mapper:
         return iet, {}
 
+    if kwargs['language'] not in petsc_languages:
+        raise ValueError(
+            f"Expected 'language' to be one of "
+            f"{petsc_languages}, but got '{kwargs['language']}'"
+        )
+
     metadata = core_metadata()
     data = FindNodes(PetscMetaData).visit(iet)
 
diff --git a/devito/petsc/initialize.py b/devito/petsc/initialize.py
@@ -3,7 +3,7 @@
 from ctypes import POINTER, cast, c_char
 import atexit
 
-from devito import Operator
+from devito import Operator, switchconfig
 from devito.types import Symbol
 from devito.types.equation import PetscEq
 from devito.petsc.types import Initialize, Finalize
@@ -20,14 +20,15 @@ def PetscInitialize():
         # to generate these "dummy_ops" instead of using the Operator class.
         # This would prevent circular imports when initializing during import
         # from the PETSc module.
-        op_init = Operator(
-            [PetscEq(dummy, Initialize(dummy))],
-            name='kernel_init', opt='noop'
-        )
-        op_finalize = Operator(
-            [PetscEq(dummy, Finalize(dummy))],
-            name='kernel_finalize', opt='noop'
-        )
+        with switchconfig(language='petsc'):
+            op_init = Operator(
+                [PetscEq(dummy, Initialize(dummy))],
+                name='kernel_init', opt='noop'
+            )
+            op_finalize = Operator(
+                [PetscEq(dummy, Finalize(dummy))],
+                name='kernel_finalize', opt='noop'
+            )
 
         # `argv_bytes` must be a list so the memory address persists
         # `os.fsencode` should be preferred over `string().encode('utf-8')`
diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py
@@ -106,10 +106,6 @@ def shape_allocated(self):
 
     @cached_property
     def _C_ctype(self):
-        # TODO: Switch to using PetscScalar instead of float/double
-        # TODO: Use cat $PETSC_DIR/$PETSC_ARCH/lib/petsc/conf/petscvariables
-        # | grep -E "PETSC_(SCALAR|PRECISION)" to determine the precision of
-        # the user's PETSc configuration.
         return POINTER(dtype_to_ctype(self.dtype))
 
     @property
diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py
@@ -1,5 +1,5 @@
 from ctypes import POINTER, c_char
-from devito.tools import CustomDtype, dtype_to_cstr, as_tuple, CustomIntType
+from devito.tools import CustomDtype, dtype_to_ctype, as_tuple, CustomIntType
 from devito.types import (LocalObject, LocalCompositeObject, ModuloDimension,
                           TimeDimension, ArrayObject, CustomDimension)
 from devito.symbolics import Byref, cast
@@ -188,7 +188,7 @@ class VecScatter(LocalObject):
 class StartPtr(LocalObject):
     def __init__(self, name, dtype):
         super().__init__(name=name)
-        self.dtype = CustomDtype(dtype_to_cstr(dtype), modifier=' *')
+        self.dtype = POINTER(dtype_to_ctype(dtype))
 
 
 class SingleIS(LocalObject):
diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py
@@ -1,6 +1,8 @@
 import sympy
 
 from devito.tools import Reconstructable, sympy_mutex
+from devito.tools.dtypes_lowering import dtype_mapper
+from devito.petsc.utils import petsc_variables
 
 
 class MetaData(sympy.Function, Reconstructable):
@@ -132,6 +134,14 @@ class FieldData:
     def __init__(self, target=None, matvecs=None, formfuncs=None, formrhs=None,
                  arrays=None, **kwargs):
         self._target = kwargs.get('target', target)
+
+        petsc_precision = dtype_mapper[petsc_variables['PETSC_PRECISION']]
+        if self._target.dtype != petsc_precision:
+            raise TypeError(
+                f"Your target dtype must match the precision of your "
+                f"PETSc configuration. "
+                f"Expected {petsc_precision}, but got {self._target.dtype}."
+            )
         self._matvecs = matvecs
         self._formfuncs = formfuncs
         self._formrhs = formrhs
diff --git a/devito/petsc/utils.py b/devito/petsc/utils.py
@@ -1,8 +1,14 @@
 import os
+import ctypes
 
+from pathlib import Path
 from devito.tools import memoized_func
 
 
+class PetscOSError(OSError):
+    pass
+
+
 solver_mapper = {
     'gmres': 'KSPGMRES',
     'jacobi': 'PCJACOBI',
@@ -12,42 +18,70 @@
 
 @memoized_func
 def get_petsc_dir():
-    # *** First try: via commonly used environment variables
-    for i in ['PETSC_DIR']:
-        petsc_dir = os.environ.get(i)
-        if petsc_dir:
-            return petsc_dir
-    # TODO: Raise error if PETSC_DIR is not set
-    return None
+    petsc_dir = os.environ.get('PETSC_DIR')
+    if petsc_dir is None:
+        raise PetscOSError("PETSC_DIR environment variable not set")
+    else:
+        petsc_dir = (Path(petsc_dir),)
 
+    petsc_arch = os.environ.get('PETSC_ARCH')
+    if petsc_arch is not None:
+        petsc_dir += (petsc_dir[0] / petsc_arch,)
 
-@memoized_func
-def get_petsc_arch():
-    # *** First try: via commonly used environment variables
-    for i in ['PETSC_ARCH']:
-        petsc_arch = os.environ.get(i)
-        if petsc_arch:
-            return petsc_arch
-    # TODO: Raise error if PETSC_ARCH is not set
-    return None
+    petsc_installed = petsc_dir[-1] / 'include' / 'petscconf.h'
+    if not petsc_installed.is_file():
+        raise PetscOSError("PETSc is not installed")
+
+    return petsc_dir
 
 
+@memoized_func
 def core_metadata():
     petsc_dir = get_petsc_dir()
-    petsc_arch = get_petsc_arch()
-
-    # Include directories
-    global_include = os.path.join(petsc_dir, 'include')
-    config_specific_include = os.path.join(petsc_dir, f'{petsc_arch}', 'include')
-    include_dirs = (global_include, config_specific_include)
 
-    # Lib directories
-    lib_dir = os.path.join(petsc_dir, f'{petsc_arch}', 'lib')
+    petsc_include = tuple([arch / 'include' for arch in petsc_dir])
+    petsc_lib = tuple([arch / 'lib' for arch in petsc_dir])
 
     return {
         'includes': ('petscsnes.h', 'petscdmda.h'),
-        'include_dirs': include_dirs,
+        'include_dirs': petsc_include,
         'libs': ('petsc'),
-        'lib_dirs': lib_dir,
-        'ldflags': ('-Wl,-rpath,%s' % lib_dir)
+        'lib_dirs': petsc_lib,
+        'ldflags': tuple([f"-Wl,-rpath,{lib}" for lib in petsc_lib])
     }
+
+
+@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: Check to see whether Petsc is compiled with
+# 32-bit or 64-bit integers
+# TODO: Check whether PetscScalar is a float or double
+# and only map the right one
+petsc_type_mappings = {ctypes.c_int: 'PetscInt',
+                       ctypes.c_float: 'PetscScalar',
+                       ctypes.c_double: 'PetscScalar'}
+
+petsc_languages = ['petsc']
diff --git a/examples/petsc/petsc_test.py b/examples/petsc/petsc_test.py
diff --git a/tests/test_petsc.py b/tests/test_petsc.py
