draft: Proteus lambda redesign RAJA-side support

CMakeLists.txt

@@ -333,6 +333,10 @@ install(FILES
   ${PROJECT_BINARY_DIR}/include/RAJA/config.hpp
   DESTINATION "include/RAJA/")
 
+install(FILES
+  ${PROJECT_BINARY_DIR}/include/RAJA/policy/device.hpp
+  DESTINATION "include/RAJA/policy/")
+
 
 # Setup internal RAJA configuration options
 include(cmake/SetupRajaConfig.cmake)

cmake/RAJAMacros.cmake

@@ -19,7 +19,7 @@ macro(raja_link_include_proteus)
     add_proteus(${arg_NAME})
     target_include_directories(${arg_NAME}
       PUBLIC
-      "${PROTEUS_HEADERS_DIR}"
+      $<BUILD_INTERFACE:${PROTEUS_HEADERS_DIR}>
     )
   endif()
 endmacro(raja_link_include_proteus)

cmake/SetupProteus.cmake

@@ -20,7 +20,7 @@ else()
   FetchContent_Declare(
     proteus
     GIT_REPOSITORY https://github.com/Olympus-HPC/proteus.git
-    GIT_TAG        257707cf7e60452ed38161f6429421be303ddaf3
+    GIT_TAG        44415066d3c91729a23137101a0f08a170101fb5
     )
   FetchContent_MakeAvailable(proteus)
   # Re-enable tests if specified by user.
@@ -30,4 +30,16 @@ else()
 endif()
 # We don't explicitly link against ProteusPass, but it is required to be
 #available as an LLVM pass, so manually enforce order
-add_dependencies(RAJA ProteusPass)
+target_include_directories(RAJA
+  PUBLIC
+  $<BUILD_INTERFACE:${PROTEUS_HEADERS_DIR}>
+  $<INSTALL_INTERFACE:include>)
+
+# RAJA's public JIT headers include Proteus headers, so install those headers
+# into RAJA's include tree instead of exporting a build-tree path.
+if (EXISTS "${PROTEUS_HEADERS_DIR}/proteus")
+  install(DIRECTORY "${PROTEUS_HEADERS_DIR}/proteus"
+    DESTINATION include)
+endif()
+
+add_dependencies(RAJA ProteusPass LambdaPass)

examples/CMakeLists.txt

@@ -27,6 +27,12 @@ if (RAJA_ENABLE_JIT)
   raja_add_executable(
     NAME forall-jit
     SOURCES forall-jit.cpp)
