Enable proteus with mass and diffusion integrators

CMakeLists.txt

@@ -175,6 +175,18 @@ if (MFEM_USE_CUDA)
    mfem_culib_set_libraries(CUBLAS cublas)
 endif()
 
+# Proteus config
+if (ENABLE_JIT)
+  if (NOT "${PROTEUS_INSTALL_DIR}" STREQUAL "")
+    find_package(proteus REQUIRED PATHS "${PROTEUS_INSTALL_DIR}")
+    message(STATUS "${PROTEUS_INSTALL_DIR}/include")
+    include_directories("${PROTEUS_INSTALL_DIR}/include")
+  else()
+    #add_subdirectory(tpl/RAJA/tpl/proteus)
+    message(FATAL_ERROR "Couldn't find Proteus, but ENABLE_JIT was turned on")
+  endif()
+endif()
+
 if (XSDK_ENABLE_C)
    enable_language(C)
 endif()
@@ -242,6 +254,7 @@ endif()
 
 # AMD HIP
 if (MFEM_USE_HIP)
+  enable_language(HIP)
   if (HIP_ARCH)
     message(STATUS "Using HIP architecture: ${HIP_ARCH}")
     set(GPU_TARGETS "${HIP_ARCH}" CACHE STRING "HIP targets to compile for")
@@ -692,6 +705,9 @@ set(MFEM_INSTALL_DIR ${CMAKE_INSTALL_PREFIX})
 mfem_add_library(mfem ${SOURCES} ${HEADERS} ${MASTER_HEADERS})
 # message(STATUS "TPL_LIBRARIES = ${TPL_LIBRARIES}")
 target_link_libraries(mfem PUBLIC ${TPL_LIBRARIES} ${TPL_TARGETS})
+if (ENABLE_JIT)
+  add_proteus(mfem)
+endif()
 if (TPL_TARGETS)
   add_dependencies(mfem ${TPL_TARGETS})
 endif()

config/defaults.cmake

@@ -20,6 +20,7 @@ endif()
 
 # MFEM options. Set to mimic the default "defaults.mk" file.
 option(BUILD_SHARED_LIBS "Enable shared library build of MFEM" OFF)
