xref: /libCEED/include/ceed/jit-source/cuda/cuda-gen-templates.h (revision 0183ed61035d97ff853cf8c8e722c0fda76e54df)

// Copyright (c) 2017-2025, Lawrence Livermore National Security, LLC and other CEED contributors.
// All Rights Reserved. See the top-level LICENSE and NOTICE files for details.
//
// SPDX-License-Identifier: BSD-2-Clause
//
// This file is part of CEED:  http://github.com/ceed

/// @file
/// Internal header for CUDA backend macro and type definitions for JiT source
#include <ceed/types.h>

//------------------------------------------------------------------------------
// Load matrices for basis actions
//------------------------------------------------------------------------------
template <int P, int Q>
inline __device__ void LoadMatrix(SharedData_Cuda &data, const CeedScalar *__restrict__ d_B, CeedScalar *B) {
  for (CeedInt i = data.t_id; i < P * Q; i += blockDim.x * blockDim.y * blockDim.z) B[i] = d_B[i];
}

//------------------------------------------------------------------------------
// AtPoints
//------------------------------------------------------------------------------

//------------------------------------------------------------------------------
// L-vector -> single point
//------------------------------------------------------------------------------
template <int NUM_COMP, int COMP_STRIDE, int NUM_PTS>
inline __device__ void ReadPoint(SharedData_Cuda &data, const CeedInt elem, const CeedInt p, const CeedInt points_in_elem,
                                 const CeedInt *__restrict__ indices, const CeedScalar *__restrict__ d_u, CeedScalar *r_u) {
  const CeedInt ind = indices[p + elem * NUM_PTS];

  for (CeedInt comp = 0; comp < NUM_COMP; comp++) {
    r_u[comp] = d_u[ind + comp * COMP_STRIDE];
  }
}

//------------------------------------------------------------------------------
// Single point -> L-vector
//------------------------------------------------------------------------------
template <int NUM_COMP, int COMP_STRIDE, int NUM_PTS>
inline __device__ void WritePoint(SharedData_Cuda &data, const CeedInt elem, const CeedInt p, const CeedInt points_in_elem,
                                  const CeedInt *__restrict__ indices, const CeedScalar *__restrict__ r_u, CeedScalar *d_u) {
  if (p < points_in_elem) {
    const CeedInt ind = indices[p + elem * NUM_PTS];

    for (CeedInt comp = 0; comp < NUM_COMP; comp++) {
      d_u[ind + comp * COMP_STRIDE] += r_u[comp];
    }
  }
}

//------------------------------------------------------------------------------
// 1D
//------------------------------------------------------------------------------

//------------------------------------------------------------------------------
// Set E-vector value
//------------------------------------------------------------------------------
template <int NUM_COMP, int P_1D>
inline __device__ void SetEVecStandard1d_Single(SharedData_Cuda &data, const CeedInt n, const CeedScalar value, CeedScalar *__restrict__ r_v) {
  const CeedInt target_comp = n / P_1D;
  const CeedInt target_node = n % P_1D;

  if (data.t_id_x == target_node) {
    r_v[target_comp] = value;
  }
}

//------------------------------------------------------------------------------
// L-vector -> E-vector, offsets provided
//------------------------------------------------------------------------------
template <int NUM_COMP, int COMP_STRIDE, int P_1D>
inline __device__ void ReadLVecStandard1d(SharedData_Cuda &data, const CeedInt num_nodes, const CeedInt elem, const CeedInt *__restrict__ indices,
                                          const CeedScalar *__restrict__ d_u, CeedScalar *__restrict__ r_u) {
  if (data.t_id_x < P_1D) {
    const CeedInt node = data.t_id_x;
    const CeedInt ind  = indices[node + elem * P_1D];

    for (CeedInt comp = 0; comp < NUM_COMP; comp++) r_u[comp] = d_u[ind + COMP_STRIDE * comp];
  }
}

