Merge pull request #217 from mfem/enzyme_hpp

added enzyme.i and necessary changes to build
mfem · Mar 26, 2024 · 42722e5 · 42722e5
2 parents 324cd66 + 89effee
commit 42722e5
Show file tree

Hide file tree

Showing 16 changed files with 132 additions and 88 deletions.
diff --git a/examples/ex26.py b/examples/ex26.py
@@ -18,34 +18,30 @@ def run(order_refinements=2,
 
     class DiffusionMultigrid(mfem.PyGeometricMultigrid):
         def __init__(self, fespaces, ess_bdr):
-            mfem.PyGeometricMultigrid.__init__(self, fespaces)
+            mfem.PyGeometricMultigrid.__init__(self, fespaces, ess_bdr)
 
             self.smoothers = []
-            self.ConstructCoarseOperatorAndSolver(fespaces.GetFESpaceAtLevel(0),
-                                                  ess_bdr)
+            self.ConstructCoarseOperatorAndSolver(
+                fespaces.GetFESpaceAtLevel(0),)
 
             for level in range(1, fespaces.GetNumLevels()):
                 self.ConstructOperatorAndSmoother(fespaces.GetFESpaceAtLevel(level),
-                                                  ess_bdr)
+                                                  level)
 
-        def ConstructBilinearForm(self, fespace, ess_bdr):
+        def ConstructBilinearForm(self, fespace):
             form = mfem.BilinearForm(fespace)
             form.SetAssemblyLevel(mfem.AssemblyLevel_PARTIAL)
             form.AddDomainIntegrator(mfem.DiffusionIntegrator(one))
             form.Assemble()
             self.AppendBilinearForm(form)
 
-            ess = mfem.intArray()
-            self.AppendEssentialTDofs(ess)
-            fespace.GetEssentialTrueDofs(ess_bdr, ess)
+        def ConstructCoarseOperatorAndSolver(self, coarse_fespace):
 
-        def ConstructCoarseOperatorAndSolver(self, coarse_fespace, ess_bdr):
-
-            self.ConstructBilinearForm(coarse_fespace, ess_bdr)
+            self.ConstructBilinearForm(coarse_fespace)
 
             opr = mfem.OperatorPtr()
             opr.SetType(mfem.Operator.ANY_TYPE)
-            self.bfs[-1].FormSystemMatrix(self.essentialTrueDofs[-1], opr)
+            self.bfs[0].FormSystemMatrix(self.essentialTrueDofs[0], opr)
             opr.SetOperatorOwner(False)
 
             pcg = mfem.CGSolver()
@@ -58,20 +54,20 @@ def ConstructCoarseOperatorAndSolver(self, coarse_fespace, ess_bdr):
             self.smoothers.append((pcg, opr))
             self.AddLevel(opr.Ptr(), pcg, False, False)
 
-        def ConstructOperatorAndSmoother(self, fespace, ess_bdr):
-            self.ConstructBilinearForm(fespace, ess_bdr)
+        def ConstructOperatorAndSmoother(self, fespace, level):
+            self.ConstructBilinearForm(fespace)
 
             opr = mfem.OperatorPtr()
             opr.SetType(mfem.Operator.ANY_TYPE)
-            self.bfs[-1].FormSystemMatrix(self.essentialTrueDofs[-1], opr)
+
+            self.bfs[level].FormSystemMatrix(
+                self.essentialTrueDofs[level], opr)
             opr.SetOperatorOwner(False)
 
             diag = mfem.Vector(fespace.GetTrueVSize())
-            self.bfs[-1].AssembleDiagonal(diag)
-
+            self.bfs[level].AssembleDiagonal(diag)
             smoother = mfem.OperatorChebyshevSmoother(opr.Ptr(), diag,
-                                                      self.essentialTrueDofs[-1], 2)
-
+                                                      self.essentialTrueDofs[level], 2)
             self.smoothers.append((smoother, opr))
             self.AddLevel(opr.Ptr(), smoother, False, False)
 

diff --git a/examples/ex26p.py b/examples/ex26p.py
@@ -23,36 +23,31 @@ def run(order_refinements=2,
 
     class DiffusionMultigrid(mfem.PyGeometricMultigrid):
         def __init__(self, fespaces, ess_bdr):
-            mfem.PyGeometricMultigrid.__init__(self, fespaces)
+            mfem.PyGeometricMultigrid.__init__(self, fespaces, ess_bdr)
 
             self.smoothers = []
-            self.ConstructCoarseOperatorAndSolver(fespaces.GetFESpaceAtLevel(0),
-                                                  ess_bdr)
+            self.ConstructCoarseOperatorAndSolver(
+                fespaces.GetFESpaceAtLevel(0),)
 
             for level in range(1, fespaces.GetNumLevels()):
                 self.ConstructOperatorAndSmoother(fespaces.GetFESpaceAtLevel(level),
-                                                  ess_bdr)
+                                                  level)
 