+option(ENABLE_JIT "Enable JIT compilation via proteus via either submodule or install" On)
 option(MFEM_USE_MPI "Enable MPI parallel build" OFF)
 option(MFEM_USE_METIS "Enable METIS usage" ${MFEM_USE_MPI})
 set(MFEM_PRECISION "double" CACHE STRING

fem/integ/bilininteg_diffusion_kernels.hpp

@@ -108,7 +108,7 @@ inline void PADiffusionDiagonal2D(const int NE,
    // store necessary entries
    auto D = Reshape(d.Read(), Q1D*Q1D, symmetric ? 3 : 4, NE);
    auto Y = Reshape(y.ReadWrite(), D1D, D1D, NE);
-   mfem::forall(NE, [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall(NE, [=, symmetric = proteus::jit_variable(symmetric)] MFEM_HOST_DEVICE (int e)
    {
       const int D1D = T_D1D ? T_D1D : d1d;
       const int Q1D = T_Q1D ? T_Q1D : q1d;
@@ -191,7 +191,7 @@ inline void SmemPADiffusionDiagonal2D(const int NE,
    auto g = Reshape(g_.Read(), Q1D, D1D);
    auto D = Reshape(d_.Read(), Q1D*Q1D, symmetric ? 3 : 4, NE);
    auto Y = Reshape(y_.ReadWrite(), D1D, D1D, NE);
-   mfem::forall_2D_batch(NE, Q1D, Q1D, NBZ, [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall_2D_batch(NE, Q1D, Q1D, NBZ, [=, symmetric = proteus::jit_variable(symmetric)] MFEM_HOST_DEVICE (int e)
    {
       const int tidz = MFEM_THREAD_ID(z);
       const int D1D = T_D1D ? T_D1D : d1d;
@@ -285,7 +285,7 @@ inline void PADiffusionDiagonal3D(const int NE,
    auto G = Reshape(g.Read(), Q1D, D1D);
    auto Q = Reshape(d.Read(), Q1D*Q1D*Q1D, symmetric ? 6 : 9, NE);
    auto Y = Reshape(y.ReadWrite(), D1D, D1D, D1D, NE);
-   mfem::forall(NE, [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall(NE, [=, symmetric = proteus::jit_variable(symmetric)] MFEM_HOST_DEVICE (int e)
    {
       const int D1D = T_D1D ? T_D1D : d1d;
       const int Q1D = T_Q1D ? T_Q1D : q1d;
@@ -386,7 +386,7 @@ inline void SmemPADiffusionDiagonal3D(const int NE,
    auto g = Reshape(g_.Read(), Q1D, D1D);
    auto D = Reshape(d_.Read(), Q1D*Q1D*Q1D, symmetric ? 6 : 9, NE);
    auto Y = Reshape(y_.ReadWrite(), D1D, D1D, D1D, NE);
-   mfem::forall_3D(NE, Q1D, Q1D, Q1D, [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall_3D(NE, Q1D, Q1D, Q1D, [=, symmetric = proteus::jit_variable(symmetric)] MFEM_HOST_DEVICE (int e)
    {
       const int tidz = MFEM_THREAD_ID(z);
       const int D1D = T_D1D ? T_D1D : d1d;
@@ -534,7 +534,7 @@ inline void PADiffusionApply2D(const int NE,
    auto D = Reshape(d_.Read(), Q1D*Q1D, symmetric ? 3 : 4, NE);
    auto X = Reshape(x_.Read(), D1D, D1D, NE);
    auto Y = Reshape(y_.ReadWrite(), D1D, D1D, NE);
-   mfem::forall(NE, [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall(NE, [=, symmetric = proteus::jit_variable(symmetric)] MFEM_HOST_DEVICE (int e)
    {
       const int D1D = T_D1D ? T_D1D : d1d;
       const int Q1D = T_Q1D ? T_Q1D : q1d;
@@ -658,7 +658,7 @@ inline void SmemPADiffusionApply2D(const int NE,
    auto D = Reshape(d_.Read(), Q1D*Q1D, symmetric ? 3 : 4, NE);
    auto x = Reshape(x_.Read(), D1D, D1D, NE);
    auto Y = Reshape(y_.ReadWrite(), D1D, D1D, NE);
-   mfem::forall_2D_batch(NE, Q1D, Q1D, NBZ, [=] MFEM_HOST_DEVICE(int e)
+   mfem::forall_2D_batch(NE, Q1D, Q1D, NBZ, [=, symmetric = proteus::jit_variable(symmetric)] MFEM_HOST_DEVICE(int e)
    {
       const int tidz = MFEM_THREAD_ID(z);
       const int D1D = T_D1D ? T_D1D : d1d;
@@ -815,7 +815,7 @@ inline void PADiffusionApply3D(const int NE,
    auto D = Reshape(d_.Read(), Q1D*Q1D*Q1D, symmetric ? 6 : 9, NE);
    auto X = Reshape(x_.Read(), D1D, D1D, D1D, NE);
    auto Y = Reshape(y_.ReadWrite(), D1D, D1D, D1D, NE);
-   mfem::forall(NE, [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall(NE, [=, symmetric = proteus::jit_variable(symmetric)] MFEM_HOST_DEVICE (int e)
    {
       const int D1D = T_D1D ? T_D1D : d1d;
       const int Q1D = T_Q1D ? T_Q1D : q1d;
@@ -1010,7 +1010,7 @@ inline void SmemPADiffusionApply3D(const int NE,
    auto x = Reshape(x_.Read(), D1D, D1D, D1D, NE);
    auto y = Reshape(y_.ReadWrite(), D1D, D1D, D1D, NE);
    MFEM_VERIFY(D1D <= Q1D, "THREAD_DIRECT requires D1D <= Q1D");
-   mfem::forall_3D(NE, Q1D, Q1D, Q1D, [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall_3D(NE, Q1D, Q1D, Q1D, [=, symmetric = proteus::jit_variable(symmetric)] MFEM_HOST_DEVICE (int e)
    {
       const int D1D = T_D1D ? T_D1D : d1d;
       const int Q1D = T_Q1D ? T_Q1D : q1d;

fem/integ/bilininteg_mass_kernels.hpp

@@ -38,7 +38,9 @@ static void PAMassAssembleDiagonal1D(const int NE,
    auto B = Reshape(b.Read(), Q1D, D1D);
    auto D = Reshape(d.Read(), Q1D, NE);
    auto Y = Reshape(y.ReadWrite(), D1D, NE);
-   mfem::forall(NE, [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall(NE, [=,
+         D1D = proteus::jit_variable(D1D),
+      Q1D = proteus::jit_variable(Q1D)] MFEM_HOST_DEVICE (int e)
    {
       for (int dx = 0; dx < D1D; ++dx)
       {
@@ -142,10 +144,11 @@ inline void PAMassAssembleDiagonal2D(const int NE,
    auto B = Reshape(b.Read(), Q1D, D1D);
    auto D = Reshape(d.Read(), Q1D, Q1D, NE);
    auto Y = Reshape(y.ReadWrite(), D1D, D1D, NE);
-   mfem::forall(NE, [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall(NE, [=,D1D = proteus::jit_variable(D1D),
+      Q1D = proteus::jit_variable(Q1D)] MFEM_HOST_DEVICE (int e)
    {
-      const int D1D = T_D1D ? T_D1D : d1d;
-      const int Q1D = T_Q1D ? T_Q1D : q1d;
+      // const int D1D = T_D1D ? T_D1D : d1d;
+      // const int Q1D = T_Q1D ? T_Q1D : q1d;
       constexpr int MD1 = T_D1D ? T_D1D : DofQuadLimits::MAX_D1D;
       constexpr int MQ1 = T_Q1D ? T_Q1D : DofQuadLimits::MAX_Q1D;
       real_t QD[MQ1][MD1];
@@ -203,11 +206,13 @@ inline void SmemPAMassAssembleDiagonal2D(const int NE,
    auto b = Reshape(b_.Read(), Q1D, D1D);
    auto D = Reshape(d_.Read(), Q1D, Q1D, NE);
    auto Y = Reshape(y_.ReadWrite(), D1D, D1D, NE);
-   mfem::forall_2D_batch(NE, Q1D, Q1D, NBZ, [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall_2D_batch(NE, Q1D, Q1D, NBZ, [=,
+         D1D = proteus::jit_variable(D1D),
+      Q1D = proteus::jit_variable(Q1D)] MFEM_HOST_DEVICE (int e)
    {
       const int tidz = MFEM_THREAD_ID(z);
-      const int D1D = T_D1D ? T_D1D : d1d;
-      const int Q1D = T_Q1D ? T_Q1D : q1d;
+      // const int D1D = T_D1D ? T_D1D : d1d;
+      // const int Q1D = T_Q1D ? T_Q1D : q1d;
       constexpr int MQ1 = T_Q1D ? T_Q1D : DofQuadLimits::MAX_Q1D;
       constexpr int MD1 = T_D1D ? T_D1D : DofQuadLimits::MAX_D1D;
       MFEM_SHARED real_t B[MQ1][MD1];
@@ -266,10 +271,11 @@ inline void PAMassAssembleDiagonal3D(const int NE,
    auto B = Reshape(b.Read(), Q1D, D1D);
    auto D = Reshape(d.Read(), Q1D, Q1D, Q1D, NE);
    auto Y = Reshape(y.ReadWrite(), D1D, D1D, D1D, NE);
-   mfem::forall(NE, [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall(NE, [=,      D1D = proteus::jit_variable(D1D),
+      Q1D = proteus::jit_variable(Q1D)] MFEM_HOST_DEVICE (int e)
    {
-      const int D1D = T_D1D ? T_D1D : d1d;
-      const int Q1D = T_Q1D ? T_Q1D : q1d;
+      // const int D1D = T_D1D ? T_D1D : d1d;
+      // const int Q1D = T_Q1D ? T_Q1D : q1d;
       constexpr int MD1 = T_D1D ? T_D1D : DofQuadLimits::MAX_D1D;
       constexpr int MQ1 = T_Q1D ? T_Q1D : DofQuadLimits::MAX_Q1D;
       real_t QQD[MQ1][MQ1][MD1];
@@ -338,11 +344,14 @@ inline void SmemPAMassAssembleDiagonal3D(const int NE,
    auto b = Reshape(b_.Read(), Q1D, D1D);
    auto D = Reshape(d_.Read(), Q1D, Q1D, Q1D, NE);
    auto Y = Reshape(y_.ReadWrite(), D1D, D1D, D1D, NE);
-   mfem::forall_3D(NE, Q1D, Q1D, Q1D, [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall_3D(NE, Q1D, Q1D, Q1D, [=,
+      D1D = proteus::jit_variable(D1D),
+      Q1D = proteus::jit_variable(Q1D)
+   ] MFEM_HOST_DEVICE (int e)
    {
       const int tidz = MFEM_THREAD_ID(z);
-      const int D1D = T_D1D ? T_D1D : d1d;
-      const int Q1D = T_Q1D ? T_Q1D : q1d;
+      //const int D1D = T_D1D ? T_D1D : d1d;
+      //const int Q1D = T_Q1D ? T_Q1D : q1d;
       constexpr int MD1 = T_D1D ? T_D1D : DofQuadLimits::MAX_D1D;
       constexpr int MQ1 = T_Q1D ? T_Q1D : DofQuadLimits::MAX_Q1D;
       MFEM_SHARED real_t B[MQ1][MD1];
@@ -1208,7 +1217,7 @@ inline void EAMassAssemble2D(const int NE,
    auto D = Reshape(padata.Read(), Q1D, Q1D, NE);
    auto M = Reshape(add ? eadata.ReadWrite() : eadata.Write(), D1D, D1D, D1D, D1D,
                     NE);
-   mfem::forall_2D(NE, D1D, D1D, [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall_2D(NE, D1D, D1D, [=, add = proteus::jit_variable(add)] MFEM_HOST_DEVICE (int e)
    {
       const int D1D = T_D1D ? T_D1D : d1d;
       const int Q1D = T_Q1D ? T_Q1D : q1d;
@@ -1281,7 +1290,7 @@ inline void EAMassAssemble3D(const int NE,
    auto D = Reshape(padata.Read(), Q1D, Q1D, Q1D, NE);
    auto M = Reshape(add ? eadata.ReadWrite() : eadata.Write(), D1D, D1D, D1D, D1D,
                     D1D, D1D, NE);
-   mfem::forall_3D(NE, D1D, D1D, D1D, [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall_3D(NE, D1D, D1D, D1D, [=, add = proteus::jit_variable(add)] MFEM_HOST_DEVICE (int e)
    {
       const int D1D = T_D1D ? T_D1D : d1d;
       const int Q1D = T_Q1D ? T_Q1D : q1d;

fem/integ/bilininteg_mass_pa.cpp

@@ -69,7 +69,10 @@ void MassIntegrator::AssemblePA(const FiniteElementSpace &fes)
       const auto C =
          const_c ? Reshape(coeff.Read(), 1, 1) : Reshape(coeff.Read(), NQ, NE);
       auto v = Reshape(pa_data.Write(), NQ, NE);
-      mfem::forall(NQ, NE, [=] MFEM_HOST_DEVICE(int q, int e)
+      mfem::forall(NQ, NE, [=,
+         const_c = proteus::jit_variable(const_c),
+         by_val  = proteus::jit_variable(by_val)]
+         MFEM_HOST_DEVICE(int q, int e)
       {
          const real_t detJ = J(q, e);
          const real_t coeff = const_c ? C(0, 0) : C(q, e);
@@ -115,7 +118,8 @@ void MassIntegrator::AssemblePABoundary(const FiniteElementSpace &fes)
       const auto C = const_c ? Reshape(coeff.Read(), 1, 1)
                      : Reshape(coeff.Read(), NQ, NE);
       auto v = Reshape(pa_data.Write(), NQ, NE);
-      mfem::forall(NQ, NE, [=] MFEM_HOST_DEVICE(int q, int e)
+      mfem::forall(NQ, NE, [=, const_c = proteus::jit_variable(const_c),
+         by_val  = proteus::jit_variable(by_val)] MFEM_HOST_DEVICE(int q, int e)
       {
          const real_t detJ = J(q, e);
          const real_t coeff = const_c ? C(0, 0) : C(q, e);

fem/integ/bilininteg_vecdiffusion_pa.cpp

@@ -86,7 +86,9 @@ static void PAVectorDiffusionSetup2D(const int Q1D,
                   Reshape(c.Read(), NQ, NE);
 
 
-   mfem::forall(NE, [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall(NE, [=,
+      const_c = proteus::jit_variable(const_c),
+      NQ = proteus::jit_variable(NQ)] MFEM_HOST_DEVICE (int e)
    {
       for (int q = 0; q < NQ; ++q)
       {
@@ -122,7 +124,9 @@ static void PAVectorDiffusionSetup3D(const int Q1D,
                   Reshape(c.Read(), NQ,NE);
 
 
-   mfem::forall(NE, [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall(NE, [=,
+      const_c = proteus::jit_variable(const_c),
+      NQ = proteus::jit_variable(NQ)] MFEM_HOST_DEVICE (int e)
    {
       for (int q = 0; q < NQ; ++q)
       {
@@ -242,7 +246,7 @@ void VectorDiffusionIntegrator::AssemblePA(const FiniteElementSpace &fes)
       const auto C = const_c ? Reshape(coeff.Read(), 1,1) :
                      Reshape(coeff.Read(), NQ,ne);
 
-      mfem::forall(ne, [=] MFEM_HOST_DEVICE (int e)
+      mfem::forall(ne, [=, const_c = proteus::jit_variable(const_c)] MFEM_HOST_DEVICE (int e)
       {
          for (int q = 0; q < NQ; ++q)
          {

general/forall.hpp

@@ -19,6 +19,7 @@
 #include "device.hpp"
 #include "mem_manager.hpp"
 #include "../linalg/dtensor.hpp"
+#include <proteus/JitInterface.hpp>
 #ifdef MFEM_USE_MPI
 #include <_hypre_utilities.h>
 #endif
@@ -661,7 +662,7 @@ struct CuWrap<3>
 /// HIP backend
 #if defined(MFEM_USE_HIP) && defined(__HIP__)
 
-template <typename BODY> __global__ static
+template <typename BODY> __global__ static __attribute__((annotate("jit", 1)))
 void HipKernel1D(const int N, BODY body)
 {
    const int k = hipBlockDim_x*hipBlockIdx_x + hipThreadIdx_x;
@@ -669,7 +670,7 @@ void HipKernel1D(const int N, BODY body)
    body(k);
 }
 
-template <typename BODY> __global__ static
+template <typename BODY> __global__ static __attribute__((annotate("jit", 1)))
 void HipKernel2D(const int N, BODY body)
 {
    const int k = hipBlockIdx_x*hipBlockDim_z + hipThreadIdx_z;
@@ -677,7 +678,7 @@ void HipKernel2D(const int N, BODY body)
    body(k);
 }
 
-template <typename BODY> __global__ static
+template <typename BODY> __global__ static __attribute__((annotate("jit", 1)))
 void HipKernel3D(const int N, BODY body)
 {
    for (int k = hipBlockIdx_x; k < N; k += hipGridDim_x) { body(k); }
@@ -876,6 +877,7 @@ inline void forall(int Nx, int Ny, int Nz, lambda &&body)
 {
    if (Device::Allows(Backend::DEVICE_MASK))
    {
+      proteus::register_lambda(body);
       forall(Nx * Ny * Nz, [=] MFEM_HOST_DEVICE(int idx)
       {
          int i = idx % Nx;

tests/benchmarks/CMakeLists.txt

@@ -32,16 +32,20 @@ function(add_benchmark name)
     endif(MFEM_USE_CUDA)
 
     if (MFEM_USE_HIP)
-        set_property(SOURCE ${${NAME}_BENCH_SRCS} PROPERTY LANGUAGE 
-                     HIP_SOURCE_PROPERTY_FORMAT TRUE)
+        # set_property(SOURCE ${${NAME}_BENCH_SRCS} PROPERTY LANGUAGE
+        #              HIP)
+        # set_property(SOURCE ${${NAME}_BENCH_SRCS} PROPERTY LANGUAGE
+        #              HIP_SOURCE_PROPERTY_FORMAT TRUE)
+        # set_property(SOURCE ${${NAME}_BENCH_SRCS} PROPERTY LANGUAGE
+        #              CMAKE_HIP_SOURCE_PROPERTY_FORMAT TRUE)
     endif(MFEM_USE_HIP)
-
-    add_executable(bench_${name} ${${NAME}_BENCH_SRCS})
+    # set(CMAKE_HIP_SOURCE_PROPERTY_FORMAT TRUE)
+    mfem_add_executable(bench_${name} ${${NAME}_BENCH_SRCS})
     target_link_libraries(bench_${name} mfem pthread)
     add_dependencies(${MFEM_ALL_BENCHMARKS_TARGET_NAME} bench_${name})
 
-    add_test(NAME bench_${name}_cpu
-             COMMAND bench_${name} --benchmark_context=device=cpu)
+    # add_test(NAME bench_${name}_cpu
+    #          COMMAND bench_${name} --benchmark_context=device=cpu)
 
     if (MFEM_USE_CUDA)
         add_test(NAME bench_${name}_cuda

tests/benchmarks/bench_ceed.cpp

@@ -132,19 +132,19 @@ struct Problem : public BakeOff<VDIM, GLL>
 
 /// Bake-off Problems (BPs)
 #define BakeOff_Problem(i, Kernel, VDIM, p_eq_q)                     \
-   static void BP##i(bm::State &state)                               \
+   static void BP##i##Kernel(bm::State &state)                               \
    {                                                                 \
       Problem<Kernel##Integrator, VDIM, p_eq_q> ker(state.range(0)); \
       while (state.KeepRunning()) { ker.benchmark(); }               \
       state.counters["MDof/s"] =                                     \
          bm::Counter(ker.SumMdofs(), bm::Counter::kIsRate);          \
    }                                                                 \
-   BENCHMARK(BP##i)->DenseRange(1, 6)->Unit(bm::kMillisecond);
+   BENCHMARK(BP##i##Kernel)->DenseRange(1, 8)->Unit(bm::kMillisecond);
 
-/// BP1: scalar PCG with mass matrix, q=p+2
+// BP1: scalar PCG with mass matrix, q=p+2
 BakeOff_Problem(1, Mass, 1, false)
 
-/// BP2: vector PCG with mass matrix, q=p+2
+// /// BP2: vector PCG with mass matrix, q=p+2
 BakeOff_Problem(2, VectorMass, 3, false)
 
 /// BP3: scalar PCG with stiffness matrix, q=p+2
@@ -191,14 +191,14 @@ struct Kernel : public BakeOff<VDIM, GLL>
 
 /// Generic CEED BKi
 #define BakeOff_Kernel(i, KER, VDIM, GLL)                      \
-   static void BK##i(bm::State &state)                         \
+   static void BK##i##KER(bm::State &state)                         \
    {                                                           \
       Kernel<KER##Integrator, VDIM, GLL> ker(state.range(0));  \
       while (state.KeepRunning()) { ker.benchmark(); }         \
       state.counters["MDof/s"] =                               \
          bm::Counter(ker.SumMdofs(), bm::Counter::kIsRate);    \
    }                                                           \
-   BENCHMARK(BK##i)->DenseRange(1, 6)->Unit(bm::kMillisecond);
+   BENCHMARK(BK##i##KER)->DenseRange(1, 8)->Unit(bm::kMillisecond);
 
 /// BK1: scalar E-vector-to-E-vector evaluation of mass matrix, q=p+2
 BakeOff_Kernel(1, Mass, 1, false)