//------------------------------------------------------------------------------
// L-vector -> E-vector, strided
//------------------------------------------------------------------------------
template <int NUM_COMP, int P_1D, int STRIDES_NODE, int STRIDES_COMP, int STRIDES_ELEM>
inline __device__ void ReadLVecStrided1d(SharedData_Cuda &data, const CeedInt elem, const CeedScalar *__restrict__ d_u,
                                         CeedScalar *__restrict__ r_u) {
  if (data.t_id_x < P_1D) {
    const CeedInt node = data.t_id_x;
    const CeedInt ind  = node * STRIDES_NODE + elem * STRIDES_ELEM;

    for (CeedInt comp = 0; comp < NUM_COMP; comp++) r_u[comp] = d_u[ind + comp * STRIDES_COMP];
  }
}

//------------------------------------------------------------------------------
// E-vector -> L-vector, offsets provided
//------------------------------------------------------------------------------
template <int NUM_COMP, int COMP_STRIDE, int P_1D>
inline __device__ void WriteLVecStandard1d(SharedData_Cuda &data, const CeedInt num_nodes, const CeedInt elem, const CeedInt *__restrict__ indices,
                                           const CeedScalar *__restrict__ r_v, CeedScalar *__restrict__ d_v) {
  if (data.t_id_x < P_1D) {
    const CeedInt node = data.t_id_x;
    const CeedInt ind  = indices[node + elem * P_1D];

    for (CeedInt comp = 0; comp < NUM_COMP; comp++) atomicAdd(&d_v[ind + COMP_STRIDE * comp], r_v[comp]);
  }
}

template <int NUM_COMP, int COMP_STRIDE, int P_1D>
inline __device__ void WriteLVecStandard1d_Single(SharedData_Cuda &data, const CeedInt num_nodes, const CeedInt elem, const CeedInt n,
                                                  const CeedInt *__restrict__ indices, const CeedScalar *__restrict__ r_v,
                                                  CeedScalar *__restrict__ d_v) {
  const CeedInt target_comp = n / P_1D;
  const CeedInt target_node = n % P_1D;

  if (data.t_id_x == target_node) {
    const CeedInt ind = indices[target_node + elem * P_1D];

    atomicAdd(&d_v[ind + COMP_STRIDE * target_comp], r_v[target_comp]);
  }
}

//------------------------------------------------------------------------------
// E-vector -> L-vector, strided
//------------------------------------------------------------------------------
template <int NUM_COMP, int P_1D, int STRIDES_NODE, int STRIDES_COMP, int STRIDES_ELEM>
inline __device__ void WriteLVecStrided1d(SharedData_Cuda &data, const CeedInt elem, const CeedScalar *__restrict__ r_v,
                                          CeedScalar *__restrict__ d_v) {
  if (data.t_id_x < P_1D) {
    const CeedInt node = data.t_id_x;
    const CeedInt ind  = node * STRIDES_NODE + elem * STRIDES_ELEM;

    for (CeedInt comp = 0; comp < NUM_COMP; comp++) d_v[ind + comp * STRIDES_COMP] += r_v[comp];
  }
}

//------------------------------------------------------------------------------
// 2D
//------------------------------------------------------------------------------

//------------------------------------------------------------------------------
// Set E-vector value
//------------------------------------------------------------------------------
template <int NUM_COMP, int P_1D>
inline __device__ void SetEVecStandard2d_Single(SharedData_Cuda &data, const CeedInt n, const CeedScalar value, CeedScalar *__restrict__ r_v) {
  const CeedInt target_comp   = n / (P_1D * P_1D);
  const CeedInt target_node_x = n % P_1D;
  const CeedInt target_node_y = (n % (P_1D * P_1D)) / P_1D;

  if (data.t_id_x == target_node_x && data.t_id_y == target_node_y) {
    r_v[target_comp] = value;
  }
}

