Problem Statement

After a naive initial port of femto's kernels to CUDA, we had an implementation that looked something like:

xxxxxxxxxx
// start off by setting up our shared memory buffers
uint32_t shmem_offset = 0;
extern __shared__ char shmem[];

uint32_t local_elem = threadIdx.z;
uint32_t elem_per_block = blockDim.z;

nd::view< Connection > shr_connectivity(
  (Connection *)(shmem + shmem_offset), 
  {connectivity.shape[1]}
);
shmem_offset += round_up_to_multiple_of_128(shr_connectivity.size() * sizeof(Connection));

uint32_t nodes_per_elem = el.num_nodes();
nd::view< uint32_t > shr_node_ids(
  (uint32_t *)(shmem + shmem_offset), 
  {nodes_per_elem}
);
shmem_offset += round_up_to_multiple_of_128(shr_node_ids.size() * sizeof(uint32_t));

nd::view< grad_type > shr_du_dxi_q(
  (grad_type *)(shmem + shmem_offset), 
  {qpts_per_elem}
);
shmem_offset += round_up_to_multiple_of_128(shr_du_dxi_q.size() * sizeof(grad_type));

nd::view< double > shr_u_e(
  (double *)(shmem + shmem_offset), 
  {nodes_per_elem}
);
shmem_offset += round_up_to_multiple_of_128(shr_u_e.size() * sizeof(double));

nd::view< double, n > shr_shape_fn_grads(
  (double *)(shmem + shmem_offset), 
  shape_functions.shape
);
shmem_offset += round_up_to_multiple_of_128(shr_shape_fn_grads.size() * sizeof(double));

nd::view< double > shr_scratch(
  (double *)(shmem + shmem_offset), 
  {scratch_size}
);

int elem_tid = threadIdx.x + blockDim.x * threadIdx.y;
int elem_stride = blockDim.x * blockDim.y;

int block_tid = threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z);
int block_stride = blockDim.x * blockDim.y * blockDim.z;

// 0. load shape function evaluations into shared memory
for (int i = block_tid; i < shr_shape_fn_grads.size(); i += block_stride) {
  shr_shape_fn_grads[i] = shape_functions[i];
}

// 1. then figure out which element we're integrating and load its connectivity info
uint32_t e = blockIdx.x;
uint32_t elem_id = elements[e];

for (int i = elem_tid; i < shr_connectivity.size(); i += elem_stride) {
  shr_connectivity[i] = connectivity(elem_id, i);
}
__syncthreads();

// 2. from the connectivity, figure out which nodes belong to this element 
if (elem_tid == 0) {
  el.indices(u_offsets, shr_connectivity.data(), shr_node_ids.data());
}
__syncthreads();

int num_components = u.shape[1];
for (int c = 0; c < num_components; c++) {

  // 3. load the nodal values for this element
  for (int i = elem_tid; i < nodes_per_elem; i += elem_stride) {
    shr_u_e(i) = u(shr_node_ids(i), c);
  }
  __syncthreads();

  // 4. interpolate the quadrature point values for this element
  el.cuda_gradient(shr_du_dxi_q, shr_u_e, shr_shape_fn_grads, shr_scratch.data());

  // 5. write those quadrature values out to global memory
  for (int q = elem_tid; q < qpts_per_elem; q += elem_stride) {
    du_dxi_q(qpts_per_elem * e + q, c) = shr_du_dxi_q[q];
  }

}

This implementation already shows significant speedups for many of the different combinations of element geometry and polynomial order, but it has an issue: it is written so that each thread block processes a single element. This may work okay when that element is a high order tetrahedron and hexahedron, but for low order triangles and quadrilaterals this amounts to launching a kernel with a block size as low as 3.

This may not be a big deal when running on a CPU, but it is problematic for GPUs. On a GPU, calculations are carried out by "warps" (groups of 32 parallel threads), which consequently can operate on up to 32 pieces of data a time. However, if we launch a kernel with a blocksize of only 3, then most of the threads in a given warp will be inactive, which means we only attain a fraction of the performance that the hardware is capable of.

So, how can we better utilize the GPU hardware for elements that only involve 3-4 values per element?

Well, fundamentally there is not enough work in a single 3-node triangle to keep a warp busy, but we can modify the kernel to process multiple elements at a time, rather than just one. That way, we can ensure that each thread block has enough work to keep the hardware busy.

Here's a sketch of how to do that:

1. Modify the kernel launch parameters

   // before: one element / block
   int block_size = nodes_per_element;
   int grid_size = num_elements;
   my_kernel<<<grid_size, block_size>>>(...);
   
   // after: multiple element / block
   dim3 block_size = {nodes_per_element, 1, elem_per_block};
   int grid_size = (num_elements + elem_per_block - 1) / elem_per_block;
   my_kernel<<<grid_size, block_size>>>(...);
   

2. Inspect threadIdx.z in the kernel to figure out which element within the block to work on

3. Load data and perform calculations on the element corresponding to threadIdx.z

In order to make step 3 work, we also have to increase the amount of shared memory we allocate, to make sure each element has space for its own intermediates. Some of the data (like shape function evaluations) are common to all of the elements in the thread block, so those don't need to be duplicated. Here's some code for the updated body of the kernel, taking into account processing multiple elements per block:

x
// start off by setting up our shared memory buffers
uint32_t shmem_offset = 0;
extern __shared__ char shmem[];

uint32_t elem_per_block = blockDim.z;