-        def ConstructBilinearForm(self, fespace, ess_bdr, partial_assembly):
+        def ConstructBilinearForm(self, fespace, partial_assembly):
             form = mfem.ParBilinearForm(fespace)
             if partial_assembly:
                 form.SetAssemblyLevel(mfem.AssemblyLevel_PARTIAL)
             form.AddDomainIntegrator(mfem.DiffusionIntegrator(one))
             form.Assemble()
             self.AppendBilinearForm(form)
 
-            ess = mfem.intArray()
-            self.AppendEssentialTDofs(ess)
-            fespace.GetEssentialTrueDofs(ess_bdr, ess)
+        def ConstructCoarseOperatorAndSolver(self, coarse_fespace,):
 
-        def ConstructCoarseOperatorAndSolver(self, coarse_fespace,
-                                             ess_bdr):
-
-            self.ConstructBilinearForm(coarse_fespace, ess_bdr, False)
+            self.ConstructBilinearForm(coarse_fespace, False)
 
             hypreCoarseMat = mfem.HypreParMatrix()
-            self.bfs[-1].FormSystemMatrix(self.essentialTrueDofs[-1],
-                                          hypreCoarseMat)
+            self.bfs[0].FormSystemMatrix(self.essentialTrueDofs[0],
+                                         hypreCoarseMat)
 
             amg = mfem.HypreBoomerAMG(hypreCoarseMat)
             amg.SetPrintLevel(-1)