//------------------------------------------------------------------------------
// L-vector -> E-vector, offsets provided
//------------------------------------------------------------------------------
template <int NUM_COMP, int COMP_STRIDE, int P_1D>
inline __device__ void ReadLVecStandard2d(SharedData_Cuda &data, const CeedInt num_nodes, const CeedInt elem, const CeedInt *__restrict__ indices,
                                          const CeedScalar *__restrict__ d_u, CeedScalar *__restrict__ r_u) {
  if (data.t_id_x < P_1D && data.t_id_y < P_1D) {
    const CeedInt node = data.t_id_x + data.t_id_y * P_1D;
    const CeedInt ind  = indices[node + elem * P_1D * P_1D];

    for (CeedInt comp = 0; comp < NUM_COMP; comp++) r_u[comp] = d_u[ind + COMP_STRIDE * comp];
  }
}

//------------------------------------------------------------------------------
// L-vector -> E-vector, strided
//------------------------------------------------------------------------------
template <int NUM_COMP, int P_1D, int STRIDES_NODE, int STRIDES_COMP, int STRIDES_ELEM>
inline __device__ void ReadLVecStrided2d(SharedData_Cuda &data, const CeedInt elem, const CeedScalar *__restrict__ d_u,
                                         CeedScalar *__restrict__ r_u) {
  if (data.t_id_x < P_1D && data.t_id_y < P_1D) {
    const CeedInt node = data.t_id_x + data.t_id_y * P_1D;
    const CeedInt ind  = node * STRIDES_NODE + elem * STRIDES_ELEM;

    for (CeedInt comp = 0; comp < NUM_COMP; comp++) r_u[comp] = d_u[ind + comp * STRIDES_COMP];
  }
}

//------------------------------------------------------------------------------
// E-vector -> L-vector, offsets provided
//------------------------------------------------------------------------------
template <int NUM_COMP, int COMP_STRIDE, int P_1D>
inline __device__ void WriteLVecStandard2d(SharedData_Cuda &data, const CeedInt num_nodes, const CeedInt elem, const CeedInt *__restrict__ indices,
                                           const CeedScalar *__restrict__ r_v, CeedScalar *__restrict__ d_v) {
  if (data.t_id_x < P_1D && data.t_id_y < P_1D) {
    const CeedInt node = data.t_id_x + data.t_id_y * P_1D;
    const CeedInt ind  = indices[node + elem * P_1D * P_1D];

    for (CeedInt comp = 0; comp < NUM_COMP; comp++) atomicAdd(&d_v[ind + COMP_STRIDE * comp], r_v[comp]);
  }
}

template <int NUM_COMP, int COMP_STRIDE, int P_1D>
inline __device__ void WriteLVecStandard2d_Single(SharedData_Cuda &data, const CeedInt num_nodes, const CeedInt elem, const CeedInt n,
                                                  const CeedInt *__restrict__ indices, const CeedScalar *__restrict__ r_v,
                                                  CeedScalar *__restrict__ d_v) {
  const CeedInt target_comp   = n / (P_1D * P_1D);
  const CeedInt target_node_x = n % P_1D;
  const CeedInt target_node_y = (n % (P_1D * P_1D)) / P_1D;

  if (data.t_id_x == target_node_x && data.t_id_y == target_node_y) {
    const CeedInt node = data.t_id_x + data.t_id_y * P_1D;
    const CeedInt ind  = indices[node + elem * P_1D * P_1D];

    atomicAdd(&d_v[ind + COMP_STRIDE * target_comp], r_v[target_comp]);
  }
}