+  raja_add_executable(
+    NAME kernel-jit
+    SOURCES kernel-jit.cpp)
+  raja_add_executable(
+    NAME launch-jit
+    SOURCES launch-jit.cpp)
 endif()
 
 raja_add_executable(

examples/kernel-jit.cpp

@@ -0,0 +1,152 @@
+//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
+// Copyright (c) Lawrence Livermore National Security, LLC and other
+// RAJA Project Developers. See top-level LICENSE and COPYRIGHT
+// files for dates and other details. No copyright assignment is required
+// to contribute to RAJA.
+//
+// SPDX-License-Identifier: (BSD-3-Clause)
+//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
+// This example mirrors forall-jit.cpp, but uses RAJA::kernel to drive the
+// outer parallel loop.
+
+#include <cstdlib>
+#include <iostream>
+
+#include "RAJA/RAJA.hpp"
+#include "RAJA/util/Timer.hpp"
+
+#if defined(RAJA_ENABLE_CUDA)
+using kernel_policy = RAJA::KernelPolicy<
+    RAJA::statement::CudaKernelFixed<256,
+        RAJA::statement::Tile<0, RAJA::tile_fixed<256>,
+                              RAJA::cuda_block_x_direct,
+            RAJA::statement::For<0, RAJA::cuda_thread_x_direct,
+                RAJA::statement::Lambda<0>>>>>;
+#elif defined(RAJA_ENABLE_HIP)
+using kernel_policy = RAJA::KernelPolicy<
+    RAJA::statement::HipKernelFixed<256,
+        RAJA::statement::Tile<0, RAJA::tile_fixed<256>,
+                              RAJA::hip_block_x_direct,
+            RAJA::statement::For<0, RAJA::hip_thread_x_direct,
+                RAJA::statement::Lambda<0>>>>>;
+#elif defined(RAJA_ENABLE_OPENMP)
+using kernel_policy = RAJA::KernelPolicy<
+    RAJA::statement::For<0, RAJA::omp_parallel_for_exec,
+        RAJA::statement::Lambda<0>>>;
+#else
+using kernel_policy = RAJA::KernelPolicy<
+    RAJA::statement::For<0, RAJA::seq_exec, RAJA::statement::Lambda<0>>>;
+#endif
+
+int main(int argc, char **argv) {
+  if (argc < 5) {
+    std::cout
+        << "Expected ./bin/kernel-jit <num rows> <num cols> <num matrices> "
+           "<boolean accum>, for example ./bin/kernel-jit 24 16 1000000 1\n";
+    return 1;
+  }
+
+  int a = std::stoi(argv[1]);
+  int b = std::stoi(argv[2]);
+  int N = std::stoi(argv[3]);
+  bool accum = std::stoi(argv[4]);
+
+#ifdef RAJA_ENABLE_CUDA
+  using resource_policy = RAJA::cuda_exec<256>;
+#elif defined(RAJA_ENABLE_HIP)
+  using resource_policy = RAJA::hip_exec<256>;
+#elif defined(RAJA_ENABLE_OPENMP)
+  using resource_policy = RAJA::omp_parallel_for_exec;
+#else
+  using resource_policy = RAJA::seq_exec;
+#endif
+
+  auto res = RAJA::resources::get_default_resource<resource_policy>();
+  double *A_ptr = res.template allocate<double>(N * a * b);
+  auto A = RAJA::make_permuted_view<RAJA::layout_right>(A_ptr, N, a, b);
+  double *B_ptr = res.template allocate<double>(N * a * b);
+  auto B = RAJA::make_permuted_view<RAJA::layout_right>(B_ptr, N, b, a);
+  double *C_ptr = res.template allocate<double>(N * a * a);
+  auto C = RAJA::make_permuted_view<RAJA::layout_right>(C_ptr, N, a, a);
+
+  auto Seg = RAJA::RangeSegment(0, N);
+
+  RAJA::kernel<kernel_policy>(RAJA::make_tuple(Seg), [=](int i) {
+    for (int row = 0; row < a; ++row) {
+      for (int col = 0; col < b; ++col) {
+        A(i, row, col) = 0;
+        B(i, row, col) = 0;
+        C(i, row, col) = 0;
+      }
+    }
+  });
+
+  RAJA::Timer aot_timer;
+  proteus::disable();
+  aot_timer.start();
+
+  RAJA::kernel<kernel_policy>(RAJA::make_tuple(Seg), [=](int i) {
+    for (int row = 0; row < a; ++row) {
+      for (int col = 0; col < b; ++col) {
+        A(i, row, col) = i % row;
+        B(i, row, col) = i % col;
+        C(i, row, col) = i % col;
+      }
+    }
+  });
+
+  RAJA::kernel<kernel_policy>(RAJA::make_tuple(Seg), [=](int i) {
+    for (int row = 0; row < a; ++row) {
+      for (int col = 0; col < b; ++col) {
+        if (!accum) {
+          C(i, row, col) = A(i, row, col) * B(i, col, row);
+        } else {
+          C(i, row, col) += A(i, row, col) * B(i, col, row);
+        }
+      }
+    }
+  });
+
+  aot_timer.stop();
+  proteus::enable();
+  std::cout << "aot total time = " << aot_timer.elapsed() << "\n";
+
+  RAJA::Timer jit_timer;
+  jit_timer.start();
+
+  RAJA::kernel<kernel_policy>(
+      RAJA::make_tuple(Seg),
+      [=, a = RAJA_JIT_VARIABLE(a), b = RAJA_JIT_VARIABLE(b)](int i)
+          RAJA_JIT_COMPILE {
+            for (int row = 0; row < a; ++row) {
+              for (int col = 0; col < b; ++col) {
+                A(i, row, col) = i % row;
+                B(i, row, col) = i % col;
+                C(i, row, col) = i % col;
+              }
+            }
+          });
+
+  RAJA::kernel<kernel_policy>(
+      RAJA::make_tuple(Seg),
+      [=, a = RAJA_JIT_VARIABLE(a), b = RAJA_JIT_VARIABLE(b),
+       accum = RAJA_JIT_VARIABLE(accum)](int i) RAJA_JIT_COMPILE {
+        for (int row = 0; row < a; ++row) {
+          for (int col = 0; col < b; ++col) {
+            if (!accum) {
+              C(i, row, col) = A(i, row, col) * B(i, col, row);
+            } else {
+              C(i, row, col) += A(i, row, col) * B(i, col, row);
+            }
+          }
+        }
+      });
+
+  jit_timer.stop();
+
+  std::cout << "jit total time = " << jit_timer.elapsed() << "\n";
+  std::cout << "speedup = " << aot_timer.elapsed() / jit_timer.elapsed()
+            << "\n";
+
+  return 0;
+}