@@ -69,20 +64,21 @@ def ConstructCoarseOperatorAndSolver(self, coarse_fespace,
             self.AddLevel(hypreCoarseMat, pcg, False, False)
             self.amg = amg
 
-        def ConstructOperatorAndSmoother(self, fespace, ess_bdr):
-            self.ConstructBilinearForm(fespace, ess_bdr, True)
+        def ConstructOperatorAndSmoother(self, fespace, level):
+            self.ConstructBilinearForm(fespace, True)
 
             opr = mfem.OperatorPtr()
             opr.SetType(mfem.Operator.ANY_TYPE)
 
-            self.bfs[-1].FormSystemMatrix(self.essentialTrueDofs[-1], opr)
+            self.bfs[level].FormSystemMatrix(
+                self.essentialTrueDofs[level], opr)
             opr.SetOperatorOwner(False)
 
             diag = mfem.Vector(fespace.GetTrueVSize())
-            self.bfs[-1].AssembleDiagonal(diag)
+            self.bfs[level].AssembleDiagonal(diag)
 
             smoother = mfem.OperatorChebyshevSmoother(opr.Ptr(), diag,
-                                                      self.essentialTrueDofs[-1],
+                                                      self.essentialTrueDofs[level],
                                                       2,
                                                       fespace.GetParMesh().GetComm())
 

diff --git a/mfem/_par/array.i b/mfem/_par/array.i
@@ -174,3 +174,6 @@ INSTANTIATE_ARRAY_NUMPYARRAY(uint, unsigned int, NPY_UINT)       // 32bit
 %template(intArray2D) mfem::Array2D<int>;
 %template(doubleArray2D) mfem::Array2D<double>;
 
+/* Array< Array<int> *> */
+IGNORE_ARRAY_METHODS(mfem::Array<int> *)
+INSTANTIATE_ARRAY2(Array<int> *, Array<int>, intArray, 1)
diff --git a/mfem/_par/bilinearform.i b/mfem/_par/bilinearform.i
@@ -119,6 +119,10 @@ namespace mfem {
 
 } //end of namespace
 
+%import "../common/array_instantiation_macro.i"
+IGNORE_ARRAY_METHODS(mfem::BilinearForm *)
+INSTANTIATE_ARRAY0(BilinearForm *, BilinearForm, 1)
+
 %include "fem/bilinearform.hpp"
 
 // instatitate template methods

diff --git a/mfem/_par/device.i b/mfem/_par/device.i
@@ -15,8 +15,7 @@ import_array();
 %include "../common/typemap_macros.i"
 %include "../common/exception.i"
 
-
-
+%import "enzyme.i"
 %import "mem_manager.i"
 
 //  enforce that device class is created only once

diff --git a/mfem/_par/enzyme.i b/mfem/_par/enzyme.i
@@ -0,0 +1,13 @@
+%module(package="mfem._par") enzyme
+%{
+#include "mfem.hpp"
+%}
+
+%include "../common/existing_mfem_headers.i"
+#ifdef FILE_EXISTS_GENERAL_ENZYME
+
+%import "exception.i"
+
+%include "general/enzyme.hpp"
+
+#endif
diff --git a/mfem/_par/multigrid.i b/mfem/_par/multigrid.i
@@ -30,27 +30,24 @@ LIST_TO_MFEMOBJ_ARRAY_IN(const mfem::Array<mfem::Operator*>&,
 LIST_TO_MFEMOBJ_ARRAY_IN(const mfem::Array<mfem::Solver*>&,
 			   mfem::Solver*)
 LIST_TO_MFEMOBJ_BOOLARRAY_IN(const mfem::Array<bool>& )
-//%feature("director") mfem::PyGeometricMultigrid;
 
 %pythonprepend mfem::PyGeometricMultigrid::AppendBilinearForm %{
-   if not hasattr(self, "_forms"): self._forms = []
-   self._forms.append(form)
-   form.thisown = 0
-%}
-%pythonprepend mfem::PyGeometricMultigrid::AppendEssentialTDofs %{
-   if not hasattr(self, "_esss"): self._esss = []
-   self._esss.append(ess)
-   ess.thisown = 0
+  if not hasattr(self, "forms"):
+      self.forms=[]
+  self.forms.append(form)
+  form.thisown=0
 %}
+
 %feature("shadow") mfem::PyGeometricMultigrid::_pybfs %{
   @property
   def bfs(self):
-     return self._forms
- %}
+     return self._GetBilinearFormArray()
+%}
+
 %feature("shadow") mfem::PyGeometricMultigrid::_pyess %{
   @property
   def essentialTrueDofs(self):
-     return self._esss
+     return self._GetEssentialTrueDofs()
 %}
 
 %include "fem/multigrid.hpp"
@@ -60,15 +57,25 @@ LIST_TO_MFEMOBJ_BOOLARRAY_IN(const mfem::Array<bool>& )
 class PyGeometricMultigrid : public GeometricMultigrid
 {
 public:
- PyGeometricMultigrid(const FiniteElementSpaceHierarchy& fespaces_)
+  PyGeometricMultigrid(const FiniteElementSpaceHierarchy& fespaces_)
    : GeometricMultigrid(fespaces_){}
 
+  PyGeometricMultigrid(const FiniteElementSpaceHierarchy& fespaces_,
+                      const Array<int> &ess_bdr_)
+   : GeometricMultigrid(fespaces_, ess_bdr_){}
+
   void AppendBilinearForm(BilinearForm *form){
     bfs.Append(form);
   }
-  void AppendEssentialTDofs(Array<int> *ess){
-      essentialTrueDofs.Append(ess);
+
+  Array<BilinearForm*> *_GetBilinearFormArray(){
+    return &bfs;
+  }
+
+  Array<Array<int> *> *_GetEssentialTrueDofs(){
+    return &essentialTrueDofs;
   }
+
   void _pybfs(void){}
   void _pyess(void){}
 };

diff --git a/mfem/_par/navier_solver.i b/mfem/_par/navier_solver.i
@@ -44,10 +44,11 @@ OSTREAM_TYPEMAP(std::ostream&)
 %typemap(argout) double &time {
   %append_output(PyFloat_FromDouble(*$1));
 }
+\
 
-%ignore *::AddVelDirichletBC(VecFuncT *f, Array &);
-%ignore *::AddPresDirichletBC(ScalarFuncT *f, Array &);
-%ignore *::AddAccelTerm(VecFuncT *f, Array &);
+%ignore mfem::navier::NavierSolver::AddVelDirichletBC(VecFuncT *, Array<int> &);
+%ignore mfem::navier::NavierSolver::AddPresDirichletBC(ScalarFuncT *f, Array<int> &);
+%ignore mfem::navier::NavierSolver::AddAccelTerm(VecFuncT *f, Array<int> &);
 
 %include "miniapps/navier/navier_solver.hpp"
 
diff --git a/mfem/_par/setup.py b/mfem/_par/setup.py
@@ -121,7 +121,7 @@ def get_extensions():
                "tmop", "tmop_amr", "tmop_tools", "qspace", "qfunction",
                "quadinterpolator", "quadinterpolator_face",
                "submesh", "transfermap", "staticcond","sidredatacollection",
-               "psubmesh", "ptransfermap","navier_solver"]
+               "psubmesh", "ptransfermap", "enzyme"]
 
     if add_pumi == '1':
         from setup_local import puminc, pumilib
@@ -140,6 +140,8 @@ def get_extensions():
     if add_cuda == '1':
         from setup_local import cudainc
         include_dirs.append(cudainc)
+    else:
+        modules.append("navier_solver")
 
     if add_libceed == '1':
         from setup_local import libceedinc

