xref: /libCEED/include/ceed/jit-source/magma/magma-basis-grad-1d.h (revision d4cc18453651bd0f94c1a2e078b2646a92dafdcc)

// Copyright (c) 2017-2026, 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 MAGMA tensor basis gradient in 1D
#include "magma-common-tensor.h"

// macros to abstract access of shared memory and reg. file
#define sT(i, j) sT[(j) * P + (i)]

////////////////////////////////////////////////////////////////////////////////
// grad basis action (1D)
template <typename T, int DIM, int NUM_COMP, int P, int Q>
static __device__ __inline__ void magma_grad_1d_device(const T *sT, T *sU[NUM_COMP], T *sV[NUM_COMP], const int tx) {
  // Assumptions
  // 1. 1D threads of size max(P,Q)
  // 2. sU[i] is 1xP: in shared memory
  // 3. sV[i] is 1xQ: in shared memory
  // 4. P_roduct per component is one row (1xP) times T matrix (PxQ) => one row (1xQ)
  // 5. Each thread computes one entry in sV[i]
  // 6. Must sync before and after call
  // 7. Note that the layout for U and V is different from 2D/3D problem

  if (tx < Q) {
    for (int comp = 0; comp < NUM_COMP; comp++) {
      T rv = 0.0;
      for (int i = 0; i < P; i++) {
        rv += sU[comp][i] * sT(i, tx);
      }
      sV[comp][tx] = rv;
    }
  }
}

////////////////////////////////////////////////////////////////////////////////
extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q, MAGMA_MAXTHREADS_1D)) __global__
    void magma_gradn_1d_kernel(const CeedScalar *dTinterp, const CeedScalar *dTgrad, const CeedScalar *dU, const int estrdU, const int cstrdU,
                               const int dstrdU, CeedScalar *dV, const int estrdV, const int cstrdV, const int dstrdV, const int nelem) {
  MAGMA_DEVICE_SHARED(CeedScalar, shared_data)

  const int tx      = threadIdx.x;
  const int ty      = threadIdx.y;
  const int elem_id = (blockIdx.x * blockDim.y) + ty;

  if (elem_id >= nelem) return;

  CeedScalar *sU[BASIS_NUM_COMP];
  CeedScalar *sV[BASIS_NUM_COMP];

  // shift global memory pointers by elem stride
  dU += elem_id * estrdU;
  dV += elem_id * estrdV;

  // assign shared memory pointers
  CeedScalar *sT = (CeedScalar *)shared_data;
  CeedScalar *sW = sT + BASIS_P * BASIS_Q;
  sU[0]          = sW + ty * BASIS_NUM_COMP * (BASIS_P + BASIS_Q);
  sV[0]          = sU[0] + (BASIS_NUM_COMP * 1 * BASIS_P);
  for (int comp = 1; comp < BASIS_NUM_COMP; comp++) {
    sU[comp] = sU[comp - 1] + (1 * BASIS_P);
    sV[comp] = sV[comp - 1] + (1 * BASIS_Q);
  }

  // read T
  if (ty == 0) {
    read_T_notrans_gm2sm<BASIS_P, BASIS_Q>(tx, dTgrad, sT);
  }

  // read U
  read_1d<CeedScalar, BASIS_P, BASIS_NUM_COMP>(dU, cstrdU, sU, tx);

  __syncthreads();
  magma_grad_1d_device<CeedScalar, BASIS_DIM, BASIS_NUM_COMP, BASIS_P, BASIS_Q>(sT, sU, sV, tx);
  __syncthreads();

  // write V
  write_1d<CeedScalar, BASIS_Q, BASIS_NUM_COMP>(sV, dV, cstrdV, tx);
}