//------------------------------------------------------------------------------
// E-vector -> L-vector, strided
//------------------------------------------------------------------------------
template <int NUM_COMP, int P_1D, int STRIDES_NODE, int STRIDES_COMP, int STRIDES_ELEM>
inline __device__ void WriteLVecStrided2d(SharedData_Cuda &data, const CeedInt elem, const CeedScalar *__restrict__ r_v,
                                          CeedScalar *__restrict__ d_v) {
  if (data.t_id_x < P_1D && data.t_id_y < P_1D) {
    const CeedInt node = data.t_id_x + data.t_id_y * P_1D;
    const CeedInt ind  = node * STRIDES_NODE + elem * STRIDES_ELEM;

    for (CeedInt comp = 0; comp < NUM_COMP; comp++) d_v[ind + comp * STRIDES_COMP] += r_v[comp];
  }
}

//------------------------------------------------------------------------------
// 3D
//------------------------------------------------------------------------------

//------------------------------------------------------------------------------
// Set E-vector value
//------------------------------------------------------------------------------
template <int NUM_COMP, int P_1D>
inline __device__ void SetEVecStandard3d_Single(SharedData_Cuda &data, const CeedInt n, const CeedScalar value, CeedScalar *__restrict__ r_v) {
  const CeedInt target_comp   = n / (P_1D * P_1D * P_1D);
  const CeedInt target_node_x = n % P_1D;
  const CeedInt target_node_y = (n % (P_1D * P_1D)) / P_1D;
  const CeedInt target_node_z = (n % (P_1D * P_1D * P_1D)) / (P_1D * P_1D);

  if (data.t_id_x == target_node_x && data.t_id_y == target_node_y) {
    r_v[target_node_z + target_comp * P_1D] = value;
  }
}

//------------------------------------------------------------------------------
// L-vector -> E-vector, offsets provided
//------------------------------------------------------------------------------
template <int NUM_COMP, int COMP_STRIDE, int P_1D>
inline __device__ void ReadLVecStandard3d(SharedData_Cuda &data, const CeedInt num_nodes, const CeedInt elem, const CeedInt *__restrict__ indices,
                                          const CeedScalar *__restrict__ d_u, CeedScalar *__restrict__ r_u) {
  if (data.t_id_x < P_1D && data.t_id_y < P_1D) {
    for (CeedInt z = 0; z < P_1D; z++) {
      const CeedInt node = data.t_id_x + data.t_id_y * P_1D + z * P_1D * P_1D;
      const CeedInt ind  = indices[node + elem * P_1D * P_1D * P_1D];

      for (CeedInt comp = 0; comp < NUM_COMP; comp++) r_u[z + comp * P_1D] = d_u[ind + COMP_STRIDE * comp];
    }
  }
}

//------------------------------------------------------------------------------
// L-vector -> E-vector, strided
//------------------------------------------------------------------------------
template <int NUM_COMP, int P_1D, int STRIDES_NODE, int STRIDES_COMP, int STRIDES_ELEM>
inline __device__ void ReadLVecStrided3d(SharedData_Cuda &data, const CeedInt elem, const CeedScalar *__restrict__ d_u,
                                         CeedScalar *__restrict__ r_u) {
  if (data.t_id_x < P_1D && data.t_id_y < P_1D) {
    for (CeedInt z = 0; z < P_1D; z++) {
      const CeedInt node = data.t_id_x + data.t_id_y * P_1D + z * P_1D * P_1D;
      const CeedInt ind  = node * STRIDES_NODE + elem * STRIDES_ELEM;

      for (CeedInt comp = 0; comp < NUM_COMP; comp++) r_u[z + comp * P_1D] = d_u[ind + comp * STRIDES_COMP];
    }
  }
}

//------------------------------------------------------------------------------
// E-vector -> Q-vector, offests provided
//------------------------------------------------------------------------------
template <int NUM_COMP, int COMP_STRIDE, int Q_1D>
inline __device__ void ReadEVecSliceStandard3d(SharedData_Cuda &data, const CeedInt nquads, const CeedInt elem, const CeedInt q,
                                               const CeedInt *__restrict__ indices, const CeedScalar *__restrict__ d_u,
                                               CeedScalar *__restrict__ r_u) {
  if (data.t_id_x < Q_1D && data.t_id_y < Q_1D) {
    const CeedInt node = data.t_id_x + data.t_id_y * Q_1D + q * Q_1D * Q_1D;
    const CeedInt ind  = indices[node + elem * Q_1D * Q_1D * Q_1D];

    for (CeedInt comp = 0; comp < NUM_COMP; comp++) r_u[comp] = d_u[ind + COMP_STRIDE * comp];
  }
}