diff --git a/mfem/_ser/array.i b/mfem/_ser/array.i
@@ -186,3 +186,7 @@ INSTANTIATE_ARRAY_NUMPYARRAY(uint, unsigned int, NPY_UINT)       // 32bit
 %template(intArray2D) mfem::Array2D<int>;
 %template(doubleArray2D) mfem::Array2D<double>;
 
+/* Array< Array<int> *> */
+IGNORE_ARRAY_METHODS(mfem::Array<int> *)
+INSTANTIATE_ARRAY2(Array<int> *, Array<int>, intArray, 1)
+
diff --git a/mfem/_ser/device.i b/mfem/_ser/device.i
@@ -15,6 +15,7 @@ import_array();
 %include "../common/typemap_macros.i"
 %include "../common/exception.i"
 
+%import "enzyme.i"
 %import "mem_manager.i"
 
 %include "../common/deprecation.i"

diff --git a/mfem/_ser/enzyme.i b/mfem/_ser/enzyme.i
@@ -0,0 +1,13 @@
+%module(package="mfem._ser") enzyme
+%{
+#include "mfem.hpp"
+%}
+
+%include "../common/existing_mfem_headers.i"
+#ifdef FILE_EXISTS_GENERAL_ENZYME
+
+%import "exception.i"
+
+%include "general/enzyme.hpp"
+
+#endif
diff --git a/mfem/_ser/multigrid.i b/mfem/_ser/multigrid.i
@@ -31,33 +31,26 @@ LIST_TO_MFEMOBJ_ARRAY_IN(const mfem::Array<mfem::Operator*>&,
 LIST_TO_MFEMOBJ_ARRAY_IN(const mfem::Array<mfem::Solver*>&,
 			   mfem::Solver*)
 LIST_TO_MFEMOBJ_BOOLARRAY_IN(const mfem::Array<bool>& )
- /*
-ObjectArrayInput(mfem::Solver *);
-ObjectArrayInput(mfem::Operator *);
-BoolArrayInput(bool);
- */
-//%feature("director") mfem::PyGeometricMultigrid;
 
 %pythonprepend mfem::PyGeometricMultigrid::AppendBilinearForm %{
-   if not hasattr(self, "_forms"): self._forms = []
-   self._forms.append(form)
-   form.thisown = 0
-%}
-%pythonprepend mfem::PyGeometricMultigrid::AppendEssentialTDofs %{
-   if not hasattr(self, "_esss"): self._esss = []
-   self._esss.append(ess)
-   ess.thisown = 0
+  if not hasattr(self, "forms"):
+      self.forms=[]
+  self.forms.append(form)
+  form.thisown=0
 %}
+
 %feature("shadow") mfem::PyGeometricMultigrid::_pybfs %{
   @property
   def bfs(self):
-     return self._forms
- %}
+     return self._GetBilinearFormArray()
+%}
+
 %feature("shadow") mfem::PyGeometricMultigrid::_pyess %{
   @property
   def essentialTrueDofs(self):
-     return self._esss
+     return self._GetEssentialTrueDofs()
 %}
+
 %include "fem/multigrid.hpp"
 
 %inline %{
@@ -67,18 +60,24 @@ class PyGeometricMultigrid : public GeometricMultigrid
 public:
  PyGeometricMultigrid(const FiniteElementSpaceHierarchy& fespaces_)
    : GeometricMultigrid(fespaces_){}
+ PyGeometricMultigrid(const FiniteElementSpaceHierarchy& fespaces_,
+                      const Array<int> &ess_bdr_)
+   : GeometricMultigrid(fespaces_, ess_bdr_){}
 
   void AppendBilinearForm(BilinearForm *form){
     bfs.Append(form);
   }
-  void AppendEssentialTDofs(Array<int> *ess){
-      essentialTrueDofs.Append(ess);
+
+  Array<BilinearForm*> *_GetBilinearFormArray(){
+    return &bfs;
+  }
+
+  Array<Array<int> *> *_GetEssentialTrueDofs(){
+    return &essentialTrueDofs;
   }
+
   void _pybfs(void){}
   void _pyess(void){}
 };
   } /* end of namespace */
 %}
-
-
-
diff --git a/mfem/_ser/setup.py b/mfem/_ser/setup.py
@@ -102,7 +102,8 @@ def get_extensions():
                "transfer", "std_vectors",
                "tmop", "tmop_amr", "tmop_tools", "qspace", "qfunction",
                "quadinterpolator", "quadinterpolator_face",
-               "submesh", "transfermap", "staticcond", "sidredatacollection"]
+               "submesh", "transfermap", "staticcond",
+               "sidredatacollection", "enzyme"]
 
     if add_cuda == '1':
         from setup_local import cudainc