compiler: Initialize and Finalize PETSc once and command line args

Author: ZoeLeibowitz

devito/ir/equations/equation.py

@@ -11,10 +11,10 @@
 from devito.tools import Pickable, Tag, frozendict
 from devito.types import (Eq, Inc, ReduceMax, ReduceMin,
                           relational_min)
-from devito.types.equation import InjectSolveEq
+from devito.types.equation import PetscEq
 
 __all__ = ['LoweredEq', 'ClusterizedEq', 'DummyEq', 'OpInc', 'OpMin', 'OpMax',
-           'identity_mapper', 'OpInjectSolve']
+           'identity_mapper', 'OpPetsc']
 
 
 class IREq(sympy.Eq, Pickable):
@@ -105,7 +105,7 @@ def detect(cls, expr):
             Inc: OpInc,
             ReduceMax: OpMax,
             ReduceMin: OpMin,
-            InjectSolveEq: OpInjectSolve
+            PetscEq: OpPetsc
         }
         try:
             return reduction_mapper[type(expr)]
@@ -122,7 +122,7 @@ def detect(cls, expr):
 OpInc = Operation('+')
 OpMax = Operation('max')
 OpMin = Operation('min')
-OpInjectSolve = Operation('solve')
+OpPetsc = Operation('solve')
 
 
 identity_mapper = {

devito/ir/iet/algorithms.py

@@ -3,7 +3,7 @@
 from devito.ir.iet import (Expression, Increment, Iteration, List, Conditional, SyncSpot,
                            Section, HaloSpot, ExpressionBundle)
 from devito.tools import timed_pass
-from devito.petsc.types import LinearSolveExpr
+from devito.petsc.types import MetaData
 from devito.petsc.iet.utils import petsc_iet_mapper
 
 __all__ = ['iet_build']
@@ -26,7 +26,7 @@ def iet_build(stree):
             for e in i.exprs:
                 if e.is_Increment:
                     exprs.append(Increment(e))
-                elif isinstance(e.rhs, LinearSolveExpr):
+                elif isinstance(e.rhs, MetaData):
                     exprs.append(petsc_iet_mapper[e.operation](e, operation=e.operation))
                 else:
                     exprs.append(Expression(e, operation=e.operation))

devito/petsc/iet/nodes.py

@@ -1,20 +1,12 @@
 from devito.ir.iet import Expression, Callback, FixedArgsCallable, Call
-from devito.ir.equations import OpInjectSolve
+from devito.ir.equations import OpPetsc
 
 
-class LinearSolverExpression(Expression):
+class PetscMetaData(Expression):
     """
-    Base class for general expressions required by a
-    matrix-free linear solve of the form Ax=b.
+    Base class for general expressions required to run a PETSc solver.
     """
-    pass
-
-
-class InjectSolveDummy(LinearSolverExpression):
-    """
-    Placeholder expression to run the iterative solver.
-    """
-    def __init__(self, expr, pragmas=None, operation=OpInjectSolve):
+    def __init__(self, expr, pragmas=None, operation=OpPetsc):
         super().__init__(expr, pragmas=pragmas, operation=operation)
 
 

devito/petsc/iet/passes.py

@@ -1,10 +1,15 @@
 import cgen as c
+import numpy as np
 
 from devito.passes.iet.engine import iet_pass
-from devito.ir.iet import Transformer, MapNodes, Iteration, BlankLine
+from devito.ir.iet import (Transformer, MapNodes, Iteration, BlankLine,
+                           FindNodes, Call, CallableBody)
 from devito.symbolics import Byref, Macro
-from devito.petsc.types import (PetscMPIInt, PetscErrorCode)
-from devito.petsc.iet.nodes import InjectSolveDummy
+from devito.types.basic import DataSymbol
+from devito.petsc.types import (PetscMPIInt, PetscErrorCode, 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.iet.routines import (CallbackBuilder, BaseObjectBuilder, BaseSetup,
                                        Solver, TimeDependent, NonTimeDependent)
@@ -14,14 +19,23 @@
 @iet_pass
 def lower_petsc(iet, **kwargs):
     # Check if PETScSolve was used
-    injectsolve_mapper = MapNodes(Iteration, InjectSolveDummy,
+    injectsolve_mapper = MapNodes(Iteration, PetscMetaData,
                                   'groupby').visit(iet)
 
     if not injectsolve_mapper:
         return iet, {}
 
+    metadata = core_metadata()
+
+    data = FindNodes(PetscMetaData).visit(iet)
+
+    if any(filter(lambda i: isinstance(i.expr.rhs, Initialize), data)):
+        return initialize(iet), metadata
+
+    if any(filter(lambda i: isinstance(i.expr.rhs, Finalize), data)):
+        return finalize(iet), metadata
+
     targets = [i.expr.rhs.target for (i,) in injectsolve_mapper.values()]
-    init = init_petsc(**kwargs)
 
     # Assumption is that all targets have the same grid so can use any target here
     objs = build_core_objects(targets[-1], **kwargs)
@@ -47,26 +61,37 @@ def lower_petsc(iet, **kwargs):
     iet = Transformer(subs).visit(iet)
 
     body = core + tuple(setup) + (BlankLine,) + iet.body.body
-    body = iet.body._rebuild(
-        init=init, body=body,
-        frees=(petsc_call('PetscFinalize', []),)
-    )
+    body = iet.body._rebuild(body=body)
     iet = iet._rebuild(body=body)
-    metadata = core_metadata()
     metadata.update({'efuncs': tuple(efuncs.values())})
 
     return iet, metadata
 
 
-def init_petsc(**kwargs):
-    # Initialize PETSc -> for now, assuming all solver options have to be
-    # specified via the parameters dict in PETScSolve
-    # TODO: Are users going to be able to use PETSc command line arguments?
-    # In firedrake, they have an options_prefix for each solver, enabling the use
-    # of command line options
-    initialize = petsc_call('PetscInitialize', [Null, Null, Null, Null])
+def initialize(iet):
+    # should be int because the correct type for argc is a C int
+    # and not a int32
+    argc = DataSymbol(name='argc', dtype=np.int32)
+    argv = ArgvSymbol(name='argv')
+    Help = Macro('help')
 
-    return petsc_func_begin_user, initialize
+    help_string = c.Line(r'static char help[] = "This is help text.\n";')
+
+    init_body = petsc_call('PetscInitialize', [Byref(argc), Byref(argv), Null, Help])
+    init_body = CallableBody(
+        body=(petsc_func_begin_user, help_string, init_body),
+        retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
+    )
+    return iet._rebuild(body=init_body)
+
+
+def finalize(iet):
+    finalize_body = petsc_call('PetscFinalize', [])
+    finalize_body = CallableBody(
+        body=(petsc_func_begin_user, finalize_body),
+        retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
+    )
+    return iet._rebuild(body=finalize_body)
 
 
 def make_core_petsc_calls(objs, **kwargs):
@@ -76,10 +101,7 @@ def make_core_petsc_calls(objs, **kwargs):
 
 
 def build_core_objects(target, **kwargs):
-    if kwargs['options']['mpi']:
-        communicator = target.grid.distributor._obj_comm
-    else:
-        communicator = 'PETSC_COMM_SELF'
+    communicator = 'PETSC_COMM_WORLD'
 
     return {
         'size': PetscMPIInt(name='size'),
@@ -128,9 +150,6 @@ def __init__(self, injectsolve, objs, iters, **kwargs):
         )
 
 
+# Move these to types folder
 Null = Macro('NULL')
 void = 'void'
-
-
-# TODO: Don't use c.Line here?
-petsc_func_begin_user = c.Line('PetscFunctionBeginUser;')

devito/petsc/iet/routines.py

@@ -14,7 +14,7 @@
 
 from devito.petsc.types import PETScArray
 from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback,
-                                    MatVecCallback, InjectSolveDummy)
+                                    MatVecCallback, PetscMetaData)
 from devito.petsc.iet.utils import petsc_call, petsc_struct
 from devito.petsc.utils import solver_mapper
 from devito.petsc.types import (DM, CallbackDM, Mat, LocalVec, GlobalVec, KSP, PC,
@@ -768,7 +768,7 @@ def _spatial_loop_nest(self, iters, injectsolve):
         spatial_body = []
         for tree in retrieve_iteration_tree(iters[0]):
             root = filter_iterations(tree, key=lambda i: i.dim.is_Space)[0]
-            if injectsolve in FindNodes(InjectSolveDummy).visit(root):
+            if injectsolve in FindNodes(PetscMetaData).visit(root):
                 spatial_body.append(root)
         return spatial_body
 

devito/petsc/iet/utils.py

@@ -1,5 +1,5 @@
-from devito.petsc.iet.nodes import InjectSolveDummy, PETScCall
-from devito.ir.equations import OpInjectSolve
+from devito.petsc.iet.nodes import PetscMetaData, PETScCall
+from devito.ir.equations import OpPetsc
 
 
 def petsc_call(specific_call, call_args):
@@ -19,4 +19,4 @@ def petsc_struct(name, fields, pname, liveness='lazy'):
 
 # Mapping special Eq operations to their corresponding IET Expression subclass types.
 # These operations correspond to subclasses of Eq utilised within PETScSolve.
-petsc_iet_mapper = {OpInjectSolve: InjectSolveDummy}
+petsc_iet_mapper = {OpPetsc: PetscMetaData}

devito/petsc/initialize.py

@@ -0,0 +1,42 @@
+import os
+import sys
+from ctypes import POINTER, cast, c_char
+import atexit
+
+from devito import Operator
+from devito.types import Symbol
+from devito.types.equation import PetscEq
+from devito.petsc.types import Initialize, Finalize
+
+global _petsc_initialized
+_petsc_initialized = False
+
+
+def PetscInitialize():
+    global _petsc_initialized
+    if not _petsc_initialized:
+        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
+        # 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'
+        )
+
+        # `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, sys.argv))
+        argv_pointer = (POINTER(c_char)*len(sys.argv))(
+            *map(lambda s: cast(s, POINTER(c_char)), argv_bytes)
+        )
+        op_init.apply(argc=len(sys.argv), argv=argv_pointer)
+
+        atexit.register(op_finalize.apply)
+        _petsc_initialized = True

devito/petsc/solve.py

@@ -5,7 +5,7 @@
 from devito.finite_differences.differentiable import Mul
 from devito.finite_differences.derivative import Derivative
 from devito.types import Eq, Symbol, SteppingDimension, TimeFunction
-from devito.types.equation import InjectSolveEq
+from devito.types.equation import PetscEq
 from devito.operations.solve import eval_time_derivatives
 from devito.symbolics import retrieve_functions
 from devito.tools import as_tuple
@@ -65,7 +65,7 @@ def PETScSolve(eqns, target, solver_parameters=None, **kwargs):
     )
     # Placeholder equation for inserting calls to the solver and generating
     # correct time loop etc
-    inject_solve = InjectSolveEq(target, LinearSolveExpr(
+    inject_solve = PetscEq(target, LinearSolveExpr(
         expr=tuple(funcs),
         target=target,
         solver_parameters=solver_parameters,

devito/petsc/types/macros.py

@@ -0,0 +1,5 @@
+import cgen as c
+
+
+# TODO: Don't use c.Line here?
+petsc_func_begin_user = c.Line('PetscFunctionBeginUser;')

devito/petsc/types/object.py

@@ -1,7 +1,8 @@
-from ctypes import POINTER
+from ctypes import POINTER, c_char
 
 from devito.tools import CustomDtype, dtype_to_cstr
 from devito.types import LocalObject, CCompositeObject, ModuloDimension, TimeDimension
+from devito.types.basic import DataSymbol
 from devito.symbolics import Byref
 
 from devito.petsc.iet.utils import petsc_call
@@ -209,3 +210,9 @@ class StartPtr(LocalObject):
     def __init__(self, name, dtype):
         super().__init__(name=name)
         self.dtype = CustomDtype(dtype_to_cstr(dtype), modifier=' *')
+
+
+class ArgvSymbol(DataSymbol):
+    @property
+    def _C_ctype(self):
+        return POINTER(POINTER(c_char))