nd::view< Connection, 2 > shr_connectivity(
  (Connection *)(shmem + shmem_offset), 
  {elem_per_block, connectivity.shape[1]}
);
shmem_offset += round_up_to_multiple_of_128(shr_connectivity.size() * sizeof(Connection));

uint32_t nodes_per_elem = el.num_nodes();
nd::view< uint32_t, 2 > shr_node_ids(
  (uint32_t *)(shmem + shmem_offset), 
  {elem_per_block, nodes_per_elem}
);
shmem_offset += round_up_to_multiple_of_128(shr_node_ids.size() * sizeof(uint32_t));

nd::view< grad_type, 2 > shr_du_dxi_q(
  (grad_type *)(shmem + shmem_offset), 
  {elem_per_block, qpts_per_elem}
);
shmem_offset += round_up_to_multiple_of_128(shr_du_dxi_q.size() * sizeof(grad_type));

nd::view< double, 2 > shr_u_e(
  (double *)(shmem + shmem_offset), 
  {elem_per_block, nodes_per_elem}
);
shmem_offset += round_up_to_multiple_of_128(shr_u_e.size() * sizeof(double));

nd::view< double, n > shr_shape_fn_grads(
  (double *)(shmem + shmem_offset), 
  shape_functions.shape
);
shmem_offset += round_up_to_multiple_of_128(shr_shape_fn_grads.size() * sizeof(double));

nd::view< double, 2 > shr_scratch(
  (double *)(shmem + shmem_offset), 
  {elem_per_block, scratch_size}
);

int local_tid = threadIdx.x + blockDim.x * threadIdx.y;
int local_stride = blockDim.x * blockDim.y;

int block_tid = threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z);
int block_stride = blockDim.x * blockDim.y * blockDim.z;

// 0. load shape function evaluations into shared memory
for (int i = block_tid; i < shr_shape_fn_grads.size(); i += block_stride) {
  shr_shape_fn_grads[i] = shape_functions[i];
}

// 1. then figure out which element we're integrating and load its connectivity info
uint32_t e = blockIdx.x * blockDim.z + threadIdx.z;
if (e < elements.shape[0]) {
  uint32_t elem_id = elements[e];
  uint32_t local_elem_id = threadIdx.z;

  for (int i = local_tid; i < shr_connectivity.shape[1]; i += local_stride) {
    shr_connectivity(local_elem_id, i) = connectivity(elem_id, i);
  }
  __syncthreads();

  // 2. from the connectivity, figure out which nodes belong to this element 
  if (local_tid == 0) {
    el.indices(u_offsets, &shr_connectivity(local_elem_id, 0), &shr_node_ids(local_elem_id, 0));
  }
  __syncthreads();

  int num_components = u.shape[1];
  for (int c = 0; c < num_components; c++) {

    // 3. load the nodal values for this element
    for (int i = local_tid; i < nodes_per_elem; i += local_stride) {
      shr_u_e(local_elem_id, i) = u(shr_node_ids(local_elem_id, i), c);
    }
    __syncthreads();

    // 4. interpolate the quadrature point values for this element
    el.cuda_gradient(
      shr_du_dxi_q(local_elem_id), 
      shr_u_e(local_elem_id), 
      shr_shape_fn_grads, 
      &shr_scratch(local_elem_id, 0)
    );

    // 5. write those quadrature values out to global memory
    for (int q = local_tid; q < qpts_per_elem; q += local_stride) {
      du_dxi_q(qpts_per_elem * e + q, c) = shr_du_dxi_q(local_elem_id, q);
    }

  }

}

This modification doesn't require significant changes. The main difference is that the shared memory arrays now have one extra dimension that each thread indexes by its local element id.

Great, now the kernel definition is able to handle multiple elements per block, but how many elements should each thread block process?

Intuitively, we should pick enough elements per block such that the warps in that thread block are utilized efficiently. This is true when the blocksize nodes_per_elem * elems_per_block is a multiple of 32, but that isn't a necessary condition. Also, recall that processing more elements per block, requires more shared memory. So, after a point, increasing the number of elements / block ends up using enough shared memory to hurt performance through decreased occupancy.

From those criteria (picking the smallest number that maximizes expected active threads / warp), we could guess some parameter values for each kind of element, but instead I just set up a performance benchmark and exhaustively tried out different numbers of elements/block (up to a maximum of 32 for the 2D elements and up to 16 for the 3D elements).

Tuning Data

The follow sections tabulate data from a performance benchmark involving an unstructured meshes ~1,000,000 elements of the specified geometry.

The tests were performed with as many quadrature points as nodes in the respective element, and the listed values are the timings (in microseconds) as measured on a NVIDIA GV100 GPU.

The labels "scalar" and "vector" indicate which timings are from interpolation of scalar-valued and vector-valued (with 2 components for tri/quad and 3 components for tet/hex) fields, respectively.

Triangles

Quadrilaterals

Tetrahedra

Hexahedra

Note: the entries marked with * require more shared memory than the GV100 has available

Summary

The naive CUDA port of the CPU code was a good start, but that implementation didn't perform very well for the low-order elements (which are the most commonly-used case). Luckily, it didn't make much effort to modify the kernel to accept a parameter that controls how many elements to process in a single threadblock. However, introducing a new parameter meant that we also needed to provide appropriate values for that parameter in each situation. By exhaustively sweeping through different number of elements processed per thread block, we were able to select values that significantly improved performance (7.5x for linear triangles, 5x for linear quads, ...).