Skip to content

Add LTimes tuning to represent Kripke LTimes #684

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
michaelmckinsey1 wants to merge 10 commits into
base: develop
Choose a base branch
from
Open

Add LTimes tuning to represent Kripke LTimes #684

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
7bdc07f
ell layout d-stride. Change default block size to 25 to match moments…
michaelmckinsey1 Jun 9, 2026
7a794c7
Add link
michaelmckinsey1 Jun 10, 2026
1147a77
Revert
michaelmckinsey1 Jun 10, 2026
d2f79aa
fix m stride-1
michaelmckinsey1 Jun 18, 2026
6c6c9de
Match sycl
michaelmckinsey1 Jun 18, 2026
351f2d8
Apply same changes to noview
michaelmckinsey1 Jun 18, 2026
3edfc3b
Merge branch 'develop' into ltimes-block25
michaelmckinsey1 Jun 18, 2026
2037b48
Separate tunings
michaelmckinsey1 Jul 9, 2026
dfb8ab5
Rename
michaelmckinsey1 Jul 9, 2026
02f78bc
Merge branch 'develop' into ltimes-block25
michaelmckinsey1 Jul 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
278 changes: 229 additions & 49 deletions src/apps/LTIMES-Cuda.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,27 +27,60 @@ using namespace ltimes_idx;
//
// Define thread block shape for CUDA execution
//
#define m_block_sz (32)
#define g_block_sz (integer::greater_of_squarest_factor_pair(block_size/m_block_sz))
#define z_block_sz (integer::lesser_of_squarest_factor_pair(block_size/m_block_sz))
#define block_moments_m_block_sz (block_size)
#define block_moments_g_block_sz (1)
#define block_moments_z_block_sz (1)

#define LTIMES_THREADS_PER_BLOCK_TEMPLATE_PARAMS_CUDA \
m_block_sz, g_block_sz, z_block_sz
#define LTIMES_BLOCK_MOMENTS_THREADS_PER_BLOCK_TEMPLATE_PARAMS_CUDA \
block_moments_m_block_sz, block_moments_g_block_sz, block_moments_z_block_sz

#define LTIMES_THREADS_PER_BLOCK_CUDA \
dim3 nthreads_per_block(LTIMES_THREADS_PER_BLOCK_TEMPLATE_PARAMS_CUDA); \
static_assert(m_block_sz*g_block_sz*z_block_sz == block_size, "Invalid block_size");
#define LTIMES_BLOCK_MOMENTS_THREADS_PER_BLOCK_CUDA \
dim3 nthreads_per_block(LTIMES_BLOCK_MOMENTS_THREADS_PER_BLOCK_TEMPLATE_PARAMS_CUDA); \
static_assert(block_moments_m_block_sz*block_moments_g_block_sz*block_moments_z_block_sz == block_size, "Invalid block_size");

#define LTIMES_NBLOCKS_CUDA \
dim3 nblocks(static_cast<size_t>(RAJA_DIVIDE_CEILING_INT(*num_m, m_block_sz)), \
static_cast<size_t>(RAJA_DIVIDE_CEILING_INT(*num_g, g_block_sz)), \
static_cast<size_t>(RAJA_DIVIDE_CEILING_INT(*num_z, z_block_sz)));
#define LTIMES_BLOCK_MOMENTS_NBLOCKS_CUDA \
dim3 nblocks(static_cast<size_t>(*num_z), \
static_cast<size_t>(*num_g), \
1);

#define factorized_m_block_sz (32)
#define factorized_g_block_sz (integer::greater_of_squarest_factor_pair(block_size/factorized_m_block_sz))
#define factorized_z_block_sz (integer::lesser_of_squarest_factor_pair(block_size/factorized_m_block_sz))

#define LTIMES_FACTORIZED_THREADS_PER_BLOCK_TEMPLATE_PARAMS_CUDA \
factorized_m_block_sz, factorized_g_block_sz, factorized_z_block_sz

