Skip to content

Commit 60e3989

Browse files
johnbowen42johnbowen42
committed
Add proteus header and include statements
1 parent cfd2692 commit 60e3989

9 files changed

Lines changed: 84 additions & 52 deletions

File tree

‎cmake/SetupRajaOptions.cmake‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ option(RAJA_ENABLE_RUNTIME_PLUGINS "Enable support for loading plugins at runtim
3535
option(RAJA_ALLOW_INCONSISTENT_OPTIONS "Enable inconsistent values for ENABLE_X and RAJA_ENABLE_X options" Off)
3636

3737
option(RAJA_ENABLE_DESUL_ATOMICS "Enable support of desul atomics" Off)
38+
option(RAJA_ENABLE_JIT "Enable JIT compilation for RAJA kernels" Off)
3839
set(DESUL_ENABLE_TESTS Off CACHE BOOL "")
3940

4041
set(TEST_DRIVER "" CACHE STRING "driver used to wrap test commands")

‎include/RAJA/policy/cuda/forall.hpp‎

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434

3535
#include "RAJA/util/macros.hpp"
3636
#include "RAJA/util/types.hpp"
37+
#include "RAJA/util/Jit.hpp"
3738

3839
#include "RAJA/internal/fault_tolerance.hpp"
3940

@@ -380,11 +381,12 @@ template<typename EXEC_POL,
380381
IterationMapping>::value &&
381382
(IterationGetter::block_size > 0),
382383
size_t> BlockSize = IterationGetter::block_size>
383-
__launch_bounds__(BlockSize, BlocksPerSM) __global__ RAJA_JIT_COMPILE(3)
384-
void forallp_cuda_kernel(const RAJA_CUDA_GRID_CONSTANT LOOP_BODY loop_body,
385-
const RAJA_CUDA_GRID_CONSTANT Iterator idx,
386-
const RAJA_CUDA_GRID_CONSTANT IndexType length,
387-
ForallParam f_params)
384+
__launch_bounds__(BlockSize, BlocksPerSM) __global__
385+
RAJA_JIT_COMPILE_ARGS(3) void forallp_cuda_kernel(
386+
const RAJA_CUDA_GRID_CONSTANT LOOP_BODY loop_body,
387+
const RAJA_CUDA_GRID_CONSTANT Iterator idx,
388+
const RAJA_CUDA_GRID_CONSTANT IndexType length,
389+
ForallParam f_params)
388390
{
389391
using RAJA::internal::thread_privatize;
390392
auto privatizer = thread_privatize(loop_body);
@@ -412,7 +414,7 @@ template<typename EXEC_POL,
412414
IterationMapping>::value &&
413415
(IterationGetter::block_size <= 0),
414416
size_t> RAJA_UNUSED_ARG(BlockSize) = 0>
415-
__global__ RAJA_JIT_COMPILE(3) void forallp_cuda_kernel(
417+
__global__ RAJA_JIT_COMPILE_ARGS(3) void forallp_cuda_kernel(
416418
const RAJA_CUDA_GRID_CONSTANT LOOP_BODY loop_body,
417419
const RAJA_CUDA_GRID_CONSTANT Iterator idx,
418420
const RAJA_CUDA_GRID_CONSTANT IndexType length,
@@ -447,11 +449,12 @@ template<
447449
IterationMapping>::value &&
448450
(IterationGetter::block_size > 0),
449451
size_t> BlockSize = IterationGetter::block_size>
450-
__launch_bounds__(BlockSize, BlocksPerSM) __global__ RAJA_JIT_COMPILE(3)
451-
void forallp_cuda_kernel(const RAJA_CUDA_GRID_CONSTANT LOOP_BODY loop_body,
452-
const RAJA_CUDA_GRID_CONSTANT Iterator idx,
453-
const RAJA_CUDA_GRID_CONSTANT IndexType length,
454-
ForallParam f_params)
452+
__launch_bounds__(BlockSize, BlocksPerSM) __global__
453+
RAJA_JIT_COMPILE_ARGS(3) void forallp_cuda_kernel(
454+
const RAJA_CUDA_GRID_CONSTANT LOOP_BODY loop_body,
455+
const RAJA_CUDA_GRID_CONSTANT Iterator idx,
456+
const RAJA_CUDA_GRID_CONSTANT IndexType length,
457+
ForallParam f_params)
455458
{
456459
using RAJA::internal::thread_privatize;
457460
auto privatizer = thread_privatize(loop_body);
@@ -482,7 +485,7 @@ template<
482485
IterationMapping>::value &&
483486
(IterationGetter::block_size <= 0),
484487
size_t> RAJA_UNUSED_ARG(BlockSize) = 0>
485-
__global__ RAJA_JIT_COMPILE(3) void forallp_cuda_kernel(
488+
__global__ RAJA_JIT_COMPILE_ARGS(3) void forallp_cuda_kernel(
486489
const RAJA_CUDA_GRID_CONSTANT LOOP_BODY loop_body,
487490
const RAJA_CUDA_GRID_CONSTANT Iterator idx,
488491
const RAJA_CUDA_GRID_CONSTANT IndexType length,

‎include/RAJA/policy/cuda/kernel/CudaKernel.hpp‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030

3131
#include "RAJA/util/macros.hpp"
3232
#include "RAJA/util/types.hpp"
33+
#include "RAJA/util/Jit.hpp"
3334

3435
#include "RAJA/pattern/kernel.hpp"
3536
#include "RAJA/pattern/kernel/For.hpp"
@@ -211,7 +212,8 @@ namespace internal
211212
* CUDA global function for launching CudaKernel policies
212213
*/
213214
template<typename Data, typename Exec>
214-
__global__ RAJA_JIT_COMPILE void CudaKernelLauncher(const RAJA_CUDA_GRID_CONSTANT Data data)
215+
__global__ RAJA_JIT_COMPILE void CudaKernelLauncher(
216+
const RAJA_CUDA_GRID_CONSTANT Data data)
215217
{
216218

217219
using data_t = camp::decay<Data>;
@@ -681,7 +683,7 @@ struct StatementExecutor<
681683
auto cuda_data = RAJA::cuda::make_launch_body(
682684
func, launch_dims.dims.blocks, launch_dims.dims.threads, shmem, res,
683685
data);
684-
686+
RAJA::register_lambda(func);
685687
//
686688
// Launch the kernel
687689
//

‎include/RAJA/policy/cuda/launch.hpp‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -24,14 +24,15 @@
2424
#include "RAJA/policy/cuda/MemUtils_CUDA.hpp"
2525
#include "RAJA/policy/cuda/raja_cudaerrchk.hpp"
2626
#include "RAJA/util/resource.hpp"
27+
#include "RAJA/util/Jit.hpp"
2728

2829
namespace RAJA
2930
{
3031

3132
template<typename BODY, typename ReduceParams>
32-
__global__ RAJA_JIT_COMPILE void launch_new_reduce_global_fcn(const RAJA_CUDA_GRID_CONSTANT BODY
33-
body_in,
34-
ReduceParams reduce_params)
33+
__global__ RAJA_JIT_COMPILE void launch_new_reduce_global_fcn(
34+
const RAJA_CUDA_GRID_CONSTANT BODY body_in,
35+
ReduceParams reduce_params)
3536
{
3637
LaunchContext ctx;
3738

@@ -95,7 +96,7 @@ struct LaunchExecute<
9596
if (gridSize.x > zero && gridSize.y > zero && gridSize.z > zero &&
9697
blockSize.x > zero && blockSize.y > zero && blockSize.z > zero)
9798
{
98-
99+
RAJA::register_lambda(body_in);
99100
RAJA_FT_BEGIN;
100101

101102
size_t shared_mem_size = launch_params.shared_mem_size;
@@ -206,7 +207,7 @@ struct LaunchExecute<
206207
if (gridSize.x > zero && gridSize.y > zero && gridSize.z > zero &&
207208
blockSize.x > zero && blockSize.y > zero && blockSize.z > zero)
208209
{
209-
210+
RAJA::register_lambda(body_in);
210211
RAJA_FT_BEGIN;
211212

212213
size_t shared_mem_size = launch_params.shared_mem_size;

‎include/RAJA/policy/hip/forall.hpp‎

Lines changed: 21 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535

3636
#include "RAJA/util/macros.hpp"
3737
#include "RAJA/util/types.hpp"
38+
#include "RAJA/util/Jit.hpp"
3839

3940
#include "RAJA/internal/fault_tolerance.hpp"
4041

@@ -373,11 +374,11 @@ template<typename EXEC_POL,
373374
IterationMapping>::value &&
374375
(IterationGetter::block_size > 0),
375376
size_t> BlockSize = IterationGetter::block_size>
376-
__launch_bounds__(BlockSize, 1) __global__ RAJA_JIT_COMPILE(3)
377-
void forallp_hip_kernel(const LOOP_BODY loop_body,
378-
const Iterator idx,
379-
const IndexType length,
380-
ForallParam f_params)
377+
__launch_bounds__(BlockSize, 1) __global__
378+
RAJA_JIT_COMPILE_ARGS(3) void forallp_hip_kernel(const LOOP_BODY loop_body,
379+
const Iterator idx,
380+
const IndexType length,
381+
ForallParam f_params)
381382
{
382383
using RAJA::internal::thread_privatize;
383384
auto privatizer = thread_privatize(loop_body);
@@ -403,10 +404,11 @@ template<typename EXEC_POL,
403404
IterationMapping>::value &&
404405
(IterationGetter::block_size <= 0),
405406
size_t> RAJA_UNUSED_ARG(BlockSize) = 0>
406-
__global__ RAJA_JIT_COMPILE(3) void forallp_hip_kernel(const LOOP_BODY loop_body,
407-
const Iterator idx,
408-
const IndexType length,
409-
ForallParam f_params)
407+
__global__ RAJA_JIT_COMPILE_ARGS(3) void forallp_hip_kernel(
408+
const LOOP_BODY loop_body,
409+
const Iterator idx,
410+
const IndexType length,
411+
ForallParam f_params)
410412
{
411413
using RAJA::internal::thread_privatize;
412414
auto privatizer = thread_privatize(loop_body);
@@ -434,11 +436,11 @@ template<
434436
IterationMapping>::value &&
435437
(IterationGetter::block_size > 0),
436438
size_t> BlockSize = IterationGetter::block_size>
437-
__launch_bounds__(BlockSize, 1) __global__ RAJA_JIT_COMPILE(3)
438-
void forallp_hip_kernel(const LOOP_BODY loop_body,
439-
const Iterator idx,
440-
const IndexType length,
441-
ForallParam f_params)
439+
__launch_bounds__(BlockSize, 1) __global__
440+
RAJA_JIT_COMPILE_ARGS(3) void forallp_hip_kernel(const LOOP_BODY loop_body,
441+
const Iterator idx,
442+
const IndexType length,
443+
ForallParam f_params)
442444
{
443445
using RAJA::internal::thread_privatize;
444446
auto privatizer = thread_privatize(loop_body);
@@ -466,10 +468,11 @@ template<
466468
IterationMapping>::value &&
467469
(IterationGetter::block_size <= 0),
468470
size_t> RAJA_UNUSED_ARG(BlockSize) = 0>
469-
__global__ RAJA_JIT_COMPILE(3) void forallp_hip_kernel(const LOOP_BODY loop_body,
470-
const Iterator idx,
471-
const IndexType length,
472-
ForallParam f_params)
471+
__global__ RAJA_JIT_COMPILE_ARGS(3) void forallp_hip_kernel(
472+
const LOOP_BODY loop_body,
473+
const Iterator idx,
474+
const IndexType length,
475+
ForallParam f_params)
473476
{
474477
using RAJA::internal::thread_privatize;
475478
auto privatizer = thread_privatize(loop_body);

‎include/RAJA/policy/hip/kernel/HipKernel.hpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030

3131
#include "RAJA/util/macros.hpp"
3232
#include "RAJA/util/types.hpp"
33+
#include "RAJA/util/Jit.hpp"
3334

3435
#include "RAJA/pattern/kernel.hpp"
3536
#include "RAJA/pattern/kernel/For.hpp"
@@ -639,7 +640,7 @@ struct StatementExecutor<
639640
auto hip_data = RAJA::hip::make_launch_body(
640641
func, launch_dims.dims.blocks, launch_dims.dims.threads, shmem, res,
641642
data);
642-
643+
RAJA::register_lambda(func);
643644
//
644645
// Launch the kernel
645646
//

‎include/RAJA/policy/hip/launch.hpp‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,13 +24,15 @@
2424
#include "RAJA/policy/hip/MemUtils_HIP.hpp"
2525
#include "RAJA/policy/hip/raja_hiperrchk.hpp"
2626
#include "RAJA/util/resource.hpp"
27+
#include "RAJA/util/Jit.hpp"
2728

2829
namespace RAJA
2930
{
3031

3132
template<typename BODY, typename ReduceParams>
32-
__global__ RAJA_JIT_COMPILE void launch_new_reduce_global_fcn(const BODY body_in,
33-
ReduceParams reduce_params)
33+
__global__ RAJA_JIT_COMPILE void launch_new_reduce_global_fcn(
34+
const BODY body_in,
35+
ReduceParams reduce_params)
3436
{
3537
LaunchContext ctx;
3638

@@ -260,7 +262,7 @@ struct LoopExecute<
260262
BODY const& body)
261263
{
262264
const diff_t i = IndexMapper::template index<diff_t>();
263-
RAJA::register_lambda(body);
265+
264266
body(*(segment.begin() + i));
265267
}
266268
};

‎include/RAJA/util/Jit.hpp‎

Lines changed: 25 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3,17 +3,32 @@
33
#include "proteus/JitInterface.hpp"
44
#endif
55

6-
#ifndef RAJA_plugins_HPP
7-
#define RAJA_plugins_HPP
6+
#ifndef RAJA_jit_HPP
7+
#define RAJA_jit_HPP
88

9-
namespace RAJA {
10-
template<typename Lambda>
11-
inline auto register_lambda(Lambda&& lambda) {
12-
#if defined RAJA_ENABLE_JIT
13-
return proteus::register_lambda(std::forward<Lambda>(lambda));
14-
#endif
15-
return std::forward<Lambda>(lambda);
16-
}
9+
namespace RAJA
10+
{
11+
template<typename Lambda>
12+
inline auto register_lambda(Lambda&& lambda)
13+
{
14+
#if defined RAJA_ENABLE_JIT
15+
return proteus::register_lambda(std::forward<Lambda>(lambda));
16+
#else
17+
return std::forward<Lambda>(lambda);
18+
#endif
19+
}
20+
21+
template<typename T>
22+
inline auto jit_variable(T arg)
23+
{
24+
#if defined RAJA_ENABLE_JIT
25+
return proteus::jit_variable(std::forward<T>(arg));
26+
#else
27+
return std::forward<T>(arg);
28+
#endif
1729
}
1830

31+
32+
} // namespace RAJA
33+
1934
#endif

‎include/RAJA/util/macros.hpp‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -225,7 +225,11 @@ inline void RAJA_ABORT_OR_THROW(const char* str)
225225
#endif
226226

227227
#if defined RAJA_ENABLE_JIT
228-
#define RAJA_JIT_COMPILE(...) __attribute__((annotate("jit", __VA_ARGS__)))
228+
#define RAJA_JIT_COMPILE_ARGS(...) __attribute__((annotate("jit", __VA_ARGS__)))
229+
#define RAJA_JIT_COMPILE __attribute__((annotate("jit")))
230+
#else
231+
#define RAJA_JIT_COMPILE_ARGS(...)
232+
#define RAJA_JIT_COMPILE
229233
#endif
230234

231235
#endif /* RAJA_INTERNAL_MACROS_HPP */

0 commit comments

Comments
 (0)