//------------------------------------------------------------------------------
// E-vector -> Q-vector, strided
//------------------------------------------------------------------------------
template <int NUM_COMP, int Q_1D, int STRIDES_NODE, int STRIDES_COMP, int STRIDES_ELEM>
inline __device__ void ReadEVecSliceStrided3d(SharedData_Cuda &data, const CeedInt elem, const CeedInt q, const CeedScalar *__restrict__ d_u,
                                              CeedScalar *__restrict__ r_u) {
  if (data.t_id_x < Q_1D && data.t_id_y < Q_1D) {
    const CeedInt node = data.t_id_x + data.t_id_y * Q_1D + q * Q_1D * Q_1D;
    const CeedInt ind  = node * STRIDES_NODE + elem * STRIDES_ELEM;

    for (CeedInt comp = 0; comp < NUM_COMP; comp++) r_u[comp] = d_u[ind + comp * STRIDES_COMP];
  }
}

//------------------------------------------------------------------------------
// E-vector -> L-vector, offsets provided
//------------------------------------------------------------------------------
template <int NUM_COMP, int COMP_STRIDE, int P_1D>
inline __device__ void WriteLVecStandard3d(SharedData_Cuda &data, const CeedInt num_nodes, const CeedInt elem, const CeedInt *__restrict__ indices,
                                           const CeedScalar *__restrict__ r_v, CeedScalar *__restrict__ d_v) {
  if (data.t_id_x < P_1D && data.t_id_y < P_1D) {
    for (CeedInt z = 0; z < P_1D; z++) {
      const CeedInt node = data.t_id_x + data.t_id_y * P_1D + z * P_1D * P_1D;
      const CeedInt ind  = indices[node + elem * P_1D * P_1D * P_1D];

      for (CeedInt comp = 0; comp < NUM_COMP; comp++) atomicAdd(&d_v[ind + COMP_STRIDE * comp], r_v[z + comp * P_1D]);
    }
  }
}

template <int NUM_COMP, int COMP_STRIDE, int P_1D>
inline __device__ void WriteLVecStandard3d_Single(SharedData_Cuda &data, const CeedInt num_nodes, const CeedInt elem, const CeedInt n,
                                                  const CeedInt *__restrict__ indices, const CeedScalar *__restrict__ r_v,
                                                  CeedScalar *__restrict__ d_v) {
  const CeedInt target_comp   = n / (P_1D * P_1D * P_1D);
  const CeedInt target_node_x = n % P_1D;
  const CeedInt target_node_y = (n % (P_1D * P_1D)) / P_1D;
  const CeedInt target_node_z = (n % (P_1D * P_1D * P_1D)) / (P_1D * P_1D);

  if (data.t_id_x == target_node_x && data.t_id_y == target_node_y) {
    const CeedInt node = data.t_id_x + data.t_id_y * P_1D + target_node_z * P_1D * P_1D;
    const CeedInt ind  = indices[node + elem * P_1D * P_1D * P_1D];

    atomicAdd(&d_v[ind + COMP_STRIDE * target_comp], r_v[target_node_z + target_comp * P_1D]);
  }
}