#define LTIMES_FACTORIZED_THREADS_PER_BLOCK_CUDA \
dim3 nthreads_per_block(LTIMES_FACTORIZED_THREADS_PER_BLOCK_TEMPLATE_PARAMS_CUDA); \
static_assert(factorized_m_block_sz*factorized_g_block_sz*factorized_z_block_sz == block_size, "Invalid block_size");

#define LTIMES_FACTORIZED_NBLOCKS_CUDA \
dim3 nblocks(static_cast<size_t>(RAJA_DIVIDE_CEILING_INT(*num_m, factorized_m_block_sz)), \
static_cast<size_t>(RAJA_DIVIDE_CEILING_INT(*num_g, factorized_g_block_sz)), \
static_cast<size_t>(RAJA_DIVIDE_CEILING_INT(*num_z, factorized_z_block_sz)));


template < size_t block_size >
__launch_bounds__(block_size)
__global__ void ltimes_block_moments(PHI_VIEW phi, ELL_VIEW ell, PSI_VIEW psi,
ID num_d, IM num_m, IG num_g, IZ num_z)
{
IG g(blockIdx.y);
IZ z(blockIdx.x);

if (g < num_g && z < num_z) {
for (IM m(threadIdx.x); m < num_m; m += block_size) {
for (ID d(0); d < num_d; ++d ) {
LTIMES_BODY;
}
}
}
}

template < size_t m_block_size, size_t g_block_size, size_t z_block_size >
__launch_bounds__(m_block_size*g_block_size*z_block_size)
__global__ void ltimes(PHI_VIEW phi, ELL_VIEW ell, PSI_VIEW psi,
ID num_d, IM num_m, IG num_g, IZ num_z)
__global__ void ltimes_factorized(PHI_VIEW phi, ELL_VIEW ell, PSI_VIEW psi,
ID num_d, IM num_m, IG num_g, IZ num_z)
{
IM m(blockIdx.x * m_block_size + threadIdx.x);
IG g(blockIdx.y * g_block_size + threadIdx.y);
Expand All @@ -60,10 +93,25 @@ __global__ void ltimes(PHI_VIEW phi, ELL_VIEW ell, PSI_VIEW psi,
}
}

template < size_t block_size, typename Lambda >
__launch_bounds__(block_size)
__global__ void ltimes_lam_block_moments(IM num_m, IG num_g, IZ num_z,
Lambda body)
{
IG g(blockIdx.y);
IZ z(blockIdx.x);

if (g < num_g && z < num_z) {
for (IM m(threadIdx.x); m < num_m; m += block_size) {
body(z, g, m);
}
}
}

template < size_t m_block_size, size_t g_block_size, size_t z_block_size, typename Lambda >
__launch_bounds__(m_block_size*g_block_size*z_block_size)
__global__ void ltimes_lam(IM num_m, IG num_g, IZ num_z,
Lambda body)
__global__ void ltimes_lam_factorized(IM num_m, IG num_g, IZ num_z,
Lambda body)
{
IM m(blockIdx.x * m_block_size + threadIdx.x);
IG g(blockIdx.y * g_block_size + threadIdx.y);
Expand Down Expand Up @@ -92,16 +140,29 @@ void LTIMES::runCudaVariantImpl(VariantID vid)
// Loop counter increment uses macro to quiet C++20 compiler warning
for (RepIndex_type irep = 0; irep < run_reps; RP_REPCOUNTINC(irep)) {

LTIMES_THREADS_PER_BLOCK_CUDA;
LTIMES_NBLOCKS_CUDA;
constexpr size_t shmem = 0;

RPlaunchCudaKernel(
(ltimes<LTIMES_THREADS_PER_BLOCK_TEMPLATE_PARAMS_CUDA>),
nblocks, nthreads_per_block,
shmem, res.get_stream(),
phi, ell, psi,
num_d, num_m, num_g, num_z );
if constexpr (tune_idx == 1) {
LTIMES_FACTORIZED_THREADS_PER_BLOCK_CUDA;
LTIMES_FACTORIZED_NBLOCKS_CUDA;

RPlaunchCudaKernel(
(ltimes_factorized<LTIMES_FACTORIZED_THREADS_PER_BLOCK_TEMPLATE_PARAMS_CUDA>),
nblocks, nthreads_per_block,
shmem, res.get_stream(),
phi, ell, psi,
num_d, num_m, num_g, num_z );
} else {
LTIMES_BLOCK_MOMENTS_THREADS_PER_BLOCK_CUDA;
LTIMES_BLOCK_MOMENTS_NBLOCKS_CUDA;

RPlaunchCudaKernel(
(ltimes_block_moments<block_size>),
nblocks, nthreads_per_block,
shmem, res.get_stream(),
phi, ell, psi,
num_d, num_m, num_g, num_z );
}

}
stopTimer();
Expand All @@ -118,17 +179,30 @@ void LTIMES::runCudaVariantImpl(VariantID vid)
}
};