////////////////////////////////////////////////////////////////////////////////
extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q, MAGMA_MAXTHREADS_1D)) __global__
    void magma_gradt_1d_kernel(const CeedScalar *dTinterp, const CeedScalar *dTgrad, const CeedScalar *dU, const int estrdU, const int cstrdU,
                               const int dstrdU, CeedScalar *dV, const int estrdV, const int cstrdV, const int dstrdV, const int nelem) {
  MAGMA_DEVICE_SHARED(CeedScalar, shared_data)

  const int tx      = threadIdx.x;
  const int ty      = threadIdx.y;
  const int elem_id = (blockIdx.x * blockDim.y) + ty;

  if (elem_id >= nelem) return;

  CeedScalar *sU[BASIS_NUM_COMP];
  CeedScalar *sV[BASIS_NUM_COMP];

  // shift global memory pointers by elem stride
  dU += elem_id * estrdU;
  dV += elem_id * estrdV;

  // assign shared memory pointers
  CeedScalar *sT = (CeedScalar *)shared_data;
  CeedScalar *sW = sT + BASIS_Q * BASIS_P;
  sU[0]          = sW + ty * BASIS_NUM_COMP * (BASIS_Q + BASIS_P);
  sV[0]          = sU[0] + (BASIS_NUM_COMP * 1 * BASIS_Q);
  for (int comp = 1; comp < BASIS_NUM_COMP; comp++) {
    sU[comp] = sU[comp - 1] + (1 * BASIS_Q);
    sV[comp] = sV[comp - 1] + (1 * BASIS_P);
  }

  // read T
  if (ty == 0) {
    read_T_trans_gm2sm<BASIS_Q, BASIS_P>(tx, dTgrad, sT);
  }

  // read U
  read_1d<CeedScalar, BASIS_Q, BASIS_NUM_COMP>(dU, cstrdU, sU, tx);

  __syncthreads();
  magma_grad_1d_device<CeedScalar, BASIS_DIM, BASIS_NUM_COMP, BASIS_Q, BASIS_P>(sT, sU, sV, tx);
  __syncthreads();

  // write V
  write_1d<CeedScalar, BASIS_P, BASIS_NUM_COMP>(sV, dV, cstrdV, tx);
}

////////////////////////////////////////////////////////////////////////////////
extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q, MAGMA_MAXTHREADS_1D)) __global__
    void magma_gradta_1d_kernel(const CeedScalar *dTinterp, const CeedScalar *dTgrad, const CeedScalar *dU, const int estrdU, const int cstrdU,
                                const int dstrdU, CeedScalar *dV, const int estrdV, const int cstrdV, const int dstrdV, const int nelem) {
  MAGMA_DEVICE_SHARED(CeedScalar, shared_data)

  const int tx      = threadIdx.x;
  const int ty      = threadIdx.y;
  const int elem_id = (blockIdx.x * blockDim.y) + ty;

  if (elem_id >= nelem) return;

  CeedScalar *sU[BASIS_NUM_COMP];
  CeedScalar *sV[BASIS_NUM_COMP];

  // shift global memory pointers by elem stride
  dU += elem_id * estrdU;
  dV += elem_id * estrdV;

  // assign shared memory pointers
  CeedScalar *sT = (CeedScalar *)shared_data;
  CeedScalar *sW = sT + BASIS_Q * BASIS_P;
  sU[0]          = sW + ty * BASIS_NUM_COMP * (BASIS_Q + BASIS_P);
  sV[0]          = sU[0] + (BASIS_NUM_COMP * 1 * BASIS_Q);
  for (int comp = 1; comp < BASIS_NUM_COMP; comp++) {
    sU[comp] = sU[comp - 1] + (1 * BASIS_Q);
    sV[comp] = sV[comp - 1] + (1 * BASIS_P);
  }

  // read T
  if (ty == 0) {
    read_T_trans_gm2sm<BASIS_Q, BASIS_P>(tx, dTgrad, sT);
  }

  // read U
  read_1d<CeedScalar, BASIS_Q, BASIS_NUM_COMP>(dU, cstrdU, sU, tx);

  __syncthreads();
  magma_grad_1d_device<CeedScalar, BASIS_DIM, BASIS_NUM_COMP, BASIS_Q, BASIS_P>(sT, sU, sV, tx);
  __syncthreads();

  // sum into V
  sum_1d<CeedScalar, BASIS_P, BASIS_NUM_COMP>(sV, dV, cstrdV, tx);
}