//------------------------------------------------------------------------------
// E-vector -> L-vector, strided
//------------------------------------------------------------------------------
template <int NUM_COMP, int P_1D, int STRIDES_NODE, int STRIDES_COMP, int STRIDES_ELEM>
inline __device__ void WriteLVecStrided3d(SharedData_Cuda &data, const CeedInt elem, const CeedScalar *__restrict__ r_v,
                                          CeedScalar *__restrict__ d_v) {
  if (data.t_id_x < P_1D && data.t_id_y < P_1D) {
    for (CeedInt z = 0; z < P_1D; z++) {
      const CeedInt node = data.t_id_x + data.t_id_y * P_1D + z * P_1D * P_1D;
      const CeedInt ind  = node * STRIDES_NODE + elem * STRIDES_ELEM;

      for (CeedInt comp = 0; comp < NUM_COMP; comp++) d_v[ind + comp * STRIDES_COMP] += r_v[z + comp * P_1D];
    }
  }
}

//------------------------------------------------------------------------------
// 3D collocated derivatives computation
//------------------------------------------------------------------------------
template <int NUM_COMP, int Q_1D, int T_1D>
inline __device__ void GradColloSlice3d(SharedData_Cuda &data, const CeedInt q, const CeedScalar *__restrict__ r_U, const CeedScalar *c_G,
                                        CeedScalar *__restrict__ r_V) {
  if (data.t_id_x < Q_1D && data.t_id_y < Q_1D) {
    for (CeedInt comp = 0; comp < NUM_COMP; comp++) {
      __syncthreads();
      data.slice[data.t_id_x + data.t_id_y * T_1D] = r_U[q + comp * Q_1D];
      __syncthreads();
      // X derivative
      r_V[comp + 0 * NUM_COMP] = 0.0;
      for (CeedInt i = 0; i < Q_1D; i++) {
        r_V[comp + 0 * NUM_COMP] += c_G[i + data.t_id_x * Q_1D] * data.slice[i + data.t_id_y * T_1D];
      }
      // Y derivative
      r_V[comp + 1 * NUM_COMP] = 0.0;
      for (CeedInt i = 0; i < Q_1D; i++) {
        r_V[comp + 1 * NUM_COMP] += c_G[i + data.t_id_y * Q_1D] * data.slice[data.t_id_x + i * T_1D];
      }
      // Z derivative
      r_V[comp + 2 * NUM_COMP] = 0.0;
      for (CeedInt i = 0; i < Q_1D; i++) {
        r_V[comp + 2 * NUM_COMP] += c_G[i + q * Q_1D] * r_U[i + comp * Q_1D];
      }
    }
  }
}

//------------------------------------------------------------------------------
// 3D collocated derivatives transpose
//------------------------------------------------------------------------------
template <int NUM_COMP, int Q_1D, int T_1D>
inline __device__ void GradColloSliceTranspose3d(SharedData_Cuda &data, const CeedInt q, const CeedScalar *__restrict__ r_U, const CeedScalar *c_G,
                                                 CeedScalar *__restrict__ r_V) {
  if (data.t_id_x < Q_1D && data.t_id_y < Q_1D) {
    for (CeedInt comp = 0; comp < NUM_COMP; comp++) {
      __syncthreads();
      data.slice[data.t_id_x + data.t_id_y * T_1D] = r_U[comp + 0 * NUM_COMP];
      __syncthreads();
      // X derivative
      for (CeedInt i = 0; i < Q_1D; i++) {
        r_V[q + comp * Q_1D] += c_G[data.t_id_x + i * Q_1D] * data.slice[i + data.t_id_y * T_1D];
      }
      // Y derivative
      __syncthreads();
      data.slice[data.t_id_x + data.t_id_y * T_1D] = r_U[comp + 1 * NUM_COMP];
      __syncthreads();
      for (CeedInt i = 0; i < Q_1D; i++) {
        r_V[q + comp * Q_1D] += c_G[data.t_id_y + i * Q_1D] * data.slice[data.t_id_x + i * T_1D];
      }
      // Z derivative
      for (CeedInt i = 0; i < Q_1D; i++) {
        r_V[i + comp * Q_1D] += c_G[i + q * Q_1D] * r_U[comp + 2 * NUM_COMP];
      }
    }
  }
}