LTIMES_THREADS_PER_BLOCK_CUDA;
LTIMES_NBLOCKS_CUDA;
constexpr size_t shmem = 0;

RPlaunchCudaKernel(
(ltimes_lam<LTIMES_THREADS_PER_BLOCK_TEMPLATE_PARAMS_CUDA,
decltype(ltimes_lambda)>),
nblocks, nthreads_per_block,
shmem, res.get_stream(),
num_m, num_g, num_z,
ltimes_lambda );
if constexpr (tune_idx == 1) {
LTIMES_FACTORIZED_THREADS_PER_BLOCK_CUDA;
LTIMES_FACTORIZED_NBLOCKS_CUDA;

RPlaunchCudaKernel(
(ltimes_lam_factorized<LTIMES_FACTORIZED_THREADS_PER_BLOCK_TEMPLATE_PARAMS_CUDA,
decltype(ltimes_lambda)>),
nblocks, nthreads_per_block,
shmem, res.get_stream(),
num_m, num_g, num_z,
ltimes_lambda );
} else {
LTIMES_BLOCK_MOMENTS_THREADS_PER_BLOCK_CUDA;
LTIMES_BLOCK_MOMENTS_NBLOCKS_CUDA;

RPlaunchCudaKernel(
(ltimes_lam_block_moments<block_size, decltype(ltimes_lambda)>),
nblocks, nthreads_per_block,
shmem, res.get_stream(),
num_m, num_g, num_z,
ltimes_lambda );
}

}
stopTimer();
Expand All @@ -139,10 +213,10 @@ void LTIMES::runCudaVariantImpl(VariantID vid)

using EXEC_POL =
RAJA::KernelPolicy<
RAJA::statement::CudaKernelFixedAsync<m_block_sz*g_block_sz*z_block_sz,
RAJA::statement::For<1, RAJA::cuda_global_size_z_direct<z_block_sz>, //z
RAJA::statement::For<2, RAJA::cuda_global_size_y_direct<g_block_sz>, //g
RAJA::statement::For<3, RAJA::cuda_global_size_x_direct<m_block_sz>, //m
RAJA::statement::CudaKernelAsync<

@rchen20 rchen20 Jun 19, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note that I've reverted LTimes to launch synchronously in Kripke, for correctness. This is fine though because the direction loop is inner-most, which should avoid race conditions.

@michaelmckinsey1 michaelmckinsey1 Jun 24, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It has been async in RAJAPerf, I just changed it from CudaKernelFixedAsync to CudaKernelAsync, but for completeness I can make it CudaKernel. I don't this would matter for performance in RAJAPerf.

RAJA::statement::For<1, RAJA::cuda_block_x_loop, // z
RAJA::statement::For<2, RAJA::cuda_block_y_loop, // g
RAJA::statement::For<3, RAJA::cuda_thread_x_loop, // m

@MrBurmark MrBurmark Jun 19, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I assume this is a non-size loop policy because it is in ltimes. Here we know the block size at compile time, is that also true in kripke?

@rchen20 rchen20 Jun 19, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In Kripke, the block sizes are determined by parameters passed in at runtime. For example, like in this version of LTimes, it will be blocked on zones and groups. The exact parameters which will be blocked are not always the same for each Kripke run. For instance, if we use the DZG layout at runtime, then the loops will be blocked with directions and zones (while groups are threaded).

RAJA::statement::For<0, RAJA::seq_exec, //d
RAJA::statement::Lambda<0>
>
Expand Down Expand Up @@ -172,31 +246,114 @@ void LTIMES::runCudaVariantImpl(VariantID vid)

} else if constexpr (tune_idx == 1) {

using EXEC_POL =
RAJA::KernelPolicy<
RAJA::statement::CudaKernelFixedAsync<factorized_m_block_sz*factorized_g_block_sz*factorized_z_block_sz,
RAJA::statement::For<1, RAJA::cuda_global_size_z_direct<factorized_z_block_sz>, // z
RAJA::statement::For<2, RAJA::cuda_global_size_y_direct<factorized_g_block_sz>, // g
RAJA::statement::For<3, RAJA::cuda_global_size_x_direct<factorized_m_block_sz>, // m
RAJA::statement::For<0, RAJA::seq_exec, // d
RAJA::statement::Lambda<0>
>
>
>
>
>
>;

startTimer();
// Loop counter increment uses macro to quiet C++20 compiler warning
for (RepIndex_type irep = 0; irep < run_reps; RP_REPCOUNTINC(irep)) {

RAJA::kernel_resource<EXEC_POL>(
RAJA::make_tuple(IDRange(0, *num_d),
IZRange(0, *num_z),
IGRange(0, *num_g),
IMRange(0, *num_m)),
res,
[=] __device__ (ID d, IZ z, IG g, IM m) {
LTIMES_BODY;
}
);

}
stopTimer();

} else if constexpr (tune_idx == 2) {

constexpr bool async = true;

using launch_policy =
RAJA::LaunchPolicy<RAJA::cuda_launch_t<async, block_size>>;

using z_policy = RAJA::LoopPolicy<RAJA::cuda_block_x_loop>;

using g_policy = RAJA::LoopPolicy<RAJA::cuda_block_y_loop>;

using m_policy = RAJA::LoopPolicy<RAJA::cuda_thread_x_loop>;

using d_policy = RAJA::LoopPolicy<RAJA::seq_exec>;

startTimer();
// Loop counter increment uses macro to quiet C++20 compiler warning
for (RepIndex_type irep = 0; irep < run_reps; RP_REPCOUNTINC(irep)) {

RAJA::launch<launch_policy>( res,
RAJA::LaunchParams(RAJA::Teams(*num_z, *num_g, 1),
RAJA::Threads(block_size, 1, 1)),
[=] RAJA_HOST_DEVICE(RAJA::LaunchContext ctx) {

@artv3 artv3 Jul 10, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It might be worth looking into the CachePolicy for the launch_context:

https://github.com/llnl/RAJAPerf/blob/develop/src/apps/MASS3DEA-Cuda.cpp#L206-L209

We got some speedups here, this might help the nested loops

RAJA::loop<z_policy>(ctx, IZRange(0, *num_z),
[&](IZ z) {
RAJA::loop<g_policy>(ctx, IGRange(0, *num_g),
[&](IG g) {
RAJA::loop<m_policy>(ctx, IMRange(0, *num_m),
[&](IM m) {
RAJA::loop<d_policy>(ctx, IDRange(0, *num_d),
[&](ID d) {
LTIMES_BODY
}
); // RAJA::loop<d_policy>
}
); // RAJA::loop<m_policy>
}
); // RAJA::loop<g_policy>
}
); // RAJA::loop<z_policy>

} // outer lambda (ctx)
); // RAJA::launch

} // loop over kernel reps
stopTimer();
} else if constexpr (tune_idx == 3) {

constexpr bool async = true;

using launch_policy = RAJA::LaunchPolicy<RAJA::cuda_launch_t<async, m_block_sz*g_block_sz*z_block_sz>>;
using launch_policy =
RAJA::LaunchPolicy<RAJA::cuda_launch_t<async, factorized_m_block_sz*factorized_g_block_sz*factorized_z_block_sz>>;

using z_policy = RAJA::LoopPolicy<RAJA::cuda_global_size_z_loop<z_block_sz>>;
using z_policy = RAJA::LoopPolicy<RAJA::cuda_global_size_z_loop<factorized_z_block_sz>>;

using g_policy = RAJA::LoopPolicy<RAJA::cuda_global_size_y_loop<g_block_sz>>;
using g_policy = RAJA::LoopPolicy<RAJA::cuda_global_size_y_loop<factorized_g_block_sz>>;

using m_policy = RAJA::LoopPolicy<RAJA::cuda_global_size_x_loop<m_block_sz>>;
using m_policy = RAJA::LoopPolicy<RAJA::cuda_global_size_x_loop<factorized_m_block_sz>>;

using d_policy = RAJA::LoopPolicy<RAJA::seq_exec>;

const size_t z_grid_sz = RAJA_DIVIDE_CEILING_INT(*num_z, z_block_sz);
const size_t z_grid_sz = RAJA_DIVIDE_CEILING_INT(*num_z, factorized_z_block_sz);

const size_t g_grid_sz = RAJA_DIVIDE_CEILING_INT(*num_g, g_block_sz);
const size_t g_grid_sz = RAJA_DIVIDE_CEILING_INT(*num_g, factorized_g_block_sz);

const size_t m_grid_sz = RAJA_DIVIDE_CEILING_INT(*num_m, m_block_sz);
const size_t m_grid_sz = RAJA_DIVIDE_CEILING_INT(*num_m, factorized_m_block_sz);

startTimer();
// Loop counter increment uses macro to quiet C++20 compiler warning
for (RepIndex_type irep = 0; irep < run_reps; RP_REPCOUNTINC(irep)) {

RAJA::launch<launch_policy>( res,
RAJA::LaunchParams(RAJA::Teams(m_grid_sz, g_grid_sz, z_grid_sz),
RAJA::Threads(m_block_sz, g_block_sz, z_block_sz)),
RAJA::Threads(factorized_m_block_sz, factorized_g_block_sz, factorized_z_block_sz)),
[=] RAJA_HOST_DEVICE(RAJA::LaunchContext ctx) {

RAJA::loop<z_policy>(ctx, IZRange(0, *num_z),
Expand Down Expand Up @@ -242,12 +399,35 @@ void LTIMES::defineCudaVariantTunings()

if (vid == RAJA_CUDA) {
addVariantTuning<&LTIMES::runCudaVariantImpl<block_size, 0>>(
vid, "kernel_"+std::to_string(block_size));
addVariantTuning<&LTIMES::runCudaVariantImpl<block_size, 1>>(
vid, "launch_"+std::to_string(block_size));
vid, "kernel_m");
addVariantTuning<&LTIMES::runCudaVariantImpl<block_size, 2>>(
vid, "launch_m");
} else {
addVariantTuning<&LTIMES::runCudaVariantImpl<block_size, 0>>(
vid, "block_"+std::to_string(block_size));
vid, "block_m");
}

}

});

}

for (VariantID vid : {Base_CUDA, Lambda_CUDA, RAJA_CUDA}) {

seq_for(factorized_gpu_block_sizes_type{}, [&](auto block_size) {

if (run_params.numValidGPUBlockSize() == 0u ||
run_params.validGPUBlockSize(block_size)) {

if (vid == RAJA_CUDA) {
addVariantTuning<&LTIMES::runCudaVariantImpl<block_size, 1>>(
vid, "kernel_m32_gBigFact_zLowFact_"+std::to_string(block_size));
addVariantTuning<&LTIMES::runCudaVariantImpl<block_size, 3>>(
vid, "launch_m32_gBigFact_zLowFact_"+std::to_string(block_size));
} else {
addVariantTuning<&LTIMES::runCudaVariantImpl<block_size, 1>>(
vid, "block_m32_gBigFact_zLowFact_"+std::to_string(block_size));
}

}
Expand Down
Loading
Loading